From d17e7375582a3c85f771ac1b1821e0bd29d72a7d Mon Sep 17 00:00:00 2001 From: Olaoluwa Osuntokun Date: Tue, 23 Jan 2024 18:57:00 -0800 Subject: [PATCH] protofsm: add optional daemon event on init In this commit, we add an optional daemon event that can be specified to dispatch during init. This is useful for instances where before we start, we want to make sure we have a registered spend/conf notification before normal operation starts. We also add new unit tests to cover this, and the prior spend/conf event additions. --- protofsm/state_machine.go | 24 +++++++++++++-- protofsm/state_machine_test.go | 54 +++++++++++++++++++++++++++++++--- 2 files changed, 71 insertions(+), 7 deletions(-) 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()