forked from c9s/bbgo
-
Notifications
You must be signed in to change notification settings - Fork 0
/
graceful_shutdown.go
44 lines (35 loc) · 1.31 KB
/
graceful_shutdown.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
package bbgo
import (
"context"
"sync"
"github.com/sirupsen/logrus"
)
type ShutdownHandler func(ctx context.Context, wg *sync.WaitGroup)
//go:generate callbackgen -type GracefulShutdown
type GracefulShutdown struct {
shutdownCallbacks []ShutdownHandler
}
// Shutdown is a blocking call to emit all shutdown callbacks at the same time.
// The context object here should not be canceled context, you need to create a todo context.
func (g *GracefulShutdown) Shutdown(shutdownCtx context.Context) {
var wg sync.WaitGroup
wg.Add(len(g.shutdownCallbacks))
go g.EmitShutdown(shutdownCtx, &wg)
wg.Wait()
}
// OnShutdown helps you register your shutdown handler
// the first context object is where you want to register your shutdown handler, where the context has the isolated storage.
// in your handler, you will get another context for the timeout context.
func OnShutdown(ctx context.Context, f ShutdownHandler) {
isolatedContext := GetIsolationFromContext(ctx)
isolatedContext.gracefulShutdown.OnShutdown(f)
}
func Shutdown(shutdownCtx context.Context) {
isolatedContext := GetIsolationFromContext(shutdownCtx)
if isolatedContext == defaultIsolation {
logrus.Infof("bbgo shutting down...")
} else {
logrus.Infof("bbgo shutting down (custom isolation)...")
}
isolatedContext.gracefulShutdown.Shutdown(shutdownCtx)
}