1// Copyright 2016 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 errgroup provides synchronization, error propagation, and Context
  6// cancelation for groups of goroutines working on subtasks of a common task.
  7//
  8// [errgroup.Group] is related to [sync.WaitGroup] but adds handling of tasks
  9// returning errors.
 10package errgroup
 11
 12import (
 13	"context"
 14	"fmt"
 15	"sync"
 16)
 17
 18type token struct{}
 19
 20// A Group is a collection of goroutines working on subtasks that are part of
 21// the same overall task. A Group should not be reused for different tasks.
 22//
 23// A zero Group is valid, has no limit on the number of active goroutines,
 24// and does not cancel on error.
 25type Group struct {
 26	cancel func(error)
 27
 28	wg sync.WaitGroup
 29
 30	sem chan token
 31
 32	errOnce sync.Once
 33	err     error
 34}
 35
 36func (g *Group) done() {
 37	if g.sem != nil {
 38		<-g.sem
 39	}
 40	g.wg.Done()
 41}
 42
 43// WithContext returns a new Group and an associated Context derived from ctx.
 44//
 45// The derived Context is canceled the first time a function passed to Go
 46// returns a non-nil error or the first time Wait returns, whichever occurs
 47// first.
 48func WithContext(ctx context.Context) (*Group, context.Context) {
 49	ctx, cancel := context.WithCancelCause(ctx)
 50	return &Group{cancel: cancel}, ctx
 51}
 52
 53// Wait blocks until all function calls from the Go method have returned, then
 54// returns the first non-nil error (if any) from them.
 55func (g *Group) Wait() error {
 56	g.wg.Wait()
 57	if g.cancel != nil {
 58		g.cancel(g.err)
 59	}
 60	return g.err
 61}
 62
 63// Go calls the given function in a new goroutine.
 64// The first call to Go must happen before a Wait.
 65// It blocks until the new goroutine can be added without the number of
 66// active goroutines in the group exceeding the configured limit.
 67//
 68// The first call to return a non-nil error cancels the group's context, if the
 69// group was created by calling WithContext. The error will be returned by Wait.
 70func (g *Group) Go(f func() error) {
 71	if g.sem != nil {
 72		g.sem <- token{}
 73	}
 74
 75	g.wg.Add(1)
 76	go func() {
 77		defer g.done()
 78
 79		if err := f(); err != nil {
 80			g.errOnce.Do(func() {
 81				g.err = err
 82				if g.cancel != nil {
 83					g.cancel(g.err)
 84				}
 85			})
 86		}
 87	}()
 88}
 89
 90// TryGo calls the given function in a new goroutine only if the number of
 91// active goroutines in the group is currently below the configured limit.
 92//
 93// The return value reports whether the goroutine was started.
 94func (g *Group) TryGo(f func() error) bool {
 95	if g.sem != nil {
 96		select {
 97		case g.sem <- token{}:
 98			// Note: this allows barging iff channels in general allow barging.
 99		default:
100			return false
101		}
102	}
103
104	g.wg.Add(1)
105	go func() {
106		defer g.done()
107
108		if err := f(); err != nil {
109			g.errOnce.Do(func() {
110				g.err = err
111				if g.cancel != nil {
112					g.cancel(g.err)
113				}
114			})
115		}
116	}()
117	return true
118}
119
120// SetLimit limits the number of active goroutines in this group to at most n.
121// A negative value indicates no limit.
122// A limit of zero will prevent any new goroutines from being added.
123//
124// Any subsequent call to the Go method will block until it can add an active
125// goroutine without exceeding the configured limit.
126//
127// The limit must not be modified while any goroutines in the group are active.
128func (g *Group) SetLimit(n int) {
129	if n < 0 {
130		g.sem = nil
131		return
132	}
133	if len(g.sem) != 0 {
134		panic(fmt.Errorf("errgroup: modify limit while %v goroutines in the group are still active", len(g.sem)))
135	}
136	g.sem = make(chan token, n)
137}