bbgo_origin/pkg/bbgo/graceful_shutdown.go

package bbgo

import (
	"context"
	"sync"
	"time"

	"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.
func (g *GracefulShutdown) Shutdown(ctx context.Context) {
	var wg sync.WaitGroup
	wg.Add(len(g.shutdownCallbacks))

	// for each shutdown callback, we give them 10 second
	shtCtx, cancel := context.WithTimeout(ctx, 10*time.Second)

	go g.EmitShutdown(shtCtx, &wg)

	wg.Wait()
	cancel()
}

func OnShutdown(ctx context.Context, f ShutdownHandler) {
	isolatedContext := GetIsolationFromContext(ctx)
	isolatedContext.gracefulShutdown.OnShutdown(f)
}

func Shutdown(ctx context.Context) {
	logrus.Infof("shutting down...")

	isolatedContext := GetIsolationFromContext(ctx)
	todo := context.WithValue(context.TODO(), IsolationContextKey, isolatedContext)

	timeoutCtx, cancel := context.WithTimeout(todo, 30*time.Second)
	defaultIsolation.gracefulShutdown.Shutdown(timeoutCtx)
	cancel()
}