diff options
Diffstat (limited to 'src/pkg/sync/waitgroup.go')
-rw-r--r-- | src/pkg/sync/waitgroup.go | 97 |
1 files changed, 97 insertions, 0 deletions
diff --git a/src/pkg/sync/waitgroup.go b/src/pkg/sync/waitgroup.go new file mode 100644 index 000000000..a4c9b7e43 --- /dev/null +++ b/src/pkg/sync/waitgroup.go @@ -0,0 +1,97 @@ +// Copyright 2011 The Go Authors. All rights reserved. +// Use of this source code is governed by a BSD-style +// license that can be found in the LICENSE file. + +package sync + +import ( + "runtime" + "sync/atomic" +) + +// A WaitGroup waits for a collection of goroutines to finish. +// The main goroutine calls Add to set the number of +// goroutines to wait for. Then each of the goroutines +// runs and calls Done when finished. At the same time, +// Wait can be used to block until all goroutines have finished. +// +// For example: +// +// for i := 0; i < n; i++ { +// if !condition(i) { +// continue +// } +// wg.Add(1) +// go func() { +// // Do something. +// wg.Done() +// }() +// } +// wg.Wait() +// +type WaitGroup struct { + m Mutex + counter int32 + waiters int32 + sema *uint32 +} + +// WaitGroup creates a new semaphore each time the old semaphore +// is released. This is to avoid the following race: +// +// G1: Add(1) +// G1: go G2() +// G1: Wait() // Context switch after Unlock() and before Semacquire(). +// G2: Done() // Release semaphore: sema == 1, waiters == 0. G1 doesn't run yet. +// G3: Wait() // Finds counter == 0, waiters == 0, doesn't block. +// G3: Add(1) // Makes counter == 1, waiters == 0. +// G3: go G4() +// G3: Wait() // G1 still hasn't run, G3 finds sema == 1, unblocked! Bug. + +// Add adds delta, which may be negative, to the WaitGroup counter. +// If the counter becomes zero, all goroutines blocked on Wait() are released. +func (wg *WaitGroup) Add(delta int) { + v := atomic.AddInt32(&wg.counter, int32(delta)) + if v < 0 { + panic("sync: negative WaitGroup count") + } + if v > 0 || atomic.LoadInt32(&wg.waiters) == 0 { + return + } + wg.m.Lock() + for i := int32(0); i < wg.waiters; i++ { + runtime.Semrelease(wg.sema) + } + wg.waiters = 0 + wg.sema = nil + wg.m.Unlock() +} + +// Done decrements the WaitGroup counter. +func (wg *WaitGroup) Done() { + wg.Add(-1) +} + +// Wait blocks until the WaitGroup counter is zero. +func (wg *WaitGroup) Wait() { + if atomic.LoadInt32(&wg.counter) == 0 { + return + } + wg.m.Lock() + atomic.AddInt32(&wg.waiters, 1) + // This code is racing with the unlocked path in Add above. + // The code above modifies counter and then reads waiters. + // We must modify waiters and then read counter (the opposite order) + // to avoid missing an Add. + if atomic.LoadInt32(&wg.counter) == 0 { + atomic.AddInt32(&wg.waiters, -1) + wg.m.Unlock() + return + } + if wg.sema == nil { + wg.sema = new(uint32) + } + s := wg.sema + wg.m.Unlock() + runtime.Semacquire(s) +} |