bbgo_origin/pkg/bbgo/graceful_shutdown.go

45 lines
1.3 KiB
Go
Raw Permalink Normal View History

package bbgo
import (
"context"
"sync"
"github.com/sirupsen/logrus"
)
2022-07-19 09:13:35 +00:00
type ShutdownHandler func(ctx context.Context, wg *sync.WaitGroup)
2022-10-02 11:34:52 +00:00
//go:generate callbackgen -type GracefulShutdown
type GracefulShutdown struct {
2022-07-19 09:13:35 +00:00
shutdownCallbacks []ShutdownHandler
}
// Shutdown is a blocking call to emit all shutdown callbacks at the same time.
2022-10-11 06:22:46 +00:00
// 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))
2022-10-11 06:22:46 +00:00
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.
2022-10-03 08:01:08 +00:00
func OnShutdown(ctx context.Context, f ShutdownHandler) {
isolatedContext := GetIsolationFromContext(ctx)
2022-10-03 08:01:08 +00:00
isolatedContext.gracefulShutdown.OnShutdown(f)
}
2022-10-11 06:22:46 +00:00
func Shutdown(shutdownCtx context.Context) {
2023-02-08 09:30:33 +00:00
isolatedContext := GetIsolationFromContext(shutdownCtx)
if isolatedContext == defaultIsolation {
logrus.Infof("bbgo shutting down...")
} else {
logrus.Infof("bbgo shutting down (custom isolation)...")
}
2023-02-08 09:30:33 +00:00
isolatedContext.gracefulShutdown.Shutdown(shutdownCtx)
}