From cb46f3daa07a3ab0bdebe011edf4f84c997f71b4 Mon Sep 17 00:00:00 2001 From: Olaoluwa Osuntokun Date: Wed, 16 Apr 2025 18:05:26 -0700 Subject: [PATCH 1/3] protofsm: add new ConfMapper similar to SpendMapper for conf events In this commit, we add a new ConfMapper which is useful for state machines that want to project some of the conf attributes into a new event to be sent post conf. --- protofsm/daemon_events.go | 11 ++- protofsm/state_machine.go | 19 ++-- protofsm/state_machine_test.go | 157 +++++++++++++++++++++++++++++++++ 3 files changed, 175 insertions(+), 12 deletions(-) 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 } From c580666e284ba1b1f77557d01b494f1615254240 Mon Sep 17 00:00:00 2001 From: Olaoluwa Osuntokun Date: Wed, 16 Apr 2025 18:26:03 -0700 Subject: [PATCH 2/3] protofsm: add unit tests for SpendMapper --- protofsm/state_machine_test.go | 183 +++++++++++++++++++++++++++++++-- 1 file changed, 174 insertions(+), 9 deletions(-) diff --git a/protofsm/state_machine_test.go b/protofsm/state_machine_test.go index 52b27b985..09caa0c97 100644 --- a/protofsm/state_machine_test.go +++ b/protofsm/state_machine_test.go @@ -54,6 +54,20 @@ 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 } @@ -88,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{}, @@ -180,14 +194,59 @@ func (d *dummyStateStart) ProcessEvent(event dummyEvents, env *dummyEnv, // 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. + // We received the mapped confirmation details, transition to + // the confirmed state. return &StateTransition[dummyEvents, *dummyEnv]{ NextState: &dummyStateConfirmed{ - blockHash: eventDetails.blockHash, - blockHeight: eventDetails.blockHeight, + 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 } @@ -240,6 +299,28 @@ 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 +} + func assertState[Event any, Env Environment](t *testing.T, m *StateMachine[Event, Env], expectedState State[Event, Env]) { @@ -565,8 +646,92 @@ func TestStateMachineConfMapper(t *testing.T) { // 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) + require.Equal(t, + *simulatedConf.BlockHash, finalStateDetails.blockHash, + ) + require.Equal(t, + simulatedConf.BlockHeight, finalStateDetails.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, err := stateMachine.CurrentState() + require.NoError(t, err) + require.IsType(t, &dummyStateSpent{}, finalState) + + // Assert that the details from the spend event were correctly + // propagated to the final state. + finalStateDetails := finalState.(*dummyStateSpent) + require.Equal(t, + *simulatedSpend.SpenderTxHash, finalStateDetails.spenderTxHash, + ) + require.Equal(t, + simulatedSpend.SpendingHeight, finalStateDetails.spendingHeight, + ) adapters.AssertExpectations(t) env.AssertExpectations(t) From 946ae4c8043af845aed4cd00285587c0b1649396 Mon Sep 17 00:00:00 2001 From: Olaoluwa Osuntokun Date: Wed, 16 Apr 2025 18:39:51 -0700 Subject: [PATCH 3/3] protofsm: add generic type assertion to state machine tests This commit introduces a new generic type assertion function `assertState` to the state machine tests. This function asserts that the state machine is currently in the expected state type and returns the state cast to that type. This allows us to directly access the fields of the state without having to perform a type assertion manually. --- protofsm/state_machine_test.go | 30 ++++++++++++++++-------------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/protofsm/state_machine_test.go b/protofsm/state_machine_test.go index 09caa0c97..c5b42fefa 100644 --- a/protofsm/state_machine_test.go +++ b/protofsm/state_machine_test.go @@ -321,12 +321,20 @@ func (d *dummyStateSpent) IsTerminal() bool { return true } -func assertState[Event any, Env Environment](t *testing.T, - m *StateMachine[Event, Env], expectedState State[Event, Env]) { +// 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]( @@ -639,18 +647,15 @@ func TestStateMachineConfMapper(t *testing.T) { assertStateTransitions(t, stateSub, expectedStates) // Final state assertion. - finalState, err := stateMachine.CurrentState() - require.NoError(t, err) - require.IsType(t, &dummyStateConfirmed{}, finalState) + finalState := assertState(t, &stateMachine, &dummyStateConfirmed{}) // 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, + *simulatedConf.BlockHash, finalState.blockHash, ) require.Equal(t, - simulatedConf.BlockHeight, finalStateDetails.blockHeight, + simulatedConf.BlockHeight, finalState.blockHeight, ) adapters.AssertExpectations(t) @@ -719,18 +724,15 @@ func TestStateMachineSpendMapper(t *testing.T) { assertStateTransitions(t, stateSub, expectedStates) // Final state assertion. - finalState, err := stateMachine.CurrentState() - require.NoError(t, err) - require.IsType(t, &dummyStateSpent{}, finalState) + finalState := assertState(t, &stateMachine, &dummyStateSpent{}) // Assert that the details from the spend event were correctly // propagated to the final state. - finalStateDetails := finalState.(*dummyStateSpent) require.Equal(t, - *simulatedSpend.SpenderTxHash, finalStateDetails.spenderTxHash, + *simulatedSpend.SpenderTxHash, finalState.spenderTxHash, ) require.Equal(t, - simulatedSpend.SpendingHeight, finalStateDetails.spendingHeight, + simulatedSpend.SpendingHeight, finalState.spendingHeight, ) adapters.AssertExpectations(t)