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