30 lines
389 B
Go
30 lines
389 B
Go
|
package utils
|
||
|
|
|
||
|
|
import (
|
||
|
|
"sync"
|
||
|
|
"sync/atomic"
|
||
|
|
)
|
||
|
|
|
||
|
|
type CountWaitGroup struct {
|
||
|
|
wg sync.WaitGroup
|
||
|
|
num atomic.Uint64
|
||
|
|
}
|
||
|
|
|
||
|
|
func (c *CountWaitGroup) Add(delta uint64) {
|
||
|
|
c.wg.Add(int(delta))
|
||
|
|
c.num.Add(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 c.num.Load()
|
||
|
|
}
|