app/common/termination.go

package common

import (
	"context"
	"errors"
	"fmt"
	"log/slog"
	"sync"
	"time"

	health "github.com/AppsFlyer/go-sundheit"
	"github.com/oklog/run"

	"github.com/openmeterio/openmeter/app/config"
)

type TerminationState = string

const (
	TerminationStateRunning     TerminationState = "running"
	TerminationStateTerminating TerminationState = "terminating"
)

type status struct {
	State     TerminationState `json:"state"`
	Reason    error            `json:"reason,omitempty"`
	Timestamp time.Time        `json:"timestamp,omitempty"`
}

var _ health.Check = (*TerminationChecker)(nil)

type TerminationChecker struct {
	status

	mu sync.RWMutex

	propagationTimeout time.Duration
}

func (s *TerminationChecker) setDetails(state TerminationState, reason error) {
	s.mu.Lock()
	defer s.mu.Unlock()

	s.status.State = state
	s.status.Reason = reason
	s.status.Timestamp = time.Now()
}

func (s *TerminationChecker) Terminate(reason error) {
	s.setDetails(TerminationStateTerminating, reason)
}

func (s *TerminationChecker) IsRunning() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.status.State == TerminationStateRunning
}

func (s *TerminationChecker) IsTerminating() bool {
	s.mu.RLock()
	defer s.mu.RUnlock()

	return s.status.State == TerminationStateTerminating
}

func (s *TerminationChecker) Name() string {
	return "termination.check"
}

func (s *TerminationChecker) Execute(_ context.Context) (interface{}, error) {
	details := struct {
		State     TerminationState `json:"state"`
		Reason    error            `json:"reason,omitempty"`
		Timestamp string           `json:"timestamp"`
	}{
		State:     s.status.State,
		Reason:    s.status.Reason,
		Timestamp: s.status.Timestamp.UTC().Format(time.RFC3339),
	}

	if s.IsTerminating() {
		return details, fmt.Errorf("termination was initiated [state=%s time=%s]: %w", details.State, details.Timestamp, details.Reason)
	}

	return details, nil
}

// WaitForPropagation blocks for
func (s *TerminationChecker) WaitForPropagation(ctx context.Context) error {
	if s.IsRunning() {
		return nil
	}

	propagationDeadline := s.status.Timestamp.Add(s.propagationTimeout)
	if time.Now().After(propagationDeadline) {
		return nil
	}

	ticker := time.NewTicker(time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case now := <-ticker.C:
			if now.After(propagationDeadline) {
				return nil
			}
		}
	}
}

// NewTerminationChecker returns a new TerminationChecker and registers the TerminationChecker to
// the provided health checker.
func NewTerminationChecker(conf config.TerminationConfig, healthChecker health.Health) (*TerminationChecker, error) {
	rc := &TerminationChecker{
		status: status{
			State:     TerminationStateRunning,
			Timestamp: time.Now(),
		},
		propagationTimeout: conf.PropagationTimeout,
	}

	err := healthChecker.RegisterCheck(
		rc,
		health.ExecutionPeriod(conf.CheckInterval),
	)
	if err != nil {
		return nil, err
	}

	return rc, nil
}

// NewTerminationCheckerActor returns actor functions (execute, interrupt) which can be passed to run.Add() to register it.
// Its state is set to TerminationStateStopping as soon as interrupt function is called.
func NewTerminationCheckerActor(r *TerminationChecker, logger *slog.Logger) (execute func() error, interrupt func(error), err error) {
	if r == nil {
		return nil, nil, errors.New("the TerminationChecker must not be nil")
	}

	if logger == nil {
		logger = slog.Default()
	}

	c := make(chan struct{})

	return func() error {
			// Wait until interrupt is invoked and return after
			<-c

			return nil
		}, func(err error) {
			logger.Debug("shutting down termination checker", "cause", err)

			r.Terminate(err)

			if errors.Is(err, run.ErrSignal) {
				logger.Debug("waiting for propagation of termination...", "until", r.Timestamp.Add(r.propagationTimeout))

				// Wait for propagation timeout
				if err = r.WaitForPropagation(context.Background()); err != nil {
					logger.Error("failed to wait for propagation", "error", err)
				}

				logger.Debug("waiting for propagation of termination is done", "until", r.Timestamp.Add(r.propagationTimeout))
			}

			close(c)
		}, nil
}