From 1846eb45bcb97cfa41784143be813d8943e1e9ca Mon Sep 17 00:00:00 2001
From: Ian Gulliver <ian@flamingcow.tv>
Date: Sun, 26 Sep 2021 04:21:03 +0000
Subject: [PATCH] Add ConcurrencyLimit

---
 client/concurrencylimit.go | 34 ++++++++++++++++++++++++++++++++++
 go.mod                     |  1 +
 go.sum                     |  1 +
 3 files changed, 36 insertions(+)
 create mode 100644 client/concurrencylimit.go

diff --git a/client/concurrencylimit.go b/client/concurrencylimit.go
new file mode 100644
index 0000000..9930560
--- /dev/null
+++ b/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)
+}
diff --git a/go.mod b/go.mod
index 8ef5598..8efbca4 100644
--- a/go.mod
+++ b/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
 )
diff --git a/go.sum b/go.sum
index 85ff409..4183b34 100644
--- a/go.sum
+++ b/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=