Merge pull request #9725 from Roasbeef/protofsm-conf-mapper

protofsm: add ConfMapper to allow conf attribute projection for new events

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.

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

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](
@@ -415,6 +589,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, &registerConf{})
+
+	// 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, &registerSpend{})
+
+	// 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
 }