-
Notifications
You must be signed in to change notification settings - Fork 221
/
task.go
69 lines (53 loc) · 1.32 KB
/
task.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
// Package task implements async task handling.
package task
import (
"context"
"sync"
)
type contextKey struct{}
// NewContext derives a Context that carries the given Manager from ctx.
func NewContext(ctx context.Context, m Manager) context.Context {
return context.WithValue(ctx, contextKey{}, m)
}
// FromContext returns the Manager ctx carries. It panics in case ctx carries
// no Manager.
func FromContext(ctx context.Context) Manager {
return ctx.Value(contextKey{}).(Manager)
}
// New initializes and returns a Manager which runs its tasks on the chain
// of the given parent context.
func New(parent context.Context) Manager {
ctx, cancel := context.WithCancel(parent)
return &manager{
Context: ctx,
CancelFunc: cancel,
}
}
// Task wraps the set of tasks.
type Task func(context.Context)
// Manager implements a task manager.
type Manager interface {
pkg() // internal
// Run runs the task in its own goroutine.
Run(Task)
// Shutdown instructs all the tasks to shutdown and waits until they've done
// so.
Shutdown()
}
type manager struct {
context.Context
context.CancelFunc
sync.WaitGroup
}
func (*manager) pkg() {}
func (m *manager) Run(t Task) {
m.WaitGroup.Add(1)
go func() {
defer m.WaitGroup.Done()
t(m.Context)
}()
}
func (m *manager) Shutdown() {
m.CancelFunc()
m.WaitGroup.Wait()
}