Files
NTrace-core/taskgroup/taskgroup.go
2022-04-24 10:07:01 +08:00

46 lines
657 B
Go

package taskgroup
import (
"sync"
)
type TaskGroup struct {
count int
mu sync.Mutex
done []chan struct{}
}
func New() *TaskGroup {
return &TaskGroup{
count: 0,
mu: sync.Mutex{},
done: []chan struct{}{},
}
}
func (t *TaskGroup) Add() {
t.mu.Lock()
defer t.mu.Unlock()
t.count++
}
func (t *TaskGroup) Done() {
t.mu.Lock()
defer t.mu.Unlock()
if t.count-1 == 0 {
for _, doneChannel := range t.done {
doneChannel <- struct{}{}
}
t.done = []chan struct{}{}
}
t.count--
}
func (t *TaskGroup) Wait() {
doneChannel := make(chan struct{})
t.mu.Lock()
t.done = append(t.done, doneChannel)
t.mu.Unlock()
<-doneChannel
}