bbgo_origin/pkg/bbgo/graceful_shutdown.go

45 lines
1.3 KiB
Go

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)
}