-
Notifications
You must be signed in to change notification settings - Fork 0
/
reaper.go
100 lines (83 loc) · 2.65 KB
/
reaper.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
package grim
import (
"sync"
"golang.org/x/net/context"
)
// GrimReaper is a task runner which will wait until all the tasks
// are complete until comtinuing. Reapers can also kill tasks prematurely if
// they are listening to the given context.
type GrimReaper interface {
// New creates a sub-reaper which is attached to the parent context. If
// the parent context is killed, so are the children. However, Wait must
// still be called for child reapers.
New() GrimReaper
// SpawnFunc starts a new goroutine for the given function. The task should
// return as soon as possible after the context completes.
SpawnFunc(TaskFunc)
// Spawn starts a new goroutine for the given Task. The task should
// return as soon as possible after the context completes.
Spawn(Task)
// Kill sends a message to all running tasks to stop. If child reapers
// have also been created, they will be triggered to stop as well.
Kill()
// Wait will block until all tasks have completed. This will NOT block
// until chlid reapers are finished. Each reaper must call wait
// independently.
Wait()
}
// Reaper returns an implementation of the GrimReaper interface.
func Reaper() GrimReaper {
var wg sync.WaitGroup
ctx, cancel := context.WithCancel(context.Background())
return &reaper{wg, ctx, cancel}
}
// ReaperWithContext creates a new GrimReaper implementation and uses the given
// context as the parent context.
func ReaperWithContext(c context.Context) GrimReaper {
var wg sync.WaitGroup
ctx, cancel := context.WithCancel(c)
return &reaper{wg, ctx, cancel}
}
// TaskFunc is a killable function which runs in a separate go routine. In order to fulfill the contract the function MUST listen to the context and exit if it fires.
type TaskFunc func(context.Context)
type Task interface {
Execute(context.Context)
}
// reaper implements the GrimReaper interface.
type reaper struct {
wg sync.WaitGroup
ctx context.Context
cancel context.CancelFunc
}
// SpawnFunc runs a task.
func (r *reaper) SpawnFunc(t TaskFunc) {
r.wg.Add(1)
c, _ := context.WithCancel(r.ctx)
go func(ctx context.Context) {
defer r.wg.Done()
t(ctx)
}(c)
}
// Spawn runs a task.
func (r *reaper) Spawn(t Task) {
r.wg.Add(1)
c, _ := context.WithCancel(r.ctx)
go func(ctx context.Context) {
defer r.wg.Done()
t.Execute(ctx)
}(c)
}
// Wait blocks until all tasks have completed.
func (r *reaper) Wait() {
r.wg.Wait()
}
// Kill cancels the context and waits for all tasks to exit.
func (r *reaper) Kill() {
r.cancel()
r.Wait()
}
// New creates a new reaper with the current context as the parent context of
// the new reaper.
func (r *reaper) New() GrimReaper {
return ReaperWithContext(r.ctx)
}