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
|
package kafka
import (
"context"
"sync"
)
// runGroup is a collection of goroutines working together. If any one goroutine
// stops, then all goroutines will be stopped.
//
// A zero runGroup is valid
type runGroup struct {
initOnce sync.Once
ctx context.Context
cancel context.CancelFunc
wg sync.WaitGroup
}
func (r *runGroup) init() {
if r.cancel == nil {
r.ctx, r.cancel = context.WithCancel(context.Background())
}
}
func (r *runGroup) WithContext(ctx context.Context) *runGroup {
ctx, cancel := context.WithCancel(ctx)
return &runGroup{
ctx: ctx,
cancel: cancel,
}
}
// Wait blocks until all function calls have returned.
func (r *runGroup) Wait() {
r.wg.Wait()
}
// Stop stops the goroutines and waits for them to complete
func (r *runGroup) Stop() {
r.initOnce.Do(r.init)
r.cancel()
r.Wait()
}
// Go calls the given function in a new goroutine.
//
// The first call to return a non-nil error cancels the group; its error will be
// returned by Wait.
func (r *runGroup) Go(f func(stop <-chan struct{})) {
r.initOnce.Do(r.init)
r.wg.Add(1)
go func() {
defer r.wg.Done()
defer r.cancel()
f(r.ctx.Done())
}()
}
|
