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

     1  package closer
     2  
     3  import "sync"
     4  
     5  // Naive holds the two things we need to close a goroutine and wait for it to finish: a chan
     6  // to tell the goroutine to shut down, and a WaitGroup with which to wait for it to finish shutting
     7  // down.
     8  // The happens before relationship between Add and Wait is guaranteed by the caller side.
     9  type Naive struct {
    10  	waiting sync.WaitGroup
    11  	done    chan struct{}
    12  }
    13  
    14  // NewNaive is ctor for Naive
    15  func NewNaive() *Naive {
    16  	return &Naive{done: make(chan struct{})}
    17  }
    18  
    19  // WaitGroupRef returns the reference of WaitGroup
    20  func (n *Naive) WaitGroupRef() *sync.WaitGroup {
    21  	return &n.waiting
    22  }
    23  
    24  // Add delta to WaitGroup
    25  func (n *Naive) Add(delta int) {
    26  	n.waiting.Add(delta)
    27  }
    28  
    29  // ClosedSignal gets signaled when Wait() is called.
    30  func (n *Naive) ClosedSignal() <-chan struct{} {
    31  	return n.done
    32  }
    33  
    34  // Done calls Done() on the WaitGroup.
    35  func (n *Naive) Done() {
    36  	n.waiting.Done()
    37  }
    38  
    39  // SignalAndWait closes chan and wait on the WaitGroup.
    40  // Call it more than once will panic
    41  func (n *Naive) SignalAndWait() {
    42  	close(n.done)
    43  	n.waiting.Wait()
    44  }