diff --git a/syncs/syncs.go b/syncs/syncs.go index 0139ad925..f319f6489 100644 --- a/syncs/syncs.go +++ b/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 +} diff --git a/syncs/syncs_test.go b/syncs/syncs_test.go index 9de72e22f..a6768e90b 100644 --- a/syncs/syncs_test.go +++ b/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() +}