sdk/go/arvados/contextgroup.go

   1 package arvados
   2
   3 import (
   4         "context"
   5         "sync"
   6 )
   7
   8 // A contextGroup is a context-aware variation on sync.WaitGroup. It
   9 // provides a child context for the added funcs to use, so they can
  10 // exit early if another added func returns an error. Its Wait()
  11 // method returns the first error returned by any added func.
  12 //
  13 // Example:
  14 //
  15 //      err := errors.New("oops")
  16 //      cg := newContextGroup()
  17 //      defer cg.Cancel()
  18 //      cg.Go(func() error {
  19 //              someFuncWithContext(cg.Context())
  20 //              return nil
  21 //      })
  22 //      cg.Go(func() error {
  23 //              return err // this cancels cg.Context()
  24 //      })
  25 //      return cg.Wait() // returns err after both goroutines have ended
  26 type contextGroup struct {
  27         ctx    context.Context
  28         cancel context.CancelFunc
  29         wg     sync.WaitGroup
  30         err    error
  31         mtx    sync.Mutex
  32 }
  33
  34 // newContextGroup returns a new contextGroup. The caller must
  35 // eventually call the Cancel() method of the returned contextGroup.
  36 func newContextGroup(ctx context.Context) *contextGroup {
  37         ctx, cancel := context.WithCancel(ctx)
  38         return &contextGroup{
  39                 ctx:    ctx,
  40                 cancel: cancel,
  41         }
  42 }
  43
  44 // Cancel cancels the context group.
  45 func (cg *contextGroup) Cancel() {
  46         cg.cancel()
  47 }
  48
  49 // Context returns a context.Context which will be canceled when all
  50 // funcs have succeeded or one has failed.
  51 func (cg *contextGroup) Context() context.Context {
  52         return cg.ctx
  53 }
  54
  55 // Go calls f in a new goroutine. If f returns an error, the
  56 // contextGroup is canceled.
  57 //
  58 // If f notices cg.Context() is done, it should abandon further work
  59 // and return. In this case, f's return value will be ignored.
  60 func (cg *contextGroup) Go(f func() error) {
  61         cg.mtx.Lock()
  62         defer cg.mtx.Unlock()
  63         if cg.err != nil {
  64                 return
  65         }
  66         cg.wg.Add(1)
  67         go func() {
  68                 defer cg.wg.Done()
  69                 err := f()
  70                 cg.mtx.Lock()
  71                 defer cg.mtx.Unlock()
  72                 if err != nil && cg.err == nil {
  73                         cg.err = err
  74                         cg.cancel()
  75                 }
  76         }()
  77 }
  78
  79 // Wait waits for all added funcs to return, and returns the first
  80 // non-nil error.
  81 //
  82 // If the parent context is canceled before a func returns an error,
  83 // Wait returns the parent context's Err().
  84 //
  85 // Wait returns nil if all funcs return nil before the parent context
  86 // is canceled.
  87 func (cg *contextGroup) Wait() error {
  88         cg.wg.Wait()
  89         cg.mtx.Lock()
  90         defer cg.mtx.Unlock()
  91         if cg.err != nil {
  92                 return cg.err
  93         }
  94         return cg.ctx.Err()
  95 }