libgo/go/sync/waitgroup.go

   1 // Copyright 2011 The Go Authors. All rights reserved.
   2 // Use of this source code is governed by a BSD-style
   3 // license that can be found in the LICENSE file.
   4
   5 package sync
   6
   7 import "sync/atomic"
   8
   9 // A WaitGroup waits for a collection of goroutines to finish.
  10 // The main goroutine calls Add to set the number of
  11 // goroutines to wait for.  Then each of the goroutines
  12 // runs and calls Done when finished.  At the same time,
  13 // Wait can be used to block until all goroutines have finished.
  14 type WaitGroup struct {
  15         m       Mutex
  16         counter int32
  17         waiters int32
  18         sema    *uint32
  19 }
  20
  21 // WaitGroup creates a new semaphore each time the old semaphore
  22 // is released. This is to avoid the following race:
  23 //
  24 // G1: Add(1)
  25 // G1: go G2()
  26 // G1: Wait() // Context switch after Unlock() and before Semacquire().
  27 // G2: Done() // Release semaphore: sema == 1, waiters == 0. G1 doesn't run yet.
  28 // G3: Wait() // Finds counter == 0, waiters == 0, doesn't block.
  29 // G3: Add(1) // Makes counter == 1, waiters == 0.
  30 // G3: go G4()
  31 // G3: Wait() // G1 still hasn't run, G3 finds sema == 1, unblocked! Bug.
  32
  33 // Add adds delta, which may be negative, to the WaitGroup counter.
  34 // If the counter becomes zero, all goroutines blocked on Wait() are released.
  35 // If the counter goes negative, Add panics.
  36 func (wg *WaitGroup) Add(delta int) {
  37         v := atomic.AddInt32(&wg.counter, int32(delta))
  38         if v < 0 {
  39                 panic("sync: negative WaitGroup counter")
  40         }
  41         if v > 0 || atomic.LoadInt32(&wg.waiters) == 0 {
  42                 return
  43         }
  44         wg.m.Lock()
  45         for i := int32(0); i < wg.waiters; i++ {
  46                 runtime_Semrelease(wg.sema)
  47         }
  48         wg.waiters = 0
  49         wg.sema = nil
  50         wg.m.Unlock()
  51 }
  52
  53 // Done decrements the WaitGroup counter.
  54 func (wg *WaitGroup) Done() {
  55         wg.Add(-1)
  56 }
  57
  58 // Wait blocks until the WaitGroup counter is zero.
  59 func (wg *WaitGroup) Wait() {
  60         if atomic.LoadInt32(&wg.counter) == 0 {
  61                 return
  62         }
  63         wg.m.Lock()
  64         atomic.AddInt32(&wg.waiters, 1)
  65         // This code is racing with the unlocked path in Add above.
  66         // The code above modifies counter and then reads waiters.
  67         // We must modify waiters and then read counter (the opposite order)
  68         // to avoid missing an Add.
  69         if atomic.LoadInt32(&wg.counter) == 0 {
  70                 atomic.AddInt32(&wg.waiters, -1)
  71                 wg.m.Unlock()
  72                 return
  73         }
  74         if wg.sema == nil {
  75                 wg.sema = new(uint32)
  76         }
  77         s := wg.sema
  78         wg.m.Unlock()
  79         runtime_Semacquire(s)
  80 }