chore: add statemachine event handlers (#2363)

* chore: try out state machine event handlers

chore: tests

chore: keep the PR focused

* chore: address PR feedback

* chore: address feedback

* chore: address lock holding concern

Author: JGAntunes

api/internal/statemachine/event_handler.go

@@ -0,0 +1,78 @@
+package statemachine
+
+import (
+	"context"
+	"fmt"
+	"runtime/debug"
+	"time"
+)
+
+var (
+	_ EventHandler = &eventHandler{}
+)
+
+// EventHandler is an interface for handling state transition events in the state machine.
+type EventHandler interface {
+	// TriggerHandler triggers the event handler for a state transition.
+	TriggerHandler(ctx context.Context, fromState, toState State) error
+}
+
+// EventHandlerFunc is a function that handles state transition events. Used to report state changes.
+type EventHandlerFunc func(ctx context.Context, fromState, toState State)
+
+// EventHandlerOption is a configurable state machine option.
+type EventHandlerOption func(*eventHandler)
+
+// WithHandlerTimeout sets the timeout for the event handler to complete.
+func WithHandlerTimeout(timeout time.Duration) EventHandlerOption {
+	return func(eh *eventHandler) {
+		eh.timeout = timeout
+	}
+}
+
+// NewEventHandler creates a new event handler with the provided function and options.
+func NewEventHandler(handler EventHandlerFunc, options ...EventHandlerOption) EventHandler {
+	eh := &eventHandler{
+		handler: handler,
+		timeout: 5 * time.Second, // Default timeout
+	}
+
+	for _, option := range options {
+		option(eh)
+	}
+
+	return eh
+}
+
+// eventHandler is a struct that implements the EventHandler interface. It contains a handler function that is called when a state transition occurs, and it supports a timeout for the handler to complete.
+type eventHandler struct {
+	handler EventHandlerFunc
+	timeout time.Duration // Timeout for the handler to complete, default is 5 seconds
+}
+
+// TriggerHandler triggers the event handler for a state transition. The trigger is blocking and will wait for the handler to complete or timeout.
+func (eh *eventHandler) TriggerHandler(ctx context.Context, fromState, toState State) error {
+	ctx, cancel := context.WithTimeout(ctx, eh.timeout)
+	defer cancel()
+	done := make(chan error, 1)
+
+	go func() {
+		defer func() {
+			if r := recover(); r != nil {
+				// Capture panic but don't affect the transition
+				err := fmt.Errorf("event handler panic from %s to %s: %v: %s\n", fromState, toState, r, debug.Stack())
+				done <- err
+			}
+			close(done)
+		}()
+		eh.handler(ctx, fromState, toState)
+	}()
+
+	select {
+	case err := <-done:
+		return err
+	case <-ctx.Done():
+		err := fmt.Errorf("event handler for transition from %s to %s timed out after %s", fromState, toState, eh.timeout)
+		return err
+	}
+}

api/internal/statemachine/interface.go

@@ -0,0 +1,34 @@
+package statemachine
+
+// State represents the possible states of the install process
+type State string
+
+var (
+	_ Interface = &stateMachine{}
+)
+
+// Interface is the interface for the state machine
+type Interface interface {
+	// CurrentState returns the current state
+	CurrentState() State
+	// IsFinalState checks if the current state is a final state
+	IsFinalState() bool
+	// ValidateTransition checks if a transition from the current state to a new state is valid
+	ValidateTransition(lock Lock, newState State) error
+	// Transition attempts to transition to a new state and returns an error if the transition is
+	// invalid.
+	Transition(lock Lock, nextState State) error
+	// AcquireLock acquires a lock on the state machine.
+	AcquireLock() (Lock, error)
+	// IsLockAcquired checks if a lock already exists on the state machine.
+	IsLockAcquired() bool
+	// RegisterEventHandler registers a blocking event handler for reporting events in the state machine.
+	RegisterEventHandler(targetState State, handler EventHandlerFunc, options ...EventHandlerOption)
+	// UnregisterEventHandler unregisters a blocking event handler for reporting events in the state machine.
+	UnregisterEventHandler(targetState State)
+}
+
+type Lock interface {
+	// Release releases the lock.
+	Release()
+}

api/internal/statemachine/statemachine.go

@@ -1,54 +1,48 @@
 package statemachine
 
 import (
+	"context"
 	"fmt"
 	"slices"
 	"sync"
-)
-
-// State represents the possible states of the install process
-type State string
 
-var (
-	_ Interface = &stateMachine{}
+	"github.com/replicatedhq/embedded-cluster/api/pkg/logger"
+	"github.com/sirupsen/logrus"
 )
 
-// Interface is the interface for the state machine
-type Interface interface {
-	// CurrentState returns the current state
-	CurrentState() State
-	// IsFinalState checks if the current state is a final state
-	IsFinalState() bool
-	// ValidateTransition checks if a transition from the current state to a new state is valid
-	ValidateTransition(lock Lock, newState State) error
-	// Transition attempts to transition to a new state and returns an error if the transition is
-	// invalid.
-	Transition(lock Lock, nextState State) error
-	// AcquireLock acquires a lock on the state machine.
-	AcquireLock() (Lock, error)
-	// IsLockAcquired checks if a lock already exists on the state machine.
-	IsLockAcquired() bool
-}
-
-type Lock interface {
-	// Release releases the lock.
-	Release()
-}
-
 // stateMachine manages the state transitions for the install process
 type stateMachine struct {
 	currentState          State
 	validStateTransitions map[State][]State
 	lock                  *lock
 	mu                    sync.RWMutex
+	eventHandlers         map[State][]EventHandler
+	logger                logrus.FieldLogger
 }
 
+// StateMachineOption is a configurable state machine option.
+type StateMachineOption func(*stateMachine)
+
 // New creates a new state machine starting in the given state with the given valid state
-// transitions.
-func New(currentState State, validStateTransitions map[State][]State) *stateMachine {
-	return &stateMachine{
+// transitions and options.
+func New(currentState State, validStateTransitions map[State][]State, opts ...StateMachineOption) *stateMachine {
+	sm := &stateMachine{
 		currentState:          currentState,
 		validStateTransitions: validStateTransitions,
+		logger:                logger.NewDiscardLogger(),
+		eventHandlers:         make(map[State][]EventHandler),
+	}
+
+	for _, opt := range opts {
+		opt(sm)
+	}
+
+	return sm
+}
+
+func WithLogger(logger logrus.FieldLogger) StateMachineOption {
+	return func(sm *stateMachine) {
+		sm.logger = logger
 	}
 }
 
@@ -109,9 +103,13 @@ func (sm *stateMachine) ValidateTransition(lock Lock, nextState State) error {
 	return nil
 }
 
-func (sm *stateMachine) Transition(lock Lock, nextState State) error {
+func (sm *stateMachine) Transition(lock Lock, nextState State) (finalError error) {
 	sm.mu.Lock()
-	defer sm.mu.Unlock()
+	defer func() {
+		if finalError != nil {
+			sm.mu.Unlock()
+		}
+	}()
 
 	if sm.lock == nil {
 		return fmt.Errorf("lock not acquired")
@@ -123,11 +121,43 @@ func (sm *stateMachine) Transition(lock Lock, nextState State) error {
 		return fmt.Errorf("invalid transition from %s to %s", sm.currentState, nextState)
 	}
 
+	fromState := sm.currentState
 	sm.currentState = nextState
 
+	// Trigger event handlers after successful transition
+	handlers, exists := sm.eventHandlers[nextState]
+	safeHandlers := make([]EventHandler, len(handlers))
+	copy(safeHandlers, handlers) // Copy to avoid holding the lock while calling handlers
+
+	// We can release the lock here since the transition is successful and there will be no further operations to the state machine internal state
+	sm.mu.Unlock()
+
+	if !exists || len(safeHandlers) == 0 {
+		return nil
+	}
+
+	for _, handler := range safeHandlers {
+		err := handler.TriggerHandler(context.Background(), fromState, nextState)
+		if err != nil {
+			sm.logger.WithFields(logrus.Fields{"fromState": fromState, "toState": nextState}).Errorf("event handler error: %v", err)
+		}
+	}
+
 	return nil
 }
 
+func (sm *stateMachine) RegisterEventHandler(targetState State, handler EventHandlerFunc, options ...EventHandlerOption) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+	sm.eventHandlers[targetState] = append(sm.eventHandlers[targetState], NewEventHandler(handler, options...))
+}
+
+func (sm *stateMachine) UnregisterEventHandler(targetState State) {
+	sm.mu.Lock()
+	defer sm.mu.Unlock()
+	delete(sm.eventHandlers, targetState)
+}
+
 func (sm *stateMachine) isValidTransition(currentState State, newState State) bool {
 	validTransitions, ok := sm.validStateTransitions[currentState]
 	if !ok {