@@ -54,16 +54,16 @@ func InitManager(ctx context.Context) {
54
54
})
55
55
}
56
56
57
- // CallbackWithContext is combined runnable and context to watch to see if the caller has finished
58
- type CallbackWithContext func (ctx context. Context , callback func ())
57
+ // WithCallback is a runnable to call when the caller has finished
58
+ type WithCallback func (callback func ())
59
59
60
60
// RunnableWithShutdownFns is a runnable with functions to run at shutdown and terminate
61
61
// After the callback to atShutdown is called and is complete, the main function must return.
62
62
// Similarly the callback function provided to atTerminate must return once termination is complete.
63
63
// Please note that use of the atShutdown and atTerminate callbacks will create go-routines that will wait till their respective signals
64
64
// - users must therefore be careful to only call these as necessary.
65
65
// If run is not expected to run indefinitely RunWithShutdownChan is likely to be more appropriate.
66
- type RunnableWithShutdownFns func (atShutdown , atTerminate func (context. Context , func ()))
66
+ type RunnableWithShutdownFns func (atShutdown , atTerminate func (func ()))
67
67
68
68
// RunWithShutdownFns takes a function that has both atShutdown and atTerminate callbacks
69
69
// After the callback to atShutdown is called and is complete, the main function must return.
@@ -80,17 +80,21 @@ func (g *Manager) RunWithShutdownFns(run RunnableWithShutdownFns) {
80
80
g .doShutdown ()
81
81
}
82
82
}()
83
- run (func (ctx context.Context , atShutdown func ()) {
84
- go func () {
85
- select {
86
- case <- g .IsShutdown ():
83
+ run (func (atShutdown func ()) {
84
+ g .lock .Lock ()
85
+ defer g .lock .Unlock ()
86
+ g .toRunAtShutdown = append (g .toRunAtShutdown ,
87
+ func () {
88
+ defer func () {
89
+ if err := recover (); err != nil {
90
+ log .Critical ("PANIC during RunWithShutdownFns: %v\n Stacktrace: %s" , err , log .Stack (2 ))
91
+ g .doShutdown ()
92
+ }
93
+ }()
87
94
atShutdown ()
88
- case <- ctx .Done ():
89
- return
90
- }
91
- }()
92
- }, func (ctx context.Context , atTerminate func ()) {
93
- g .RunAtTerminate (ctx , atTerminate )
95
+ })
96
+ }, func (atTerminate func ()) {
97
+ g .RunAtTerminate (atTerminate )
94
98
})
95
99
}
96
100
@@ -99,7 +103,7 @@ func (g *Manager) RunWithShutdownFns(run RunnableWithShutdownFns) {
99
103
// (Optionally IsHammer may be waited for instead however, this should be avoided if possible.)
100
104
// The callback function provided to atTerminate must return once termination is complete.
101
105
// Please note that use of the atTerminate function will create a go-routine that will wait till terminate - users must therefore be careful to only call this as necessary.
102
- type RunnableWithShutdownChan func (atShutdown <- chan struct {}, atTerminate CallbackWithContext )
106
+ type RunnableWithShutdownChan func (atShutdown <- chan struct {}, atTerminate WithCallback )
103
107
104
108
// RunWithShutdownChan takes a function that has channel to watch for shutdown and atTerminate callbacks
105
109
// After the atShutdown channel is closed, the main function must return once shutdown is complete.
@@ -115,8 +119,8 @@ func (g *Manager) RunWithShutdownChan(run RunnableWithShutdownChan) {
115
119
g .doShutdown ()
116
120
}
117
121
}()
118
- run (g .IsShutdown (), func (ctx context. Context , atTerminate func ()) {
119
- g .RunAtTerminate (ctx , atTerminate )
122
+ run (g .IsShutdown (), func (atTerminate func ()) {
123
+ g .RunAtTerminate (atTerminate )
120
124
})
121
125
}
122
126
@@ -136,60 +140,65 @@ func (g *Manager) RunWithShutdownContext(run func(context.Context)) {
136
140
}
137
141
138
142
// RunAtTerminate adds to the terminate wait group and creates a go-routine to run the provided function at termination
139
- func (g * Manager ) RunAtTerminate (ctx context. Context , terminate func ()) {
143
+ func (g * Manager ) RunAtTerminate (terminate func ()) {
140
144
g .terminateWaitGroup .Add (1 )
141
- go func () {
142
- defer g .terminateWaitGroup .Done ()
143
- defer func () {
144
- if err := recover (); err != nil {
145
- log .Critical ("PANIC during RunAtTerminate: %v\n Stacktrace: %s" , err , log .Stack (2 ))
146
- }
147
- }()
148
- select {
149
- case <- g .IsTerminate ():
145
+ g .lock .Lock ()
146
+ defer g .lock .Unlock ()
147
+ g .toRunAtTerminate = append (g .toRunAtTerminate ,
148
+ func () {
149
+ defer g .terminateWaitGroup .Done ()
150
+ defer func () {
151
+ if err := recover (); err != nil {
152
+ log .Critical ("PANIC during RunAtTerminate: %v\n Stacktrace: %s" , err , log .Stack (2 ))
153
+ }
154
+ }()
150
155
terminate ()
151
- case <- ctx .Done ():
152
- }
153
- }()
156
+ })
154
157
}
155
158
156
159
// RunAtShutdown creates a go-routine to run the provided function at shutdown
157
160
func (g * Manager ) RunAtShutdown (ctx context.Context , shutdown func ()) {
158
- go func () {
159
- defer func () {
160
- if err := recover (); err != nil {
161
- log .Critical ("PANIC during RunAtShutdown: %v\n Stacktrace: %s" , err , log .Stack (2 ))
161
+ g .lock .Lock ()
162
+ defer g .lock .Unlock ()
163
+ g .toRunAtShutdown = append (g .toRunAtShutdown ,
164
+ func () {
165
+ defer func () {
166
+ if err := recover (); err != nil {
167
+ log .Critical ("PANIC during RunAtShutdown: %v\n Stacktrace: %s" , err , log .Stack (2 ))
168
+ }
169
+ }()
170
+ select {
171
+ case <- ctx .Done ():
172
+ return
173
+ default :
174
+ shutdown ()
162
175
}
163
- }()
164
- select {
165
- case <- g .IsShutdown ():
166
- shutdown ()
167
- case <- ctx .Done ():
168
- }
169
- }()
176
+ })
170
177
}
171
178
172
179
// RunAtHammer creates a go-routine to run the provided function at shutdown
173
- func (g * Manager ) RunAtHammer (ctx context.Context , hammer func ()) {
174
- go func () {
175
- defer func () {
176
- if err := recover (); err != nil {
177
- log .Critical ("PANIC during RunAtHammer: %v\n Stacktrace: %s" , err , log .Stack (2 ))
178
- }
179
- }()
180
- select {
181
- case <- g .IsHammer ():
180
+ func (g * Manager ) RunAtHammer (hammer func ()) {
181
+ g .lock .Lock ()
182
+ defer g .lock .Unlock ()
183
+ g .toRunAtHammer = append (g .toRunAtHammer ,
184
+ func () {
185
+ defer func () {
186
+ if err := recover (); err != nil {
187
+ log .Critical ("PANIC during RunAtHammer: %v\n Stacktrace: %s" , err , log .Stack (2 ))
188
+ }
189
+ }()
182
190
hammer ()
183
- case <- ctx .Done ():
184
- }
185
- }()
191
+ })
186
192
}
187
193
func (g * Manager ) doShutdown () {
188
194
if ! g .setStateTransition (stateRunning , stateShuttingDown ) {
189
195
return
190
196
}
191
197
g .lock .Lock ()
192
- close (g .shutdown )
198
+ g .shutdownCtxCancel ()
199
+ for _ , fn := range g .toRunAtShutdown {
200
+ go fn ()
201
+ }
193
202
g .lock .Unlock ()
194
203
195
204
if setting .GracefulHammerTime >= 0 {
@@ -203,7 +212,7 @@ func (g *Manager) doShutdown() {
203
212
g .doTerminate ()
204
213
g .WaitForTerminate ()
205
214
g .lock .Lock ()
206
- close ( g . done )
215
+ g . doneCtxCancel ( )
207
216
g .lock .Unlock ()
208
217
}()
209
218
}
@@ -212,10 +221,13 @@ func (g *Manager) doHammerTime(d time.Duration) {
212
221
time .Sleep (d )
213
222
g .lock .Lock ()
214
223
select {
215
- case <- g .hammer :
224
+ case <- g .hammerCtx . Done () :
216
225
default :
217
226
log .Warn ("Setting Hammer condition" )
218
- close (g .hammer )
227
+ g .hammerCtxCancel ()
228
+ for _ , fn := range g .toRunAtHammer {
229
+ go fn ()
230
+ }
219
231
}
220
232
g .lock .Unlock ()
221
233
}
@@ -226,10 +238,13 @@ func (g *Manager) doTerminate() {
226
238
}
227
239
g .lock .Lock ()
228
240
select {
229
- case <- g .terminate :
241
+ case <- g .terminateCtx . Done () :
230
242
default :
231
243
log .Warn ("Terminating" )
232
- close (g .terminate )
244
+ g .terminateCtxCancel ()
245
+ for _ , fn := range g .toRunAtTerminate {
246
+ go fn ()
247
+ }
233
248
}
234
249
g .lock .Unlock ()
235
250
}
@@ -242,22 +257,22 @@ func (g *Manager) IsChild() bool {
242
257
// IsShutdown returns a channel which will be closed at shutdown.
243
258
// The order of closure is IsShutdown, IsHammer (potentially), IsTerminate
244
259
func (g * Manager ) IsShutdown () <- chan struct {} {
245
- return g .shutdown
260
+ return g .shutdownCtx . Done ()
246
261
}
247
262
248
263
// IsHammer returns a channel which will be closed at hammer
249
264
// The order of closure is IsShutdown, IsHammer (potentially), IsTerminate
250
265
// Servers running within the running server wait group should respond to IsHammer
251
266
// if not shutdown already
252
267
func (g * Manager ) IsHammer () <- chan struct {} {
253
- return g .hammer
268
+ return g .hammerCtx . Done ()
254
269
}
255
270
256
271
// IsTerminate returns a channel which will be closed at terminate
257
272
// The order of closure is IsShutdown, IsHammer (potentially), IsTerminate
258
273
// IsTerminate will only close once all running servers have stopped
259
274
func (g * Manager ) IsTerminate () <- chan struct {} {
260
- return g .terminate
275
+ return g .terminateCtx . Done ()
261
276
}
262
277
263
278
// ServerDone declares a running server done and subtracts one from the
@@ -314,25 +329,20 @@ func (g *Manager) InformCleanup() {
314
329
315
330
// Done allows the manager to be viewed as a context.Context, it returns a channel that is closed when the server is finished terminating
316
331
func (g * Manager ) Done () <- chan struct {} {
317
- return g .done
332
+ return g .doneCtx . Done ()
318
333
}
319
334
320
- // Err allows the manager to be viewed as a context.Context done at Terminate, it returns ErrTerminate
335
+ // Err allows the manager to be viewed as a context.Context done at Terminate
321
336
func (g * Manager ) Err () error {
322
- select {
323
- case <- g .Done ():
324
- return ErrTerminate
325
- default :
326
- return nil
327
- }
337
+ return g .doneCtx .Err ()
328
338
}
329
339
330
- // Value allows the manager to be viewed as a context.Context done at Terminate, it has no values
340
+ // Value allows the manager to be viewed as a context.Context done at Terminate
331
341
func (g * Manager ) Value (key interface {}) interface {} {
332
- return nil
342
+ return g . doneCtx . Value ( key )
333
343
}
334
344
335
345
// Deadline returns nil as there is no fixed Deadline for the manager, it allows the manager to be viewed as a context.Context
336
346
func (g * Manager ) Deadline () (deadline time.Time , ok bool ) {
337
- return
347
+ return g . doneCtx . Deadline ()
338
348
}
0 commit comments