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  	"internal/race"
     9  	"sync/atomic"
    10  	"unsafe"
    11  )
    12  
    13  // A WaitGroup is a counting semaphore typically used to wait
    14  // for a group of goroutines to finish.
    15  //
    16  // The main goroutine calls [WaitGroup.Add] to set (or increase) the number of
    17  // goroutines to wait for. Then each of the goroutines
    18  // runs and calls [WaitGroup.Done] when finished. At the same time,
    19  // [WaitGroup.Wait] can be used to block until all goroutines have finished.
    20  //
    21  // This is a typical pattern of WaitGroup usage to
    22  // synchronize 3 goroutines, each calling the function f:
    23  //
    24  //	var wg sync.WaitGroup
    25  //	for range 3 {
    26  //	   wg.Add(1)
    27  //	   go func() {
    28  //	       defer wg.Done()
    29  //	       f()
    30  //	   }()
    31  //	}
    32  //	wg.Wait()
    33  //
    34  // For convenience, the [WaitGroup.Go] method simplifies this pattern to:
    35  //
    36  //	var wg sync.WaitGroup
    37  //	for range 3 {
    38  //	   wg.Go(f)
    39  //	}
    40  //	wg.Wait()
    41  //
    42  // A WaitGroup must not be copied after first use.
    43  //
    44  // In the terminology of [the Go memory model], a call to [WaitGroup.Done]
    45  // “synchronizes before” the return of any Wait call that it unblocks.
    46  //
    47  // [the Go memory model]: https://go.dev/ref/mem
    48  type WaitGroup struct {
    49  	noCopy noCopy
    50  
    51  	state atomic.Uint64 // high 32 bits are counter, low 32 bits are waiter count.
    52  	sema  uint32
    53  }
    54  
    55  // Add adds delta, which may be negative, to the [WaitGroup] counter.
    56  // If the counter becomes zero, all goroutines blocked on [WaitGroup.Wait] are released.
    57  // If the counter goes negative, Add panics.
    58  //
    59  // Note that calls with a positive delta that occur when the counter is zero
    60  // must happen before a Wait. Calls with a negative delta, or calls with a
    61  // positive delta that start when the counter is greater than zero, may happen
    62  // at any time.
    63  // Typically this means the calls to Add should execute before the statement
    64  // creating the goroutine or other event to be waited for.
    65  // If a WaitGroup is reused to wait for several independent sets of events,
    66  // new Add calls must happen after all previous Wait calls have returned.
    67  // See the WaitGroup example.
    68  func (wg *WaitGroup) Add(delta int) {
    69  	if race.Enabled {
    70  		if delta < 0 {
    71  			// Synchronize decrements with Wait.
    72  			race.ReleaseMerge(unsafe.Pointer(wg))
    73  		}
    74  		race.Disable()
    75  		defer race.Enable()
    76  	}
    77  	state := wg.state.Add(uint64(delta) << 32)
    78  	v := int32(state >> 32)
    79  	w := uint32(state)
    80  	if race.Enabled && delta > 0 && v == int32(delta) {
    81  		// The first increment must be synchronized with Wait.
    82  		// Need to model this as a read, because there can be
    83  		// several concurrent wg.counter transitions from 0.
    84  		race.Read(unsafe.Pointer(&wg.sema))
    85  	}
    86  	if v < 0 {
    87  		panic("sync: negative WaitGroup counter")
    88  	}
    89  	if w != 0 && delta > 0 && v == int32(delta) {
    90  		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
    91  	}
    92  	if v > 0 || w == 0 {
    93  		return
    94  	}
    95  	// This goroutine has set counter to 0 when waiters > 0.
    96  	// Now there can't be concurrent mutations of state:
    97  	// - Adds must not happen concurrently with Wait,
    98  	// - Wait does not increment waiters if it sees counter == 0.
    99  	// Still do a cheap sanity check to detect WaitGroup misuse.
   100  	if wg.state.Load() != state {
   101  		panic("sync: WaitGroup misuse: Add called concurrently with Wait")
   102  	}
   103  	// Reset waiters count to 0.
   104  	wg.state.Store(0)
   105  	for ; w != 0; w-- {
   106  		runtime_Semrelease(&wg.sema, false, 0)
   107  	}
   108  }
   109  
   110  // Done decrements the [WaitGroup] counter by one.
   111  func (wg *WaitGroup) Done() {
   112  	wg.Add(-1)
   113  }
   114  
   115  // Wait blocks until the [WaitGroup] counter is zero.
   116  func (wg *WaitGroup) Wait() {
   117  	if race.Enabled {
   118  		race.Disable()
   119  	}
   120  	for {
   121  		state := wg.state.Load()
   122  		v := int32(state >> 32)
   123  		w := uint32(state)
   124  		if v == 0 {
   125  			// Counter is 0, no need to wait.
   126  			if race.Enabled {
   127  				race.Enable()
   128  				race.Acquire(unsafe.Pointer(wg))
   129  			}
   130  			return
   131  		}
   132  		// Increment waiters count.
   133  		if wg.state.CompareAndSwap(state, state+1) {
   134  			if race.Enabled && w == 0 {
   135  				// Wait must be synchronized with the first Add.
   136  				// Need to model this is as a write to race with the read in Add.
   137  				// As a consequence, can do the write only for the first waiter,
   138  				// otherwise concurrent Waits will race with each other.
   139  				race.Write(unsafe.Pointer(&wg.sema))
   140  			}
   141  			runtime_SemacquireWaitGroup(&wg.sema)
   142  			if wg.state.Load() != 0 {
   143  				panic("sync: WaitGroup is reused before previous Wait has returned")
   144  			}
   145  			if race.Enabled {
   146  				race.Enable()
   147  				race.Acquire(unsafe.Pointer(wg))
   148  			}
   149  			return
   150  		}
   151  	}
   152  }
   153  
   154  // Go calls f in a new goroutine and adds that task to the WaitGroup.
   155  // When f returns, the task is removed from the WaitGroup.
   156  //
   157  // If the WaitGroup is empty, Go must happen before a [WaitGroup.Wait].
   158  // Typically, this simply means Go is called to start tasks before Wait is called.
   159  // If the WaitGroup is not empty, Go may happen at any time.
   160  // This means a goroutine started by Go may itself call Go.
   161  // If a WaitGroup is reused to wait for several independent sets of tasks,
   162  // new Go calls must happen after all previous Wait calls have returned.
   163  //
   164  // In the terminology of [the Go memory model](https://go.dev/ref/mem),
   165  // the return from f "synchronizes before" the return of any Wait call that it unblocks.
   166  func (wg *WaitGroup) Go(f func()) {
   167  	wg.Add(1)
   168  	go func() {
   169  		defer wg.Done()
   170  		f()
   171  	}()
   172  }
   173
View as plain text