Skip to content

Commit 14820fa

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 31e7023 commit 14820fa

File tree

2 files changed

+74
-9
lines changed

2 files changed

+74
-9
lines changed

protofsm/state_machine.go

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

139139
stateQuery chan stateQuery[Event, Env]
140140

141+
initEvent fn.Option[DaemonEvent]
142+
141143
startOnce sync.Once
142144
stopOnce sync.Once
143145

144146
// TODO(roasbeef): also use that context guard here?
145147
}
146148

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

153157
return StateMachine[Event, Env]{
154158
daemon: adapters,
@@ -157,6 +161,7 @@ func NewStateMachine[Event any, Env Environment](adapters DaemonAdapters,
157161
stateQuery: make(chan stateQuery[Event, Env]),
158162
quit: make(chan struct{}),
159163
env: env,
164+
initEvent: initEvent,
160165
newStateEvents: fn.NewEventDistributor[State[Event, Env]](),
161166
}
162167
}
@@ -446,6 +451,9 @@ func (s *StateMachine[Event, Env]) applyEvents(newEvent Event) (State[Event, Env
446451
currentState = transition.NextState
447452

448453
// Notify our subscribers of the new state transition.
454+
//
455+
// TODO(roasbeef): will only give us the outer state?
456+
// * let FSMs choose which state to emit?
449457
s.newStateEvents.NotifySubscribers(currentState)
450458

451459
return nil
@@ -467,6 +475,16 @@ func (s *StateMachine[Event, Env]) driveMachine() {
467475
// TODO(roasbeef): move into env? read only to start with
468476
currentState := s.currentState
469477

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

protofsm/state_machine_test.go

Lines changed: 53 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -221,10 +221,10 @@ func TestStateMachineTerminateCleanup(t *testing.T) {
221221
env := &dummyEnv{}
222222
startingState := &dummyStateStart{}
223223

224-
adapters := &dummyAdapters{}
224+
adapters := newDaemonAdapters()
225225

226226
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
227-
adapters, startingState, env,
227+
adapters, startingState, env, fn.None[DaemonEvent](),
228228
)
229229
stateMachine.Start()
230230
defer stateMachine.Stop()
@@ -244,6 +244,53 @@ func TestStateMachineTerminateCleanup(t *testing.T) {
244244
env.AssertExpectations(t)
245245
}
246246

247+
// TestStateMachineOnInitDaemonEvent tests that the state machine will properly
248+
// execute any init-level daemon events passed into it.
249+
func TestStateMachineOnInitDaemonEvent(t *testing.T) {
250+
// First, we'll create our state machine given the env, and our
251+
// starting state.
252+
env := &dummyEnv{}
253+
startingState := &dummyStateStart{}
254+
255+
adapters := newDaemonAdapters()
256+
257+
// We'll make an init event that'll send to a peer, then transition us
258+
// to our terminal state.
259+
initEvent := &SendMsgEvent[dummyEvents]{
260+
TargetPeer: *pub1,
261+
PostSendEvent: fn.Some(dummyEvents(&goToFin{})),
262+
}
263+
264+
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
265+
adapters, startingState, env, fn.Some[DaemonEvent](initEvent),
266+
)
267+
268+
// Before we start up the state machine, we'll assert that the send
269+
// message adapter is called on start up.
270+
adapters.On("SendMessages", *pub1, mock.Anything).Return(nil)
271+
env.On("CleanUp").Return(nil)
272+
273+
stateMachine.Start()
274+
defer stateMachine.Stop()
275+
276+
// As we're triggering internal events, we'll also subscribe to the set
277+
// of new states so we can assert as we go.
278+
stateSub := stateMachine.RegisterStateEvents()
279+
defer stateMachine.RemoveStateSub(stateSub)
280+
281+
// Assert that we go from the starting state to the final state. The
282+
// state machine should now also be on the final terminal state.
283+
expectedStates := []State[dummyEvents, *dummyEnv]{
284+
&dummyStateStart{}, &dummyStateFin{},
285+
}
286+
assertStateTransitions(t, stateSub, expectedStates)
287+
288+
// We'll now assert that after the daemon was started, the send message
289+
// adapter was called above as specified in the init event.
290+
adapters.AssertExpectations(t)
291+
env.AssertExpectations(t)
292+
}
293+
247294
// TestStateMachineInternalEvents tests that the state machine is able to add
248295
// new internal events to the event queue for further processing during a state
249296
// transition.
@@ -255,10 +302,10 @@ func TestStateMachineInternalEvents(t *testing.T) {
255302
env := &dummyEnv{}
256303
startingState := &dummyStateStart{}
257304

258-
adapters := &dummyAdapters{}
305+
adapters := newDaemonAdapters()
259306

260307
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
261-
adapters, startingState, env,
308+
adapters, startingState, env, fn.None[DaemonEvent](),
262309
)
263310
stateMachine.Start()
264311
defer stateMachine.Stop()
@@ -306,10 +353,10 @@ func TestStateMachineDaemonEvents(t *testing.T) {
306353
canSend: &boolTrigger,
307354
}
308355

309-
adapters := &dummyAdapters{}
356+
adapters := newDaemonAdapters()
310357

311358
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
312-
adapters, startingState, env,
359+
adapters, startingState, env, fn.None[DaemonEvent](),
313360
)
314361
stateMachine.Start()
315362
defer stateMachine.Stop()

0 commit comments

Comments
 (0)