syncs: add Semaphore

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
pull/1575/head
Brad Fitzpatrick 4 years ago committed by Brad Fitzpatrick
parent 9643d8b34d
commit 77ec80538a

@ -5,7 +5,10 @@
// Package syncs contains additional sync types and functionality. // Package syncs contains additional sync types and functionality.
package syncs package syncs
import "sync/atomic" import (
"context"
"sync/atomic"
)
// ClosedChan returns a channel that's already closed. // ClosedChan returns a channel that's already closed.
func ClosedChan() <-chan struct{} { return closedChan } func ClosedChan() <-chan struct{} { return closedChan }
@ -79,3 +82,45 @@ func (b *AtomicBool) Set(v bool) {
func (b *AtomicBool) Get() bool { func (b *AtomicBool) Get() bool {
return atomic.LoadInt32((*int32)(b)) != 0 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
}

@ -4,7 +4,10 @@
package syncs package syncs
import "testing" import (
"context"
"testing"
)
func TestWaitGroupChan(t *testing.T) { func TestWaitGroupChan(t *testing.T) {
wg := NewWaitGroupChan() 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()
}

Loading…
Cancel
Save