bbgo_origin/pkg/bbgo/graceful_shutdown.go

package bbgo

import (
	"context"
	"sync"
	"time"

	"github.com/sirupsen/logrus"
)

var graceful = &Graceful{}

//go:generate callbackgen -type Graceful
type Graceful struct {
	shutdownCallbacks []func(ctx context.Context, wg *sync.WaitGroup)
}

// Shutdown is a blocking call to emit all shutdown callbacks at the same time.
func (g *Graceful) 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(f func(ctx context.Context, wg *sync.WaitGroup)) {
	graceful.OnShutdown(f)
}

func Shutdown() {
	logrus.Infof("shutting down...")

	ctx, cancel := context.WithTimeout(context.TODO(), 30*time.Second)
	graceful.Shutdown(ctx)
	cancel()
}