micro/sync/waitgroup.go

70 lines
871 B
Go
Raw Permalink Normal View History

package sync
import (
"context"
"sync"
)
type WaitGroup struct {
wg *sync.WaitGroup
c int
mu sync.Mutex
}
func WrapWaitGroup(wg *sync.WaitGroup) *WaitGroup {
g := &WaitGroup{
wg: wg,
}
return g
}
func NewWaitGroup() *WaitGroup {
var wg sync.WaitGroup
return WrapWaitGroup(&wg)
}
func (g *WaitGroup) Add(n int) {
g.mu.Lock()
g.c += n
g.wg.Add(n)
g.mu.Unlock()
}
func (g *WaitGroup) Done() {
g.mu.Lock()
g.c += -1
g.wg.Add(-1)
g.mu.Unlock()
}
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.mu.Lock()
g.wg.Add(-g.c)
<-done
g.wg.Add(g.c)
g.mu.Unlock()
return
case <-done:
return
}
}
func (g *WaitGroup) Waiters() int {
g.mu.Lock()
c := g.c
g.mu.Unlock()
return c
}