Add ConcurrencyLimit

client/concurrencylimit.go

@@ -0,0 +1,34 @@
+package client
+
+import "context"
+
+import "golang.org/x/sync/semaphore"
+
+type ConcurrencyLimit struct {
+	sem *semaphore.Weighted
+}
+
+func NewConcurrencyLimit(limit int64) *ConcurrencyLimit {
+	return &ConcurrencyLimit{
+		sem: semaphore.NewWeighted(limit),
+	}
+}
+
+func (cl *ConcurrencyLimit) Acquire1() {
+	cl.AcquireN(1)
+}
+
+func (cl *ConcurrencyLimit) AcquireN(cost int64) {
+	err := cl.sem.Acquire(context.TODO(), cost)
+	if err != nil {
+		panic(err)
+	}
+}
+
+func (cl *ConcurrencyLimit) Release1() {
+	cl.ReleaseN(1)
+}
+
+func (cl *ConcurrencyLimit) ReleaseN(cost int64) {
+	cl.sem.Release(cost)
+}

go.mod

@@ -5,4 +5,5 @@ go 1.13
 require (
 	cloud.google.com/go v0.94.1
 	golang.org/x/net v0.0.0-20210908191846-a5e095526f91
+	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
 )

go.sum

@@ -274,6 +274,7 @@ golang.org/x/sync v0.0.0-20200317015054-43a5402ce75a/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sync v0.0.0-20200625203802-6e8e738ad208/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20201207232520-09787c993a3a/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
+golang.org/x/sync v0.0.0-20210220032951-036812b2e83c h1:5KslGYwFpkhGh+Q16bwMP3cOontH8FOep7tGV86Y7SQ=
 golang.org/x/sync v0.0.0-20210220032951-036812b2e83c/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sys v0.0.0-20180830151530-49385e6e1522/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=