74 lines
1008 B
Go
74 lines
1008 B
Go
package sync
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
type WaitGroup struct {
|
|
wg *sync.WaitGroup
|
|
drain *atomic.Bool
|
|
c *atomic.Int64
|
|
}
|
|
|
|
func WrapWaitGroup(wg *sync.WaitGroup) *WaitGroup {
|
|
g := &WaitGroup{
|
|
wg: wg,
|
|
c: &(atomic.Int64{}),
|
|
drain: &(atomic.Bool{}),
|
|
}
|
|
return g
|
|
}
|
|
|
|
func NewWaitGroup() *WaitGroup {
|
|
var wg sync.WaitGroup
|
|
g := &WaitGroup{
|
|
wg: &wg,
|
|
c: &(atomic.Int64{}),
|
|
drain: &(atomic.Bool{}),
|
|
}
|
|
return g
|
|
}
|
|
|
|
func (g *WaitGroup) Add(n int) {
|
|
g.c.Add(int64(n))
|
|
g.wg.Add(n)
|
|
}
|
|
|
|
func (g *WaitGroup) Done() {
|
|
if g.drain.Load() {
|
|
return
|
|
}
|
|
g.c.Add(int64(-1))
|
|
g.wg.Done()
|
|
}
|
|
|
|
func (g *WaitGroup) Wait() {
|
|
g.wg.Wait()
|
|
}
|
|
|
|
func (g *WaitGroup) WaitContext(ctx context.Context) {
|
|
done := make(chan struct{})
|
|
go func() {
|
|
g.wg.Wait()
|
|
close(done)
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
g.drain.Swap(true)
|
|
for g.c.Load() > 0 {
|
|
select {
|
|
case <-done:
|
|
g.drain.Swap(false)
|
|
return
|
|
default:
|
|
g.wg.Done()
|
|
}
|
|
}
|
|
case <-done:
|
|
return
|
|
}
|
|
}
|