64 lines
1016 B
Go
64 lines
1016 B
Go
package semaphore
|
|
|
|
import (
|
|
"context"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
type Semaphore interface {
|
|
Acquire(ctx context.Context) error
|
|
Release()
|
|
}
|
|
|
|
// ChanSemaphore implementation
|
|
type ChanSemaphore struct {
|
|
C chan struct{}
|
|
}
|
|
|
|
func (s *ChanSemaphore) Acquire(ctx context.Context) error {
|
|
select {
|
|
case s.C <- struct{}{}:
|
|
return nil
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
}
|
|
}
|
|
|
|
func (s *ChanSemaphore) Release() {
|
|
<-s.C
|
|
}
|
|
|
|
// AtomicSemaphore implementation
|
|
type AtomicSemaphore struct {
|
|
acquired atomic.Bool
|
|
TryDelay time.Duration
|
|
}
|
|
|
|
func (s *AtomicSemaphore) Acquire(ctx context.Context) error {
|
|
if s.acquired.CompareAndSwap(false, true) {
|
|
return nil
|
|
}
|
|
|
|
return s.slowAcquire(ctx)
|
|
}
|
|
|
|
func (s *AtomicSemaphore) slowAcquire(ctx context.Context) error {
|
|
ticker := time.NewTicker(s.TryDelay)
|
|
|
|
for {
|
|
select {
|
|
case <-ctx.Done():
|
|
return ctx.Err()
|
|
case <-ticker.C:
|
|
if s.acquired.CompareAndSwap(false, true) {
|
|
return nil
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *AtomicSemaphore) Release() {
|
|
s.acquired.Store(false)
|
|
}
|