application.go

package cli

import (
	"context"
	"fmt"
	"math"
	"time"

	"github.com/streamingfast/shutter"
	"go.uber.org/atomic"
	"go.uber.org/zap"
)

// Application is a simple object that can be used to manage the lifecycle of an application. You
// create the application with `NewApplication` and then you can use it to supervise and start
// child processes by using `SuperviseAndStart`.
//
// You then wait for the application to terminate by calling `WaitForTermination`. This call is blocking
// and register a signal handler to be notified of SIGINT and SIGTERM.
type Application struct {
	appCtx  context.Context
	shutter *shutter.Shutter

	isSignaled *atomic.Bool
}

func NewApplication(ctx context.Context) *Application {
	shutter := shutter.New()

	appCtx, cancelApp := context.WithCancel(ctx)
	shutter.OnTerminating(func(_ error) {
		cancelApp()
	})

	return &Application{
		appCtx:     appCtx,
		shutter:    shutter,
		isSignaled: atomic.NewBool(false),
	}
}

// IsReady returns true if the application is ready to be used. When the Ctrl-C signal is received,
// the app is immediately marked as not ready and the `WaitForTermination` call will wait for the
// unreadyPeriodDelay to expire before terminating the application.
func (a *Application) IsReady() bool {
	return !a.isSignaled.Load()
}

func (a *Application) Context() context.Context {
	return a.appCtx
}

// Shutter interface over the `*shutter.Shutter` struct.
type Shutter interface {
	OnTerminated(f func(error))
	OnTerminating(f func(error))
	Shutdown(error)
}

// Runnable contracts is to be blocking and to return only when the task is done.
type Runnable interface {
	Run()
}

// RunnableError contracts is to be blocking and to return only when the task is done. We assume
// the error happens while bootstrapping the task.
type RunnableError interface {
	Run() error
}

// RunnableContext contracts is to be blocking and to return only when the task is done. The context
// must be **used** only for the bootstrap period of task, long running task should be tied to a
// `*shutter.Shutter` instance.
type RunnableContext interface {
	Run(ctx context.Context)
}

// RunnableContext contracts is to be blocking and to return only when the task is done. The context
// must be **used** only for the bootstrap period of task, long running task should be tied to a
// `*shutter.Shutter` instance. We assume the error happens while bootstrapping the task.
type RunnableContextError interface {
	Run(ctx context.Context) error
}

// Supervise the received child shutter, mainly, this ensures that on child's termination,
// the application is also terminated with the error that caused the child to terminate.
//
// If the application shuts down before the child, the child is also terminated but
// gracefully (it does **not** receive the error that caused the application to terminate).
//
// The child termination is always performed before the application fully complete, unless
// the gracecul shutdown delay has expired.
func (a *Application) Supervise(child Shutter) {
	child.OnTerminated(a.shutter.Shutdown)
	a.shutter.OnTerminating(func(_ error) {
		child.Shutdown(nil)
	})
}

// SuperviseAndStart calls [Supervise] and then starts the child in a goroutine. The received
// child must implement one of [Runnable], [RunnableContext], [RunnableError] or [RunnableContextError] to
// be able to be started correctly.
//
// The child is started in a goroutine and tied to the application lifecycle because we also
// called [Supervise]. Later the call to `WaitForTermination` will wait for the application to
// terminate which will also terminates and wait for all child.
//
// People can use either [SuperviseAndStart] or [SuperviseAndStartUsing] and passing `child.Run`
// explicitly depending on their preference, see [SuperviseAndStartUsing] for details.
func (a *Application) SuperviseAndStart(child Shutter) {
	a.Supervise(child)

	switch v := child.(type) {
	case Runnable:
		a.SuperviseAndStartUsing(child, v.Run)
	case RunnableContext:
		a.SuperviseAndStartUsing(child, v.Run)
	case RunnableError:
		a.SuperviseAndStartUsing(child, v.Run)
	case RunnableContextError:
		a.SuperviseAndStartUsing(child, v.Run)

	default:
		panic(fmt.Errorf("unsupported child type %T, must implement one of cli.Runnable, cli.RunnableContext, cli.RunnableError or cli.RunnableContextError", child))
	}
}

// SuperviseAndStartUsing calls [Supervise] and then starts the child in a goroutine using the
// provided runner function. The runner function must be one of the following types:
//
//	func()
//	func(ctx context.Context)
//	func() error
//	func(ctx context.Context) error
//
// This can be used to make it more explicit what is the runner function that is started, making
// the `Run` method apparent at call site. Compare:
//
//	a.SuperviseAndStart(child)
//
// with
//
//	a.SuperviseAndStartUsing(child, child.Run)
//
// In the second case, you can "<key>-click" to go to the `Run` method of the child
// directly, while in the first case, you have to search, the linking is not as direct.
//
// People can use either [SuperviseAndStart] or [SuperviseAndStartUsing] depending on their preference.
func (a *Application) SuperviseAndStartUsing(child Shutter, runner any) {
	a.Supervise(child)

	switch v := runner.(type) {
	case func():
		go func() {
			v()
			child.Shutdown(nil)
		}()
	case func(ctx context.Context):
		go func() {
			v(a.appCtx)
			child.Shutdown(nil)
		}()
	case func() error:
		go func() {
			err := v()
			child.Shutdown(err)
		}()
	case func(ctx context.Context) error:
		go func() {
			err := v(a.appCtx)
			child.Shutdown(err)
		}()

	default:
		panic(fmt.Errorf("unsupported child type %T, must pass a function  cli.Runnable, cli.RunnableContext, cli.RunnableError or cli.RunnableContextError", child))
	}
}

// BlockUntilTerminated waits for the application to terminate. This is a blocking call that
// will wait for the application to be terminated. After receiving the initial terminating signal, if
// gracefulShutdownDelay is not 0, the terminated signal must happen within the gracefulShutdownDelay
// otherwise the call will unblock.
func (a *Application) BlockUntilTerminated(logger *zap.Logger, gracefulShutdownDelay time.Duration) error {
	<-a.shutter.Terminating()
	logger.Info("application terminating", zap.Bool("with_error", a.shutter.Err() != nil))

	maxWaitDelay := gracefulShutdownDelay
	if maxWaitDelay == 0 {
		maxWaitDelay = time.Duration(math.MaxInt64)
	}

	logger.Info("waiting for application termination")
	select {
	case <-a.shutter.Terminated():
	case <-time.After(maxWaitDelay):
		err := fmt.Errorf("application did not terminate within graceful period of %s", maxWaitDelay)

		if appErr := a.shutter.Err(); appErr != nil {
			return fmt.Errorf("%w: %w", err, appErr)
		}

		return err
	}

	if err := a.shutter.Err(); err != nil {
		return fmt.Errorf("application terminated with error: %w", err)
	}

	logger.Info("application terminated gracefully")
	return nil
}

// WaitForTermination waits for the application to terminate. This first setup the signal handler and
// then wait for either the signal handler to be notified or the application to be terminating.
//
// On application terminating, all child registered with [Supervise] are also terminated. We then wait for
// all child to gracefully terminate. If the graceful shutdown delay is reached, we force the termination
// of the application right now.
//
// Doing Ctrl-C 4 times or more will lead to a force quit of the whole process by calling `os.Exit(1)`, this
// is performed by the signal handler code and is does **not** respect the graceful shutdown delay in this case.
func (a *Application) WaitForTermination(logger *zap.Logger, unreadyPeriodDelay, gracefulShutdownDelay time.Duration) error {
	// On any exit path, we synchronize the logger one last time
	defer func() {
		logger.Sync()
	}()

	signalHandler, isSignaled, _ := SetupSignalHandler(unreadyPeriodDelay, logger)

	// Wire the signal handler to the application
	a.isSignaled = isSignaled

	select {
	case <-signalHandler:
		go a.shutter.Shutdown(nil)
		break
	case <-a.shutter.Terminating():
		logger.Info("run terminating", zap.Bool("from_signal", isSignaled.Load()), zap.Bool("with_error", a.shutter.Err() != nil))
		break
	}

	logger.Info("waiting for run termination")
	select {
	case <-a.shutter.Terminated():
	case <-time.After(gracefulShutdownDelay):
		if gracefulShutdownDelay > 0 {
			logger.Warn("application did not terminate within graceful period of " + gracefulShutdownDelay.String() + ", forcing termination")
		}
	}

	if err := a.shutter.Err(); err != nil {
		return err
	}

	logger.Info("run terminated gracefully")
	return nil
}