This repository has been archived by the owner on Sep 25, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathgroup.go
212 lines (188 loc) · 5.05 KB
/
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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
package parallel
import (
"context"
"errors"
"fmt"
"sync"
"sync/atomic"
)
var nextTaskID int64 = 0x0bace1d000000000
// Group is a facility for running a task with several subtasks without
// inversion of control. For most ordinary use cases, use Run instead.
//
// return Run(ctx, start)
//
// ...is equivalent to:
//
// g := NewGroup(ctx)
// if err := start(g.Context(), g.Spawn); err != nil {
// g.Exit(err)
// }
// return g.Wait()
//
// Group is mostly useful in test suites where starting and finishing the group
// is controlled by test setup and teardown functions.
type Group struct {
ctx context.Context
cancel context.CancelFunc
mu sync.Mutex
running int
done chan struct{}
closing bool
err error
}
// NewGroup creates a new Group controlled by the given context
func NewGroup(ctx context.Context) *Group {
g := new(Group)
g.ctx, g.cancel = context.WithCancel(ctx)
g.done = make(chan struct{})
close(g.done)
return g
}
// NewSubgroup creates a new Group nested within another. The spawn argument is
// the spawn function of the parent group.
//
// The subgroup's context is inherited from the parent group. The entire
// subgroup is treated as a task in the parent group.
//
// Example within parallel.Run:
//
// err := parallel.Run(ctx, func(ctx context.Context, spawn parallel.SpawnFn) error {
// spawn(...)
// spawn(...)
// subgroup := parallel.NewSubgroup(spawn, "updater")
// subgroup.Spawn(...)
// subgroup.Spawn(...)
// return nil
// })
//
// Example within an explicit group:
//
// group := parallel.NewGroup(ctx)
// group.Spawn(...)
// group.Spawn(...)
// subgroup := parallel.NewSubgroup(group.Spawn, "updater")
// subgroup.Spawn(...)
// subgroup.Spawn(...)
//
func NewSubgroup(spawn SpawnFn, name string, onExit OnExit) *Group {
ch := make(chan *Group)
spawn(name, onExit, func(ctx context.Context) error {
g := NewGroup(ctx)
ch <- g
return g.Complete(ctx)
})
return <-ch
}
// Context returns the inner context of the group which controls the lifespan of
// its subtasks
func (g *Group) Context() context.Context {
return g.ctx
}
// Spawn spawns a subtask. See documentation for SpawnFn.
//
// When a subtask finishes, it sets the result of the group if it's not already
// set (unless the task returns nil and its OnExit mode is Continue).
func (g *Group) Spawn(name string, onExit OnExit, task Task) {
id := atomic.AddInt64(&nextTaskID, 1)
g.mu.Lock()
if g.running == 0 {
g.done = make(chan struct{})
}
g.running++
g.mu.Unlock()
go g.runTask(g.ctx, id, name, onExit, task)
}
// Second parameter is the task ID. It is ignored because the only reason to
// pass it is to add it to the stack trace
func (g *Group) runTask(ctx context.Context, _ int64, name string, onExit OnExit, task Task) {
err := runTask(ctx, task)
g.mu.Lock()
defer g.mu.Unlock()
if err != nil {
g.exit(err)
} else if !g.closing {
switch onExit {
case Continue:
case Exit:
g.exit(nil)
case Fail:
g.exit(fmt.Errorf("task %s terminated unexpectedly", name))
default:
g.exit(fmt.Errorf("task %s: %v", name, onExit))
}
}
g.running--
if g.running == 0 {
close(g.done)
}
}
func (g *Group) exit(err error) {
// Cancellations during shutdown are fine
if g.closing && errors.Is(err, context.Canceled) {
return
}
if g.err == nil {
g.err = err
}
if !g.closing {
g.closing = true
g.cancel()
}
}
// Exit prompts the group to shut down, if it's not already shutting down or
// finished. This causes the inner context to close, which should prompt any
// running subtasks to exit. Use Wait to block until all the subtasks actually
// finish.
//
// If the group result is not yet set, Exit sets it to err.
func (g *Group) Exit(err error) {
g.mu.Lock()
defer g.mu.Unlock()
g.exit(err)
}
// Running returns the number of running subtasks
func (g *Group) Running() int {
g.mu.Lock()
defer g.mu.Unlock()
return g.running
}
// Done returns a channel that closes when the last running subtask finishes. If
// no subtasks are running, the returned channel is already closed.
func (g *Group) Done() <-chan struct{} {
g.mu.Lock()
defer g.mu.Unlock()
return g.done
}
// Wait blocks until no subtasks are running, then returns the group result.
//
// The group result is set by finishing subtasks (see the documentation for
// OnExit modes) as well as by Exit calls.
func (g *Group) Wait() error {
<-g.Done()
return g.err
}
// Complete first waits for either the given context to close or the group to
// exit on its own, then for the group's remaining subtasks to finish.
//
// Returns the group result. If the group result is nil, returns the error from
// the given context so as to not confuse parallel.Fail if the group is empty.
//
// This is a convenience method useful when attaching a subgroup:
//
// spawn("subgroup", parallel.Fail, subgroup.Complete)
//
// ...or:
//
// group.Spawn("subgroup", parallel.Fail, subgroup.Complete)
//
func (g *Group) Complete(ctx context.Context) error {
select {
case <-ctx.Done():
case <-g.ctx.Done():
}
if err := g.Wait(); err != nil {
return err
}
return ctx.Err()
}