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 (
   8         "sync/atomic"
   9         "unsafe"
  10 )
  11
  12 // A WaitGroup waits for a collection of goroutines to finish.
  13 // The main goroutine calls Add to set the number of
  14 // goroutines to wait for.  Then each of the goroutines
  15 // runs and calls Done when finished.  At the same time,
  16 // Wait can be used to block until all goroutines have finished.
  17 type WaitGroup struct {
  18         m       Mutex
  19         counter int32
  20         waiters int32
  21         sema    *uint32
  22 }
  23
  24 // WaitGroup creates a new semaphore each time the old semaphore
  25 // is released. This is to avoid the following race:
  26 //
  27 // G1: Add(1)
  28 // G1: go G2()
  29 // G1: Wait() // Context switch after Unlock() and before Semacquire().
  30 // G2: Done() // Release semaphore: sema == 1, waiters == 0. G1 doesn't run yet.
  31 // G3: Wait() // Finds counter == 0, waiters == 0, doesn't block.
  32 // G3: Add(1) // Makes counter == 1, waiters == 0.
  33 // G3: go G4()
  34 // G3: Wait() // G1 still hasn't run, G3 finds sema == 1, unblocked! Bug.
  35
  36 // Add adds delta, which may be negative, to the WaitGroup counter.
  37 // If the counter becomes zero, all goroutines blocked on Wait are released.
  38 // If the counter goes negative, Add panics.
  39 //
  40 // Note that calls with positive delta must happen before the call to Wait,
  41 // or else Wait may wait for too small a group. Typically this means the calls
  42 // to Add should execute before the statement creating the goroutine or
  43 // other event to be waited for. See the WaitGroup example.
  44 func (wg *WaitGroup) Add(delta int) {
  45         if raceenabled {
  46                 _ = wg.m.state // trigger nil deref early
  47                 if delta < 0 {
  48                         // Synchronize decrements with Wait.
  49                         raceReleaseMerge(unsafe.Pointer(wg))
  50                 }
  51                 raceDisable()
  52                 defer raceEnable()
  53         }
  54         v := atomic.AddInt32(&wg.counter, int32(delta))
  55         if raceenabled {
  56                 if delta > 0 && v == int32(delta) {
  57                         // The first increment must be synchronized with Wait.
  58                         // Need to model this as a read, because there can be
  59                         // several concurrent wg.counter transitions from 0.
  60                         raceRead(unsafe.Pointer(&wg.sema))
  61                 }
  62         }
  63         if v < 0 {
  64                 panic("sync: negative WaitGroup counter")
  65         }
  66         if v > 0 || atomic.LoadInt32(&wg.waiters) == 0 {
  67                 return
  68         }
  69         wg.m.Lock()
  70         for i := int32(0); i < wg.waiters; i++ {
  71                 runtime_Semrelease(wg.sema)
  72         }
  73         wg.waiters = 0
  74         wg.sema = nil
  75         wg.m.Unlock()
  76 }
  77
  78 // Done decrements the WaitGroup counter.
  79 func (wg *WaitGroup) Done() {
  80         wg.Add(-1)
  81 }
  82
  83 // Wait blocks until the WaitGroup counter is zero.
  84 func (wg *WaitGroup) Wait() {
  85         if raceenabled {
  86                 _ = wg.m.state // trigger nil deref early
  87                 raceDisable()
  88         }
  89         if atomic.LoadInt32(&wg.counter) == 0 {
  90                 if raceenabled {
  91                         raceEnable()
  92                         raceAcquire(unsafe.Pointer(wg))
  93                 }
  94                 return
  95         }
  96         wg.m.Lock()
  97         w := atomic.AddInt32(&wg.waiters, 1)
  98         // This code is racing with the unlocked path in Add above.
  99         // The code above modifies counter and then reads waiters.
 100         // We must modify waiters and then read counter (the opposite order)
 101         // to avoid missing an Add.
 102         if atomic.LoadInt32(&wg.counter) == 0 {
 103                 atomic.AddInt32(&wg.waiters, -1)
 104                 if raceenabled {
 105                         raceEnable()
 106                         raceAcquire(unsafe.Pointer(wg))
 107                         raceDisable()
 108                 }
 109                 wg.m.Unlock()
 110                 if raceenabled {
 111                         raceEnable()
 112                 }
 113                 return
 114         }
 115         if raceenabled && w == 1 {
 116                 // Wait must be synchronized with the first Add.
 117                 // Need to model this is as a write to race with the read in Add.
 118                 // As a consequence, can do the write only for the first waiter,
 119                 // otherwise concurrent Waits will race with each other.
 120                 raceWrite(unsafe.Pointer(&wg.sema))
 121         }
 122         if wg.sema == nil {
 123                 wg.sema = new(uint32)
 124         }
 125         s := wg.sema
 126         wg.m.Unlock()
 127         runtime_Semacquire(s)
 128         if raceenabled {
 129                 raceEnable()
 130                 raceAcquire(unsafe.Pointer(wg))
 131         }
 132 }