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

     1  package closer
     2  
     3  import (
     4  	"errors"
     5  	"sync"
     6  	"sync/atomic"
     7  )
     8  
     9  // Strict closer is a sync.WaitGroup with state.
    10  // It guarantees no Add with positive delta will ever succeed after Wait.
    11  // Wait can only be called once, but Add and Wait can be called concurrently.
    12  // The happens before relationship between Add and Wait is taken care of automatically.
    13  type Strict struct {
    14  	mu      sync.RWMutex
    15  	cond    *sync.Cond
    16  	closed  uint32
    17  	counter int32
    18  	done    chan struct{}
    19  }
    20  
    21  // NewStrict is ctor for Strict
    22  func NewStrict() *Strict {
    23  	s := &Strict{done: make(chan struct{})}
    24  	s.cond = sync.NewCond(&s.mu)
    25  	return s
    26  }
    27  
    28  var (
    29  	errAlreadyClosed = errors.New("closer already closed")
    30  )
    31  
    32  // Add delta to wait group
    33  // Trying to Add positive delta after Wait will return non nil error
    34  func (s *Strict) Add(delta int) (err error) {
    35  	if delta > 0 {
    36  		if s.HasBeenClosed() {
    37  			err = errAlreadyClosed
    38  			return
    39  		}
    40  
    41  		s.mu.RLock()
    42  		if s.closed != 0 {
    43  			s.mu.RUnlock()
    44  			err = errAlreadyClosed
    45  			return
    46  		}
    47  	}
    48  
    49  	counter := atomic.AddInt32(&s.counter, int32(delta))
    50  
    51  	if delta > 0 {
    52  		s.mu.RUnlock()
    53  	}
    54  
    55  	if counter == 0 {
    56  		s.mu.RLock()
    57  		if s.HasBeenClosed() {
    58  			s.cond.Signal()
    59  		}
    60  		s.mu.RUnlock()
    61  	}
    62  
    63  	return
    64  }
    65  
    66  // HasBeenClosed tells whether closed
    67  func (s *Strict) HasBeenClosed() bool {
    68  	return atomic.LoadUint32(&s.closed) != 0
    69  }
    70  
    71  // ClosedSignal gets signaled when Wait() is called.
    72  func (s *Strict) ClosedSignal() <-chan struct{} {
    73  	return s.done
    74  }
    75  
    76  // Done decrements the WaitGroup counter by one.
    77  func (s *Strict) Done() {
    78  	s.Add(-1)
    79  }
    80  
    81  // SignalAndWait updates closed and blocks until the WaitGroup counter is zero.
    82  // Call it more than once will panic
    83  func (s *Strict) SignalAndWait() {
    84  	close(s.done)
    85  
    86  	s.mu.Lock()
    87  
    88  	atomic.StoreUint32(&s.closed, 1)
    89  
    90  	for atomic.LoadInt32(&s.counter) != 0 {
    91  		s.cond.Wait()
    92  	}
    93  
    94  	s.mu.Unlock()
    95  }