diff --git a/protofsm/daemon_events.go b/protofsm/daemon_events.go index bca7283d39..3c75ed9e69 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 3930fdff2b..5145327f7c 100644 --- a/protofsm/state_machine.go +++ b/protofsm/state_machine.go @@ -506,16 +506,18 @@ 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 + 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 //nolint:ll + postConfMapper.WhenSome(func(f ConfMapper[Event]) { //nolint:ll + customEvent := f(conf) + s.SendEvent(ctx, customEvent) }) return diff --git a/protofsm/state_machine_test.go b/protofsm/state_machine_test.go index ca05d336c0..3dcb35099e 100644 --- a/protofsm/state_machine_test.go +++ b/protofsm/state_machine_test.go @@ -40,6 +40,34 @@ 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 spendDetailsEvent struct { + spenderTxHash chainhash.Hash + spendingHeight int32 +} + +func (s *spendDetailsEvent) dummy() { +} + +type registerSpend struct { +} + +func (r *registerSpend) dummy() { +} + type dummyEnv struct { mock.Mock } @@ -74,7 +102,7 @@ var ( func (d *dummyStateStart) ProcessEvent(event dummyEvents, env *dummyEnv, ) (*StateTransition[dummyEvents, *dummyEnv], error) { - switch event.(type) { + switch newEvent := event.(type) { case *goToFin: return &StateTransition[dummyEvents, *dummyEnv]{ NextState: &dummyStateFin{}, @@ -127,6 +155,100 @@ 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: + // We received the mapped confirmation details, transition to + // the confirmed state. + return &StateTransition[dummyEvents, *dummyEnv]{ + NextState: &dummyStateConfirmed{ + blockHash: newEvent.blockHash, + blockHeight: newEvent.blockHeight, + }, + }, nil + + // This state will emit a RegisterSpend event which uses a mapper to + // transition to the spent state upon spend detection. + case *registerSpend: + spendMapper := func( + spend *chainntnfs.SpendDetail) dummyEvents { + + // Map the spend details into our custom event. + return &spendDetailsEvent{ + spenderTxHash: *spend.SpenderTxHash, + spendingHeight: spend.SpendingHeight, + } + } + + regSpendEvent := &RegisterSpend[dummyEvents]{ + OutPoint: wire.OutPoint{Hash: chainhash.Hash{3}}, + PkScript: []byte{0x03}, + HeightHint: 300, + PostSpendEvent: fn.Some[SpendMapper[dummyEvents]]( + spendMapper, + ), + } + + return &StateTransition[dummyEvents, *dummyEnv]{ + // Stay in the start state until the spend event is + // received and mapped. + NextState: &dummyStateStart{ + canSend: d.canSend, + }, + NewEvents: fn.Some(EmittedEvent[dummyEvents]{ + ExternalEvents: DaemonEventSet{ + regSpendEvent, + }, + }), + }, nil + + // This event contains details from the spend notification and signals + // us to transition to the spent state. + case *spendDetailsEvent: + // We received the mapped spend details, transition to the + // spent state. + return &StateTransition[dummyEvents, *dummyEnv]{ + NextState: &dummyStateSpent{ + spenderTxHash: newEvent.spenderTxHash, + spendingHeight: newEvent.spendingHeight, + }, + }, nil } return nil, fmt.Errorf("unknown event: %T", event) @@ -155,12 +277,64 @@ func (d *dummyStateFin) IsTerminal() bool { return true } -func assertState[Event any, Env Environment](t *testing.T, - m *StateMachine[Event, Env], expectedState State[Event, Env]) { +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 +} + +type dummyStateSpent struct { + spenderTxHash chainhash.Hash + spendingHeight int32 +} + +func (d *dummyStateSpent) String() string { + return "dummyStateSpent" +} + +func (d *dummyStateSpent) 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 *dummyStateSpent) IsTerminal() bool { + return true +} + +// assertState asserts that the state machine is currently in the expected +// state type and returns the state cast to that type. +func assertState[Event any, Env Environment, S State[Event, Env]](t *testing.T, + m *StateMachine[Event, Env], expectedState S) S { state, err := m.CurrentState() require.NoError(t, err) require.IsType(t, expectedState, state) + + // Perform the type assertion to return the concrete type. + concreteState, ok := state.(S) + require.True(t, ok, "state type assertion failed") + + return concreteState } func assertStateTransitions[Event any, Env Environment]( @@ -402,6 +576,156 @@ 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 := assertState(t, &stateMachine, &dummyStateConfirmed{}) + + // Assert that the details from the confirmation event were correctly + // propagated to the final state. + require.Equal(t, + *simulatedConf.BlockHash, finalState.blockHash, + ) + require.Equal(t, + simulatedConf.BlockHeight, finalState.blockHeight, + ) + + adapters.AssertExpectations(t) + env.AssertExpectations(t) +} + +// TestStateMachineSpendMapper tests that the state machine is able to properly +// map the spend event into a custom event that can be used to trigger a state +// transition. +func TestStateMachineSpendMapper(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 RegisterSpendNtfn call when we send the event. + targetOutpoint := &wire.OutPoint{Hash: chainhash.Hash{3}} + targetPkScript := []byte{0x03} + targetHeightHint := uint32(300) + adapters.On( + "RegisterSpendNtfn", targetOutpoint, targetPkScript, + targetHeightHint, + ).Return(nil) + + // Send the event that triggers RegisterSpend emission. + stateMachine.SendEvent(ctx, ®isterSpend{}) + + // 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 spend event coming back from the notifier. Populate + // it with some data to be mapped. + simulatedSpend := &chainntnfs.SpendDetail{ + SpentOutPoint: targetOutpoint, + SpenderTxHash: &chainhash.Hash{4}, + SpendingTx: &wire.MsgTx{}, + SpendingHeight: 456, + } + adapters.spendChan <- simulatedSpend + + // This should trigger the mapper and send the spendDetailsEvent, + // transitioning us to the spent state. + expectedStates = []State[dummyEvents, *dummyEnv]{&dummyStateSpent{}} + assertStateTransitions(t, stateSub, expectedStates) + + // Final state assertion. + finalState := assertState(t, &stateMachine, &dummyStateSpent{}) + + // Assert that the details from the spend event were correctly + // propagated to the final state. + require.Equal(t, + *simulatedSpend.SpenderTxHash, finalState.spenderTxHash, + ) + require.Equal(t, + simulatedSpend.SpendingHeight, finalState.spendingHeight, + ) + + adapters.AssertExpectations(t) + env.AssertExpectations(t) +} + type dummyMsgMapper struct { mock.Mock }