diff --git a/protofsm/state_machine.go b/protofsm/state_machine.go index e25d82515..9c4db8dc1 100644 --- a/protofsm/state_machine.go +++ b/protofsm/state_machine.go @@ -135,6 +135,8 @@ type StateMachine[Event any, Env Environment] struct { stateQuery chan stateQuery[Event, Env] + initEvent fn.Option[DaemonEvent] + startOnce sync.Once stopOnce sync.Once @@ -142,10 +144,12 @@ type StateMachine[Event any, Env Environment] struct { } // NewStateMachine creates a new state machine given a set of daemon adapters, -// an initial state, and an environment. +// an initial state, an environment, and an event to process as if emitted at +// the onset of the state machine. Such an event can be used to set up tracking +// state such as a txid confirmation event. func NewStateMachine[Event any, Env Environment](adapters DaemonAdapters, - initialState State[Event, Env], - env Env) StateMachine[Event, Env] { + initialState State[Event, Env], env Env, + initEvent fn.Option[DaemonEvent]) StateMachine[Event, Env] { return StateMachine[Event, Env]{ daemon: adapters, @@ -154,6 +158,7 @@ func NewStateMachine[Event any, Env Environment](adapters DaemonAdapters, stateQuery: make(chan stateQuery[Event, Env]), quit: make(chan struct{}), env: env, + initEvent: initEvent, newStateEvents: fn.NewEventDistributor[State[Event, Env]](), } } @@ -443,6 +448,9 @@ func (s *StateMachine[Event, Env]) applyEvents(newEvent Event) (State[Event, Env currentState = transition.NextState // Notify our subscribers of the new state transition. + // + // TODO(roasbeef): will only give us the outer state? + // * let FSMs choose which state to emit? s.newStateEvents.NotifySubscribers(currentState) return nil @@ -464,6 +472,16 @@ func (s *StateMachine[Event, Env]) driveMachine() { // TODO(roasbeef): move into env? read only to start with currentState := s.currentState + // Before we start, if we have an init daemon event specified, then + // we'll handle that now. + err := fn.MapOptionZ(s.initEvent, func(event DaemonEvent) error { + return s.executeDaemonEvent(event) + }) + if err != nil { + log.Errorf("unable to execute init event: %w", err) + return + } + // We just started driving the state machine, so we'll notify our // subscribers of this starting state. s.newStateEvents.NotifySubscribers(currentState) diff --git a/protofsm/state_machine_test.go b/protofsm/state_machine_test.go index 17a5a32c4..da765574e 100644 --- a/protofsm/state_machine_test.go +++ b/protofsm/state_machine_test.go @@ -211,6 +211,52 @@ func (d *dummyAdapters) RegisterSpendNtfn(outpoint *wire.OutPoint, }, err } +// TestStateMachineOnInitDaemonEvent tests that the state machine will properly +// execute any init-level daemon events passed into it. +func TestStateMachineOnInitDaemonEvent(t *testing.T) { + // First, we'll create our state machine given the env, and our + // starting state. + env := &dummyEnv{} + startingState := &dummyStateStart{} + + adapters := newDaemonAdapters() + + // We'll make an init event that'll send to a peer, then transition us + // to our terminal state. + initEvent := &SendMsgEvent[dummyEvents]{ + TargetPeer: *pub1, + PostSendEvent: fn.Some(dummyEvents(&goToFin{})), + } + + stateMachine := NewStateMachine[dummyEvents, *dummyEnv]( + adapters, startingState, env, fn.Some[DaemonEvent](initEvent), + ) + + // Before we start up the state machine, we'll assert that the send + // message adapter is called on start up. + adapters.On("SendMessages", *pub1, mock.Anything).Return(nil) + + stateMachine.Start() + defer stateMachine.Stop() + + // As we're triggering internal events, we'll also subscribe to the set + // of new states so we can assert as we go. + stateSub := stateMachine.RegisterStateEvents() + defer stateMachine.RemoveStateSub(stateSub) + + // Assert that we go from the starting state to the final state. The + // state machine should now also be on the final terminal state. + expectedStates := []State[dummyEvents, *dummyEnv]{ + &dummyStateStart{}, &dummyStateFin{}, + } + assertStateTransitions(t, stateSub, expectedStates) + + // We'll now assert that after the daemon was started, the send message + // adapter was called above as specified in the init event. + adapters.AssertExpectations(t) + env.AssertExpectations(t) +} + // TestStateMachineInternalEvents tests that the state machine is able to add // new internal events to the event queue for further processing during a state // transition. @@ -222,10 +268,10 @@ func TestStateMachineInternalEvents(t *testing.T) { env := &dummyEnv{} startingState := &dummyStateStart{} - adapters := &dummyAdapters{} + adapters := newDaemonAdapters() stateMachine := NewStateMachine[dummyEvents, *dummyEnv]( - adapters, startingState, env, + adapters, startingState, env, fn.None[DaemonEvent](), ) stateMachine.Start() defer stateMachine.Stop() @@ -269,10 +315,10 @@ func TestStateMachineDaemonEvents(t *testing.T) { canSend: &boolTrigger, } - adapters := &dummyAdapters{} + adapters := newDaemonAdapters() stateMachine := NewStateMachine[dummyEvents, *dummyEnv]( - adapters, startingState, env, + adapters, startingState, env, fn.None[DaemonEvent](), ) stateMachine.Start() defer stateMachine.Stop()