Skip to content

Commit 2ee238f

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 795a527 commit 2ee238f

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
@@ -141,17 +141,21 @@ type StateMachine[Event any, Env Environment] struct {
141141

142142
stateQuery chan stateQuery[Event, Env]
143143

144+
initEvent fn.Option[DaemonEvent]
145+
144146
startOnce sync.Once
145147
stopOnce sync.Once
146148

147149
// TODO(roasbeef): also use that context guard here?
148150
}
149151

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

156160
return StateMachine[Event, Env]{
157161
daemon: adapters,
@@ -160,6 +164,7 @@ func NewStateMachine[Event any, Env Environment](adapters DaemonAdapters,
160164
stateQuery: make(chan stateQuery[Event, Env]),
161165
quit: make(chan struct{}),
162166
env: env,
167+
initEvent: initEvent,
163168
newStateEvents: fn.NewEventDistributor[State[Event, Env]](),
164169
}
165170
}
@@ -459,6 +464,9 @@ func (s *StateMachine[Event, Env]) applyEvents(newEvent Event) (State[Event, Env
459464
currentState = transition.NextState
460465

461466
// Notify our subscribers of the new state transition.
467+
//
468+
// TODO(roasbeef): will only give us the outer state?
469+
// * let FSMs choose which state to emit?
462470
s.newStateEvents.NotifySubscribers(currentState)
463471

464472
return nil
@@ -480,6 +488,16 @@ func (s *StateMachine[Event, Env]) driveMachine() {
480488
// TODO(roasbeef): move into env? read only to start with
481489
currentState := s.currentState
482490

491+
// Before we start, if we have an init daemon event specified, then
492+
// we'll handle that now.
493+
err := fn.MapOptionZ(s.initEvent, func(event DaemonEvent) error {
494+
return s.executeDaemonEvent(event)
495+
})
496+
if err != nil {
497+
log.Errorf("unable to execute init event: %w", err)
498+
return
499+
}
500+
483501
// We just started driving the state machine, so we'll notify our
484502
// subscribers of this starting state.
485503
s.newStateEvents.NotifySubscribers(currentState)

protofsm/state_machine_test.go

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

231-
adapters := &dummyAdapters{}
231+
adapters := newDaemonAdapters()
232232

233233
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
234-
adapters, startingState, env,
234+
adapters, startingState, env, fn.None[DaemonEvent](),
235235
)
236236
stateMachine.Start()
237237
defer stateMachine.Stop()
@@ -251,6 +251,53 @@ func TestStateMachineTerminateCleanup(t *testing.T) {
251251
env.AssertExpectations(t)
252252
}
253253

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

265-
adapters := &dummyAdapters{}
312+
adapters := newDaemonAdapters()
266313

267314
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
268-
adapters, startingState, env,
315+
adapters, startingState, env, fn.None[DaemonEvent](),
269316
)
270317
stateMachine.Start()
271318
defer stateMachine.Stop()
@@ -313,10 +360,10 @@ func TestStateMachineDaemonEvents(t *testing.T) {
313360
canSend: &boolTrigger,
314361
}
315362

316-
adapters := &dummyAdapters{}
363+
adapters := newDaemonAdapters()
317364

318365
stateMachine := NewStateMachine[dummyEvents, *dummyEnv](
319-
adapters, startingState, env,
366+
adapters, startingState, env, fn.None[DaemonEvent](),
320367
)
321368
stateMachine.Start()
322369
defer stateMachine.Stop()

0 commit comments

Comments
 (0)