|
|
@@ -8,22 +8,35 @@ package errgroup
|
|
|
|
|
|
import (
|
|
|
"context"
|
|
|
+ "fmt"
|
|
|
"sync"
|
|
|
)
|
|
|
|
|
|
+type token struct{}
|
|
|
+
|
|
|
// A Group is a collection of goroutines working on subtasks that are part of
|
|
|
// the same overall task.
|
|
|
//
|
|
|
-// A zero Group is valid and does not cancel on error.
|
|
|
+// A zero Group is valid, has no limit on the number of active goroutines,
|
|
|
+// and does not cancel on error.
|
|
|
type Group struct {
|
|
|
cancel func()
|
|
|
|
|
|
wg sync.WaitGroup
|
|
|
|
|
|
+ sem chan token
|
|
|
+
|
|
|
errOnce sync.Once
|
|
|
err error
|
|
|
}
|
|
|
|
|
|
+func (g *Group) done() {
|
|
|
+ if g.sem != nil {
|
|
|
+ <-g.sem
|
|
|
+ }
|
|
|
+ g.wg.Done()
|
|
|
+}
|
|
|
+
|
|
|
// WithContext returns a new Group and an associated Context derived from ctx.
|
|
|
//
|
|
|
// The derived Context is canceled the first time a function passed to Go
|
|
|
@@ -45,14 +58,48 @@ func (g *Group) Wait() error {
|
|
|
}
|
|
|
|
|
|
// Go calls the given function in a new goroutine.
|
|
|
+// It blocks until the new goroutine can be added without the number of
|
|
|
+// active goroutines in the group exceeding the configured limit.
|
|
|
//
|
|
|
-// The first call to return a non-nil error cancels the group; its error will be
|
|
|
-// returned by Wait.
|
|
|
+// The first call to return a non-nil error cancels the group's context, if the
|
|
|
+// group was created by calling WithContext. The error will be returned by Wait.
|
|
|
func (g *Group) Go(f func() error) {
|
|
|
+ if g.sem != nil {
|
|
|
+ g.sem <- token{}
|
|
|
+ }
|
|
|
+
|
|
|
g.wg.Add(1)
|
|
|
+ go func() {
|
|
|
+ defer g.done()
|
|
|
+
|
|
|
+ if err := f(); err != nil {
|
|
|
+ g.errOnce.Do(func() {
|
|
|
+ g.err = err
|
|
|
+ if g.cancel != nil {
|
|
|
+ g.cancel()
|
|
|
+ }
|
|
|
+ })
|
|
|
+ }
|
|
|
+ }()
|
|
|
+}
|
|
|
+
|
|
|
+// TryGo calls the given function in a new goroutine only if the number of
|
|
|
+// active goroutines in the group is currently below the configured limit.
|
|
|
+//
|
|
|
+// The return value reports whether the goroutine was started.
|
|
|
+func (g *Group) TryGo(f func() error) bool {
|
|
|
+ if g.sem != nil {
|
|
|
+ select {
|
|
|
+ case g.sem <- token{}:
|
|
|
+ // Note: this allows barging iff channels in general allow barging.
|
|
|
+ default:
|
|
|
+ return false
|
|
|
+ }
|
|
|
+ }
|
|
|
|
|
|
+ g.wg.Add(1)
|
|
|
go func() {
|
|
|
- defer g.wg.Done()
|
|
|
+ defer g.done()
|
|
|
|
|
|
if err := f(); err != nil {
|
|
|
g.errOnce.Do(func() {
|
|
|
@@ -63,4 +110,23 @@ func (g *Group) Go(f func() error) {
|
|
|
})
|
|
|
}
|
|
|
}()
|
|
|
+ return true
|
|
|
+}
|
|
|
+
|
|
|
+// SetLimit limits the number of active goroutines in this group to at most n.
|
|
|
+// A negative value indicates no limit.
|
|
|
+//
|
|
|
+// Any subsequent call to the Go method will block until it can add an active
|
|
|
+// goroutine without exceeding the configured limit.
|
|
|
+//
|
|
|
+// The limit must not be modified while any goroutines in the group are active.
|
|
|
+func (g *Group) SetLimit(n int) {
|
|
|
+ if n < 0 {
|
|
|
+ g.sem = nil
|
|
|
+ return
|
|
|
+ }
|
|
|
+ if len(g.sem) != 0 {
|
|
|
+ panic(fmt.Errorf("errgroup: modify limit while %v goroutines in the group are still active", len(g.sem)))
|
|
|
+ }
|
|
|
+ g.sem = make(chan token, n)
|
|
|
}
|
Sebastiaan van Stijn