syncs: add Semaphore

Signed-off-by: Brad Fitzpatrick <[email protected]>

syncs/syncs.go

@@ -5,7 +5,10 @@
 // Package syncs contains additional sync types and functionality.
 package syncs
 
-import "sync/atomic"
+import (
+	"context"
+	"sync/atomic"
+)
 
 // ClosedChan returns a channel that's already closed.
 func ClosedChan() <-chan struct{} { return closedChan }
@@ -79,3 +82,45 @@ func (b *AtomicBool) Set(v bool) {
 func (b *AtomicBool) Get() bool {
 	return atomic.LoadInt32((*int32)(b)) != 0
 }
+
+// Semaphore is a counting semaphore.
+//
+// Use NewSemaphore to create one.
+type Semaphore struct {
+	c chan struct{}
+}
+
+// NewSemaphore returns a semaphore with resource count n.
+func NewSemaphore(n int) Semaphore {
+	return Semaphore{c: make(chan struct{}, n)}
+}
+
+// Acquire blocks until a resource is acquired.
+func (s Semaphore) Acquire() {
+	s.c <- struct{}{}
+}
+
+// AcquireContext reports whether the resource was acquired before the ctx was done.
+func (s Semaphore) AcquireContext(ctx context.Context) bool {
+	select {
+	case s.c <- struct{}{}:
+		return true
+	case <-ctx.Done():
+		return false
+	}
+}
+
+// TryAcquire reports, without blocking, whether the resource was acquired.
+func (s Semaphore) TryAcquire() bool {
+	select {
+	case s.c <- struct{}{}:
+		return true
+	default:
+		return false
+	}
+}
+
+// Release releases a resource.
+func (s Semaphore) Release() {
+	<-s.c
+}

syncs/syncs_test.go

@@ -4,7 +4,10 @@
 
 package syncs
 
-import "testing"
+import (
+	"context"
+	"testing"
+)
 
 func TestWaitGroupChan(t *testing.T) {
 	wg := NewWaitGroupChan()
@@ -48,3 +51,25 @@ func TestClosedChan(t *testing.T) {
 		}
 	}
 }
+
+func TestSemaphore(t *testing.T) {
+	s := NewSemaphore(2)
+	s.Acquire()
+	if !s.TryAcquire() {
+		t.Fatal("want true")
+	}
+	if s.TryAcquire() {
+		t.Fatal("want false")
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	if s.AcquireContext(ctx) {
+		t.Fatal("want false")
+	}
+	s.Release()
+	if !s.AcquireContext(context.Background()) {
+		t.Fatal("want true")
+	}
+	s.Release()
+	s.Release()
+}