waitgroup.go raw

   1  package conc
   2  
   3  import (
   4  	"sync"
   5  
   6  	"github.com/sourcegraph/conc/panics"
   7  )
   8  
   9  // NewWaitGroup creates a new WaitGroup.
  10  func NewWaitGroup() *WaitGroup {
  11  	return &WaitGroup{}
  12  }
  13  
  14  // WaitGroup is the primary building block for scoped concurrency.
  15  // Goroutines can be spawned in the WaitGroup with the Go method,
  16  // and calling Wait() will ensure that each of those goroutines exits
  17  // before continuing. Any panics in a child goroutine will be caught
  18  // and propagated to the caller of Wait().
  19  //
  20  // The zero value of WaitGroup is usable, just like sync.WaitGroup.
  21  // Also like sync.WaitGroup, it must not be copied after first use.
  22  type WaitGroup struct {
  23  	wg sync.WaitGroup
  24  	pc panics.Catcher
  25  }
  26  
  27  // Go spawns a new goroutine in the WaitGroup.
  28  func (h *WaitGroup) Go(f func()) {
  29  	h.wg.Add(1)
  30  	go func() {
  31  		defer h.wg.Done()
  32  		h.pc.Try(f)
  33  	}()
  34  }
  35  
  36  // Wait will block until all goroutines spawned with Go exit and will
  37  // propagate any panics spawned in a child goroutine.
  38  func (h *WaitGroup) Wait() {
  39  	h.wg.Wait()
  40  
  41  	// Propagate a panic if we caught one from a child goroutine.
  42  	h.pc.Repanic()
  43  }
  44  
  45  // WaitAndRecover will block until all goroutines spawned with Go exit and
  46  // will return a *panics.Recovered if one of the child goroutines panics.
  47  func (h *WaitGroup) WaitAndRecover() *panics.Recovered {
  48  	h.wg.Wait()
  49  
  50  	// Return a recovered panic if we caught one from a child goroutine.
  51  	return h.pc.Recovered()
  52  }
  53