45 lines
596 B
Go
45 lines
596 B
Go
package utils
|
|
|
|
import (
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
type WaitGroup interface {
|
|
Add(delta int)
|
|
Done()
|
|
Wait()
|
|
}
|
|
|
|
type CountWaitGroup struct {
|
|
wg sync.WaitGroup
|
|
num atomic.Int64
|
|
}
|
|
|
|
func (c *CountWaitGroup) Add(delta int) {
|
|
c.wg.Add(delta)
|
|
c.num.Add(int64(delta))
|
|
}
|
|
|
|
func (c *CountWaitGroup) Done() {
|
|
c.wg.Done()
|
|
c.num.Add(-1)
|
|
}
|
|
|
|
func (c *CountWaitGroup) Wait() {
|
|
c.wg.Wait()
|
|
}
|
|
|
|
func (c *CountWaitGroup) Count() uint64 {
|
|
return uint64(c.num.Load())
|
|
}
|
|
|
|
func WgWait[T WaitGroup](wg T) <-chan struct{} {
|
|
ch := make(chan struct{})
|
|
go func() {
|
|
wg.Wait()
|
|
ch <- struct{}{}
|
|
}()
|
|
return ch
|
|
}
|