github.com/go-board/x-go@v0.1.2-0.20220610024734-db1323f6cb15/concurrent/semaphore.go (about)

     1  package concurrent
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  
     7  	"golang.org/x/sync/semaphore"
     8  )
     9  
    10  type Semaphore struct {
    11  	w *semaphore.Weighted
    12  }
    13  
    14  func (s *Semaphore) SpawnContext(ctx context.Context, fn func(ctx context.Context) error) error {
    15  	if err := s.w.Acquire(ctx, 1); err != nil {
    16  		return err
    17  	}
    18  	defer s.w.Release(1)
    19  	return fn(ctx)
    20  }
    21  
    22  func (s *Semaphore) Spawn(fn func() error) error {
    23  	if err := s.w.Acquire(context.Background(), 1); err != nil {
    24  		return err
    25  	}
    26  	defer s.w.Release(1)
    27  	return fn()
    28  }
    29  
    30  func (s *Semaphore) TrySpawnContext(ctx context.Context, fn func(ctx context.Context) error) error {
    31  	if !s.w.TryAcquire(1) {
    32  		return errors.New("err: try acquire semaphore failed")
    33  	}
    34  	defer s.w.Release(1)
    35  	return fn(ctx)
    36  }
    37  
    38  func (s *Semaphore) TrySpawn(fn func() error) error {
    39  	if !s.w.TryAcquire(1) {
    40  		return errors.New("err: try acquire semaphore failed")
    41  	}
    42  	defer s.w.Release(1)
    43  	return fn()
    44  }
    45  
    46  func NewSemaphore(n int64) *Semaphore {
    47  	return &Semaphore{w: semaphore.NewWeighted(n)}
    48  }