-
Notifications
You must be signed in to change notification settings - Fork 4.4k
/
routine.go
173 lines (140 loc) · 3.74 KB
/
routine.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
package routine
import (
"context"
"os"
"sync"
"github.com/hashicorp/go-hclog"
)
type Routine func(ctx context.Context) error
// cancelCh is the ctx.Done()
// When cancel() is called, if the routine is running a blocking call (e.g. some ACL replication RPCs),
// stoppedCh won't be closed till the blocking call returns, while cancelCh will be closed immediately.
// cancelCh is used to properly detect routine running status between cancel() and close(stoppedCh)
type routineTracker struct {
cancel context.CancelFunc
cancelCh <-chan struct{} // closed when ctx is done
stoppedCh chan struct{} // closed when no longer running
}
func (r *routineTracker) running() bool {
select {
case <-r.stoppedCh:
return false
case <-r.cancelCh:
return false
default:
return true
}
}
func (r *routineTracker) wait() {
<-r.stoppedCh
}
type Manager struct {
lock sync.RWMutex
logger hclog.Logger
routines map[string]*routineTracker
}
func NewManager(logger hclog.Logger) *Manager {
if logger == nil {
logger = hclog.New(&hclog.LoggerOptions{
Output: os.Stderr,
})
}
return &Manager{
logger: logger,
routines: make(map[string]*routineTracker),
}
}
func (m *Manager) IsRunning(name string) bool {
m.lock.Lock()
defer m.lock.Unlock()
if routine, ok := m.routines[name]; ok {
return routine.running()
}
return false
}
func (m *Manager) Start(ctx context.Context, name string, routine Routine) error {
m.lock.Lock()
defer m.lock.Unlock()
if instance, ok := m.routines[name]; ok && instance.running() {
return nil
}
if ctx == nil {
ctx = context.Background()
}
rtCtx, cancel := context.WithCancel(ctx)
instance := &routineTracker{
cancel: cancel,
cancelCh: ctx.Done(),
stoppedCh: make(chan struct{}),
}
go m.execute(rtCtx, name, routine, instance.stoppedCh)
m.routines[name] = instance
m.logger.Info("started routine", "routine", name)
return nil
}
// execute will run the given routine in the foreground and close the given channel when its done executing
func (m *Manager) execute(ctx context.Context, name string, routine Routine, done chan struct{}) {
defer func() {
close(done)
}()
err := routine(ctx)
if err != nil && err != context.DeadlineExceeded && err != context.Canceled {
m.logger.Error("routine exited with error",
"routine", name,
"error", err,
)
} else {
m.logger.Info("stopped routine", "routine", name)
}
}
// Caveat: The returned stoppedCh indicates that the routine is completed
// It's possible that ctx is canceled, but stoppedCh not yet closed
// Use mgr.IsRunning(name) than this stoppedCh to tell whether the
// instance is still running (not cancelled or completed).
func (m *Manager) Stop(name string) <-chan struct{} {
instance := m.stopInstance(name)
if instance == nil {
// Fabricate a closed channel so it won't block forever.
ch := make(chan struct{})
close(ch)
return ch
}
return instance.stoppedCh
}
func (m *Manager) stopInstance(name string) *routineTracker {
m.lock.Lock()
defer m.lock.Unlock()
instance, ok := m.routines[name]
if !ok {
// no running instance
return nil
}
if !instance.running() {
return instance
}
m.logger.Info("stopping routine", "routine", name)
instance.cancel()
delete(m.routines, name)
return instance
}
// StopAll goroutines. Once StopAll is called, it is no longer safe to add no
// goroutines to the Manager.
func (m *Manager) StopAll() {
m.lock.Lock()
defer m.lock.Unlock()
for name, routine := range m.routines {
if !routine.running() {
continue
}
m.logger.Info("stopping routine", "routine", name)
routine.cancel()
}
}
// Wait for all goroutines to stop after StopAll is called.
func (m *Manager) Wait() {
m.lock.Lock()
defer m.lock.Unlock()
for _, routine := range m.routines {
routine.wait()
}
}