github.com/zhiqiangxu/util@v0.0.0-20230112053021-0a7aee056cd5/wm/max.go (about)

     1  package wm
     2  
     3  import (
     4  	"context"
     5  
     6  	"golang.org/x/sync/semaphore"
     7  )
     8  
     9  // Max watermark model
    10  type Max struct {
    11  	sema *semaphore.Weighted
    12  }
    13  
    14  // NewMax is ctor for Max
    15  func NewMax(max int64) *Max {
    16  	m := &Max{sema: semaphore.NewWeighted(max)}
    17  	return m
    18  }
    19  
    20  // Enter returns nil if not canceled
    21  func (m *Max) Enter(ctx context.Context) error {
    22  	return m.sema.Acquire(ctx, 1)
    23  }
    24  
    25  // TryEnter returns true if succeed
    26  func (m *Max) TryEnter() bool {
    27  	return m.sema.TryAcquire(1)
    28  }
    29  
    30  // Exit should only be called if Enter returns nil
    31  func (m *Max) Exit() {
    32  	m.sema.Release(1)
    33  }