diff --git a/protofsm/daemon_events.go b/protofsm/daemon_events.go index bca7283d3..3c75ed9e6 100644 --- a/protofsm/daemon_events.go +++ b/protofsm/daemon_events.go @@ -72,6 +72,10 @@ func (b *BroadcastTxn) daemonSealed() {} // custom state machine event. type SpendMapper[Event any] func(*chainntnfs.SpendDetail) Event +// ConfMapper is a function that's used to map a confirmation notification to a +// custom state machine event. +type ConfMapper[Event any] func(*chainntnfs.TxConfirmation) Event + // RegisterSpend is used to request that a certain event is sent into the state // machine once the specified outpoint has been spent. type RegisterSpend[Event any] struct { @@ -112,10 +116,9 @@ type RegisterConf[Event any] struct { // transaction needs to dispatch an event. NumConfs fn.Option[uint32] - // PostConfEvent is an event that's sent back to the requester once the - // transaction specified above has confirmed in the chain with - // sufficient depth. - PostConfEvent fn.Option[Event] + // PostConfMapper is a special conf mapper, that if present, will be + // used to map the protofsm confirmation event to a custom event. + PostConfMapper fn.Option[ConfMapper[Event]] } // daemonSealed indicates that this struct is a DaemonEvent instance. diff --git a/protofsm/state_machine.go b/protofsm/state_machine.go index 5081480de..cd0e69385 100644 --- a/protofsm/state_machine.go +++ b/protofsm/state_machine.go @@ -522,16 +522,19 @@ func (s *StateMachine[Event, Env]) executeDaemonEvent(ctx context.Context, launched := s.gm.Go(ctx, func(ctx context.Context) { for { select { - case <-confEvent.Confirmed: - // If there's a post-conf event, then + //nolint:ll + case conf, ok := <-confEvent.Confirmed: + if !ok { + return + } + + // If there's a post-conf mapper, then // we'll send that into the current // state now. - // - // TODO(roasbeef): refactor to - // dispatchAfterRecv w/ above - postConf := daemonEvent.PostConfEvent - postConf.WhenSome(func(e Event) { - s.SendEvent(ctx, e) + postConfMapper := daemonEvent.PostConfMapper + postConfMapper.WhenSome(func(f ConfMapper[Event]) { + customEvent := f(conf) + s.SendEvent(ctx, customEvent) }) return diff --git a/protofsm/state_machine_test.go b/protofsm/state_machine_test.go index dc3562365..52b27b985 100644 --- a/protofsm/state_machine_test.go +++ b/protofsm/state_machine_test.go @@ -40,6 +40,20 @@ type daemonEvents struct { func (s *daemonEvents) dummy() { } +type confDetailsEvent struct { + blockHash chainhash.Hash + blockHeight uint32 +} + +func (c *confDetailsEvent) dummy() { +} + +type registerConf struct { +} + +func (r *registerConf) dummy() { +} + type dummyEnv struct { mock.Mock } @@ -127,6 +141,55 @@ func (d *dummyStateStart) ProcessEvent(event dummyEvents, env *dummyEnv, }, }), }, nil + + // This state will emit a RegisterConf event which uses a mapper to + // transition to the final state upon confirmation. + case *registerConf: + confMapper := func( + conf *chainntnfs.TxConfirmation) dummyEvents { + + // Map the conf details into our custom event. + return &confDetailsEvent{ + blockHash: *conf.BlockHash, + blockHeight: conf.BlockHeight, + } + } + + regConfEvent := &RegisterConf[dummyEvents]{ + Txid: chainhash.Hash{1}, + PkScript: []byte{0x01}, + HeightHint: 100, + PostConfMapper: fn.Some[ConfMapper[dummyEvents]]( + confMapper, + ), + } + + return &StateTransition[dummyEvents, *dummyEnv]{ + // Stay in the start state until the conf event is + // received and mapped. + NextState: &dummyStateStart{ + canSend: d.canSend, + }, + NewEvents: fn.Some(EmittedEvent[dummyEvents]{ + ExternalEvents: DaemonEventSet{ + regConfEvent, + }, + }), + }, nil + + // This event contains details from the confirmation and signals us to + // transition to the final state. + case *confDetailsEvent: + eventDetails := event.(*confDetailsEvent) + + // We received the mapped confirmation details, transition to the + // confirmed state. + return &StateTransition[dummyEvents, *dummyEnv]{ + NextState: &dummyStateConfirmed{ + blockHash: eventDetails.blockHash, + blockHeight: eventDetails.blockHeight, + }, + }, nil } return nil, fmt.Errorf("unknown event: %T", event) @@ -155,6 +218,28 @@ func (d *dummyStateFin) IsTerminal() bool { return true } +type dummyStateConfirmed struct { + blockHash chainhash.Hash + blockHeight uint32 +} + +func (d *dummyStateConfirmed) String() string { + return "dummyStateConfirmed" +} + +func (d *dummyStateConfirmed) ProcessEvent(event dummyEvents, env *dummyEnv, +) (*StateTransition[dummyEvents, *dummyEnv], error) { + + // This is a terminal state, no further transitions. + return &StateTransition[dummyEvents, *dummyEnv]{ + NextState: d, + }, nil +} + +func (d *dummyStateConfirmed) IsTerminal() bool { + return true +} + func assertState[Event any, Env Environment](t *testing.T, m *StateMachine[Event, Env], expectedState State[Event, Env]) { @@ -415,6 +500,78 @@ func TestStateMachineDaemonEvents(t *testing.T) { env.AssertExpectations(t) } +// TestStateMachineConfMapper tests that the state machine is able to properly +// map the confirmation event into a custom event that can be used to trigger a +// state transition. +func TestStateMachineConfMapper(t *testing.T) { + t.Parallel() + ctx := context.Background() + + // Create the state machine. + env := &dummyEnv{} + startingState := &dummyStateStart{} + adapters := newDaemonAdapters() + + cfg := StateMachineCfg[dummyEvents, *dummyEnv]{ + Daemon: adapters, + InitialState: startingState, + Env: env, + } + stateMachine := NewStateMachine(cfg) + + stateSub := stateMachine.RegisterStateEvents() + defer stateMachine.RemoveStateSub(stateSub) + + stateMachine.Start(ctx) + defer stateMachine.Stop() + + // Expect the RegisterConfirmationsNtfn call when we send the event. + // We use NumConfs=1 as the default. + adapters.On( + "RegisterConfirmationsNtfn", &chainhash.Hash{1}, []byte{0x01}, + uint32(1), + ).Return(nil) + + // Send the event that triggers RegisterConf emission. + stateMachine.SendEvent(ctx, ®isterConf{}) + + // We should transition back to the starting state initially. + expectedStates := []State[dummyEvents, *dummyEnv]{ + &dummyStateStart{}, &dummyStateStart{}, + } + assertStateTransitions(t, stateSub, expectedStates) + + // Assert the registration call was made. + adapters.AssertExpectations(t) + + // Now, simulate the confirmation event coming back from the notifier. + // Populate it with some data to be mapped. + simulatedConf := &chainntnfs.TxConfirmation{ + BlockHash: &chainhash.Hash{2}, + BlockHeight: 123, + } + adapters.confChan <- simulatedConf + + // This should trigger the mapper and send the confDetailsEvent, + // transitioning us to the final state. + expectedStates = []State[dummyEvents, *dummyEnv]{&dummyStateConfirmed{}} + assertStateTransitions(t, stateSub, expectedStates) + + // Final state assertion. + finalState, err := stateMachine.CurrentState() + require.NoError(t, err) + require.IsType(t, &dummyStateConfirmed{}, finalState) + + // Assert that the details from the confirmation event were correctly + // propagated to the final state. + finalStateDetails := finalState.(*dummyStateConfirmed) + require.Equal(t, simulatedConf.BlockHash, &finalStateDetails.blockHash) + require.Equal(t, simulatedConf.BlockHeight, finalStateDetails.blockHeight) + + adapters.AssertExpectations(t) + env.AssertExpectations(t) +} + type dummyMsgMapper struct { mock.Mock }