Skip to content

Commit d17e737

Browse files
committed
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.
1 parent 7f69ceb commit d17e737

File tree

2 files changed

+71
-7
lines changed

2 files changed

+71
-7
lines changed

protofsm/state_machine.go

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -135,17 +135,21 @@ type StateMachine[Event any, Env Environment] struct {
135135

136136
stateQuery chan stateQuery[Event, Env]
137137

138+
initEvent fn.Option[DaemonEvent]
139+
138140
startOnce sync.Once
139141
stopOnce sync.Once
140142

141143
// TODO(roasbeef): also use that context guard here?
142144
}
143145

144146
// NewStateMachine creates a new state machine given a set of daemon adapters,
145-
// an initial state, and an environment.
147+
// an initial state, an environment, and an event to process as if emitted at
148+
// the onset of the state machine. Such an event can be used to set up tracking
149+
// state such as a txid confirmation event.
146150
func NewStateMachine[Event any, Env Environment](adapters DaemonAdapters,
147-
initialState State[Event, Env],
148-
env Env) StateMachine[Event, Env] {
151+
initialState State[Event, Env], env Env,
152+
initEvent fn.Option[DaemonEvent]) StateMachine[Event, Env] {
149153

150154
return StateMachine[Event, Env]{
151155
daemon: adapters,
@@ -154,6 +158,7 @@ func NewStateMachine[Event any, Env Environment](adapters DaemonAdapters,
154158
stateQuery: make(chan stateQuery[Event, Env]),
155159
quit: make(chan struct{}),
156160
env: env,
161+
initEvent: initEvent,
157162
newStateEvents: fn.NewEventDistributor[State[Event, Env]](),
158163
}
159164
}
@@ -443,6 +448,9 @@ func (s *StateMachine[Event, Env]) applyEvents(newEvent Event) (State[Event, Env
443448
currentState = transition.NextState
444449

445450
// Notify our subscribers of the new state transition.
451+
//
452+
// TODO(roasbeef): will only give us the outer state?
453+
// * let FSMs choose which state to emit?
446454
s.newStateEvents.NotifySubscribers(currentState)
447455

448456
return nil
@@ -464,6 +472,16 @@ func (s *StateMachine[Event, Env]) driveMachine() {
464472
// TODO(roasbeef): move into env? read only to start with
465473
currentState := s.currentState
466474

475+
// Before we start, if we have an init daemon event specified, then
476+
// we'll handle that now.
477+
err := fn.MapOptionZ(s.initEvent, func(event DaemonEvent) error {
478+
return s.executeDaemonEvent(event)
479+
})
480+
if err != nil {
481+
log.Errorf("unable to execute init event: %w", err)
482+
return
483+
}
484+
467485
// We just started driving the state machine, so we'll notify our
468486
// subscribers of this starting state.
469487
s.newStateEvents.NotifySubscribers(currentState)

protofsm/state_machine_test.go

Lines changed: 50 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -211,6 +211,52 @@ func (d *dummyAdapters) RegisterSpendNtfn(outpoint *wire.OutPoint,
211211
}, err
212212
}
213213

214+
// TestStateMachineOnInitDaemonEvent tests that the state machine will properly
215+
// execute any init-level daemon events passed into it.
216+
func TestStateMachineOnInitDaemonEvent(t *testing.T) {
217+
// First, we'll create our state machine given the env, and our
218+
// starting state.
219+
env := &dummyEnv{}
220+
startingState := &dummyStateStart{}
221+
222+
adapters := newDaemonAdapters()
223+
224+
// We'll make an init event that'll send to a peer, then transition us
225+
// to our terminal state.
226+
initEvent := &SendMsgEvent[dummyEvents]{
227+
TargetPeer: *pub1,
228+
PostSendEvent: fn.Some(dummyEvents(&goToFin{})),
229+
}
230+
231+
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
232+
adapters, startingState, env, fn.Some[DaemonEvent](initEvent),
233+
)
234+
235+
// Before we start up the state machine, we'll assert that the send
236+
// message adapter is called on start up.
237+
adapters.On("SendMessages", *pub1, mock.Anything).Return(nil)
238+
239+
stateMachine.Start()
240+
defer stateMachine.Stop()
241+
242+
// As we're triggering internal events, we'll also subscribe to the set
243+
// of new states so we can assert as we go.
244+
stateSub := stateMachine.RegisterStateEvents()
245+
defer stateMachine.RemoveStateSub(stateSub)
246+
247+
// Assert that we go from the starting state to the final state. The
248+
// state machine should now also be on the final terminal state.
249+
expectedStates := []State[dummyEvents, *dummyEnv]{
250+
&dummyStateStart{}, &dummyStateFin{},
251+
}
252+
assertStateTransitions(t, stateSub, expectedStates)
253+
254+
// We'll now assert that after the daemon was started, the send message
255+
// adapter was called above as specified in the init event.
256+
adapters.AssertExpectations(t)
257+
env.AssertExpectations(t)
258+
}
259+
214260
// TestStateMachineInternalEvents tests that the state machine is able to add
215261
// new internal events to the event queue for further processing during a state
216262
// transition.
@@ -222,10 +268,10 @@ func TestStateMachineInternalEvents(t *testing.T) {
222268
env := &dummyEnv{}
223269
startingState := &dummyStateStart{}
224270

225-
adapters := &dummyAdapters{}
271+
adapters := newDaemonAdapters()
226272

227273
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
228-
adapters, startingState, env,
274+
adapters, startingState, env, fn.None[DaemonEvent](),
229275
)
230276
stateMachine.Start()
231277
defer stateMachine.Stop()
@@ -269,10 +315,10 @@ func TestStateMachineDaemonEvents(t *testing.T) {
269315
canSend: &boolTrigger,
270316
}
271317

272-
adapters := &dummyAdapters{}
318+
adapters := newDaemonAdapters()
273319

274320
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
275-
adapters, startingState, env,
321+
adapters, startingState, env, fn.None[DaemonEvent](),
276322
)
277323
stateMachine.Start()
278324
defer stateMachine.Stop()

0 commit comments

Comments
 (0)