-
Notifications
You must be signed in to change notification settings - Fork 2.1k
/
error_group.go
106 lines (93 loc) · 2.96 KB
/
error_group.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
/*
Copyright 2021 The Vitess Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
package concurrency
import "context"
// ErrorGroup provides a function for waiting for N goroutines to complete with
// at least X successes and no more than Y failures, and cancelling the rest.
//
// It should be used as follows:
//
// errCh := make(chan error)
// errgroupCtx, errgroupCancel := context.WithCancel(ctx)
//
// for _, arg := range args {
// arg := arg
//
// go func() {
// err := doWork(errGroupCtx, arg)
// errCh <- err
// }()
// }
//
// errgroup := concurrency.ErrorGroup{
// NumGoroutines: len(args),
// NumRequiredSuccess: 5, // need at least 5 to respond with nil error before cancelling the rest
// NumAllowedErrors: 1, // if more than 1 responds with non-nil error, cancel the rest
// }
// errRec := errgroup.Wait(errgroupCancel, errCh)
//
// if errRec.HasErrors() {
// // ...
// }
type ErrorGroup struct {
NumGoroutines int
NumRequiredSuccesses int
NumAllowedErrors int
}
// Wait waits for a group of goroutines that are sending errors to the given
// error channel, and are cancellable by the given cancel function.
//
// Wait will cancel any outstanding goroutines under the following conditions:
//
// (1) More than NumAllowedErrors non-nil results have been consumed on the
// error channel.
//
// (2) At least NumRequiredSuccesses nil results have been consumed on the error
// channel.
//
// After the cancellation condition is triggered, Wait will continue to consume
// results off the error channel so as to not permanently block any of those
// cancelled goroutines.
//
// When finished consuming results from all goroutines, cancelled or otherwise,
// Wait returns an AllErrorRecorder that contains all errors returned by any of
// those goroutines. It does not close the error channel.
func (eg ErrorGroup) Wait(cancel context.CancelFunc, errors chan error) *AllErrorRecorder {
errCounter := 0
successCounter := 0
responseCounter := 0
rec := &AllErrorRecorder{}
if eg.NumGoroutines < 1 {
return rec
}
for err := range errors {
responseCounter++
switch err {
case nil:
successCounter++
default:
errCounter++
rec.RecordError(err)
}
// Even though we cancel in the next conditional, we need to keep
// consuming off the channel, or those goroutines will get stuck
// forever.
if responseCounter == eg.NumGoroutines {
break
}
if errCounter > eg.NumAllowedErrors || successCounter >= eg.NumRequiredSuccesses {
cancel()
}
}
return rec
}