54 строки
1,1 КиБ
Go
54 строки
1,1 КиБ
Go
package sync
|
|
|
|
import "internal/task"
|
|
|
|
type WaitGroup struct {
|
|
counter uint
|
|
waiters task.Stack
|
|
}
|
|
|
|
func (wg *WaitGroup) Add(delta int) {
|
|
if delta > 0 {
|
|
// Check for overflow.
|
|
if uint(delta) > (^uint(0))-wg.counter {
|
|
panic("sync: WaitGroup counter overflowed")
|
|
}
|
|
|
|
// Add to the counter.
|
|
wg.counter += uint(delta)
|
|
} else {
|
|
// Check for underflow.
|
|
if uint(-delta) > wg.counter {
|
|
panic("sync: negative WaitGroup counter")
|
|
}
|
|
|
|
// Subtract from the counter.
|
|
wg.counter -= uint(-delta)
|
|
|
|
// If the counter is zero, everything is done and the waiters should be resumed.
|
|
// This code assumes that the waiters cannot wake up until after this function returns.
|
|
// In the current implementation, this is always correct.
|
|
if wg.counter == 0 {
|
|
for t := wg.waiters.Pop(); t != nil; t = wg.waiters.Pop() {
|
|
scheduleTask(t)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func (wg *WaitGroup) Done() {
|
|
wg.Add(-1)
|
|
}
|
|
|
|
func (wg *WaitGroup) Wait() {
|
|
if wg.counter == 0 {
|
|
// Everything already finished.
|
|
return
|
|
}
|
|
|
|
// Push the current goroutine onto the waiter stack.
|
|
wg.waiters.Push(task.Current())
|
|
|
|
// Pause until the waiters are awoken by Add/Done.
|
|
task.Pause()
|
|
}
|