Skip to content

Commit 6545c4f

Browse files
committed
reservation: update package to use new fsm context
This commit updates the reservation package to use the new fsm context instead of the old fsm context. This is only a first step in the process of migrating the reservation package to the new fsm context. The next step should be to remove the stored context in the reservation manager.
1 parent 7b00bae commit 6545c4f

File tree

4 files changed

+39
-40
lines changed

4 files changed

+39
-40
lines changed

instantout/reservation/actions.go

Lines changed: 18 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -22,15 +22,17 @@ type InitReservationContext struct {
2222
// InitAction is the action that is executed when the reservation state machine
2323
// is initialized. It creates the reservation in the database and dispatches the
2424
// payment to the server.
25-
func (f *FSM) InitAction(eventCtx fsm.EventContext) fsm.EventType {
25+
func (f *FSM) InitAction(ctx context.Context, eventCtx fsm.EventContext,
26+
) fsm.EventType {
27+
2628
// Check if the context is of the correct type.
2729
reservationRequest, ok := eventCtx.(*InitReservationContext)
2830
if !ok {
2931
return f.HandleError(fsm.ErrInvalidContextType)
3032
}
3133

3234
keyRes, err := f.cfg.Wallet.DeriveNextKey(
33-
f.ctx, KeyFamily,
35+
ctx, KeyFamily,
3436
)
3537
if err != nil {
3638
return f.HandleError(err)
@@ -45,7 +47,7 @@ func (f *FSM) InitAction(eventCtx fsm.EventContext) fsm.EventType {
4547
ClientKey: keyRes.PubKey.SerializeCompressed(),
4648
}
4749

48-
_, err = f.cfg.ReservationClient.OpenReservation(f.ctx, request)
50+
_, err = f.cfg.ReservationClient.OpenReservation(ctx, request)
4951
if err != nil {
5052
return f.HandleError(err)
5153
}
@@ -66,7 +68,7 @@ func (f *FSM) InitAction(eventCtx fsm.EventContext) fsm.EventType {
6668
f.reservation = reservation
6769

6870
// Create the reservation in the database.
69-
err = f.cfg.Store.CreateReservation(f.ctx, reservation)
71+
err = f.cfg.Store.CreateReservation(ctx, reservation)
7072
if err != nil {
7173
return f.HandleError(err)
7274
}
@@ -77,13 +79,15 @@ func (f *FSM) InitAction(eventCtx fsm.EventContext) fsm.EventType {
7779
// SubscribeToConfirmationAction is the action that is executed when the
7880
// reservation is waiting for confirmation. It subscribes to the confirmation
7981
// of the reservation transaction.
80-
func (f *FSM) SubscribeToConfirmationAction(_ fsm.EventContext) fsm.EventType {
82+
func (f *FSM) SubscribeToConfirmationAction(ctx context.Context,
83+
_ fsm.EventContext) fsm.EventType {
84+
8185
pkscript, err := f.reservation.GetPkScript()
8286
if err != nil {
8387
return f.HandleError(err)
8488
}
8589

86-
callCtx, cancel := context.WithCancel(f.ctx)
90+
callCtx, cancel := context.WithCancel(ctx)
8791
defer cancel()
8892

8993
// Subscribe to the confirmation of the reservation transaction.
@@ -141,7 +145,7 @@ func (f *FSM) SubscribeToConfirmationAction(_ fsm.EventContext) fsm.EventType {
141145
return OnTimedOut
142146
}
143147

144-
case <-f.ctx.Done():
148+
case <-ctx.Done():
145149
return fsm.NoOp
146150
}
147151
}
@@ -150,10 +154,10 @@ func (f *FSM) SubscribeToConfirmationAction(_ fsm.EventContext) fsm.EventType {
150154
// AsyncWaitForExpiredOrSweptAction waits for the reservation to be either
151155
// expired or swept. This is non-blocking and can be used to wait for the
152156
// reservation to expire while expecting other events.
153-
func (f *FSM) AsyncWaitForExpiredOrSweptAction(_ fsm.EventContext,
154-
) fsm.EventType {
157+
func (f *FSM) AsyncWaitForExpiredOrSweptAction(ctx context.Context,
158+
_ fsm.EventContext) fsm.EventType {
155159

156-
notifCtx, cancel := context.WithCancel(f.ctx)
160+
notifCtx, cancel := context.WithCancel(ctx)
157161

158162
blockHeightChan, errEpochChan, err := f.cfg.ChainNotifier.
159163
RegisterBlockEpochNtfn(notifCtx)
@@ -184,13 +188,13 @@ func (f *FSM) AsyncWaitForExpiredOrSweptAction(_ fsm.EventContext,
184188
errSpendChan,
185189
)
186190
if err != nil {
187-
f.handleAsyncError(err)
191+
f.handleAsyncError(ctx, err)
188192
return
189193
}
190194
if op == fsm.NoOp {
191195
return
192196
}
193-
err = f.SendEvent(op, nil)
197+
err = f.SendEvent(ctx, op, nil)
194198
if err != nil {
195199
f.Errorf("Error sending %s event: %v", op, err)
196200
}
@@ -229,10 +233,10 @@ func (f *FSM) handleSubcriptions(ctx context.Context,
229233
}
230234
}
231235

232-
func (f *FSM) handleAsyncError(err error) {
236+
func (f *FSM) handleAsyncError(ctx context.Context, err error) {
233237
f.LastActionError = err
234238
f.Errorf("Error on async action: %v", err)
235-
err2 := f.SendEvent(fsm.OnError, err)
239+
err2 := f.SendEvent(ctx, fsm.OnError, err)
236240
if err2 != nil {
237241
f.Errorf("Error sending event: %v", err2)
238242
}

instantout/reservation/actions_test.go

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -144,7 +144,6 @@ func TestInitReservationAction(t *testing.T) {
144144
).Return(tc.mockStoreErr)
145145

146146
reservationFSM := &FSM{
147-
ctx: ctxb,
148147
cfg: &Config{
149148
Wallet: mockLnd.WalletKit,
150149
ChainNotifier: mockLnd.ChainNotifier,
@@ -154,7 +153,7 @@ func TestInitReservationAction(t *testing.T) {
154153
StateMachine: &fsm.StateMachine{},
155154
}
156155

157-
event := reservationFSM.InitAction(tc.eventCtx)
156+
event := reservationFSM.InitAction(ctxb, tc.eventCtx)
158157
require.Equal(t, tc.expectedEvent, event)
159158
}
160159
}
@@ -227,10 +226,10 @@ func TestSubscribeToConfirmationAction(t *testing.T) {
227226
tc := tc
228227
t.Run(tc.name, func(t *testing.T) {
229228
chainNotifier := new(MockChainNotifier)
230-
229+
ctxb := context.Background()
231230
// Create the FSM.
232231
r := NewFSMFromReservation(
233-
context.Background(), &Config{
232+
&Config{
234233
ChainNotifier: chainNotifier,
235234
},
236235
&Reservation{
@@ -296,7 +295,7 @@ func TestSubscribeToConfirmationAction(t *testing.T) {
296295
}
297296
}()
298297

299-
eventType := r.SubscribeToConfirmationAction(nil)
298+
eventType := r.SubscribeToConfirmationAction(ctxb, nil)
300299
// Assert that the return value is as expected
301300
require.Equal(t, tc.expectedEvent, eventType)
302301

@@ -335,10 +334,11 @@ func TestAsyncWaitForExpiredOrSweptAction(t *testing.T) {
335334
tc := tc
336335
t.Run(tc.name, func(t *testing.T) { // Create a mock ChainNotifier and Reservation
337336
chainNotifier := new(MockChainNotifier)
337+
ctxb := context.Background()
338338

339339
// Define your FSM
340340
r := NewFSMFromReservation(
341-
context.Background(), &Config{
341+
&Config{
342342
ChainNotifier: chainNotifier,
343343
},
344344
&Reservation{
@@ -361,7 +361,7 @@ func TestAsyncWaitForExpiredOrSweptAction(t *testing.T) {
361361
make(chan error), tc.spendErr,
362362
)
363363

364-
eventType := r.AsyncWaitForExpiredOrSweptAction(nil)
364+
eventType := r.AsyncWaitForExpiredOrSweptAction(ctxb, nil)
365365
// Assert that the return value is as expected
366366
require.Equal(t, tc.expectedEvent, eventType)
367367
})
@@ -415,7 +415,7 @@ func TestHandleSubcriptions(t *testing.T) {
415415

416416
// Create the FSM.
417417
r := NewFSMFromReservation(
418-
context.Background(), &Config{
418+
&Config{
419419
ChainNotifier: chainNotifier,
420420
},
421421
&Reservation{

instantout/reservation/fsm.go

Lines changed: 7 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -40,26 +40,21 @@ type FSM struct {
4040
cfg *Config
4141

4242
reservation *Reservation
43-
44-
ctx context.Context
4543
}
4644

4745
// NewFSM creates a new reservation FSM.
48-
func NewFSM(ctx context.Context, cfg *Config) *FSM {
46+
func NewFSM(cfg *Config) *FSM {
4947
reservation := &Reservation{
5048
State: fsm.EmptyState,
5149
}
5250

53-
return NewFSMFromReservation(ctx, cfg, reservation)
51+
return NewFSMFromReservation(cfg, reservation)
5452
}
5553

5654
// NewFSMFromReservation creates a new reservation FSM from an existing
5755
// reservation recovered from the database.
58-
func NewFSMFromReservation(ctx context.Context, cfg *Config,
59-
reservation *Reservation) *FSM {
60-
56+
func NewFSMFromReservation(cfg *Config, reservation *Reservation) *FSM {
6157
reservationFsm := &FSM{
62-
ctx: ctx,
6358
cfg: cfg,
6459
reservation: reservation,
6560
}
@@ -206,7 +201,9 @@ func (f *FSM) GetReservationStates() fsm.States {
206201

207202
// updateReservation updates the reservation in the database. This function
208203
// is called after every new state transition.
209-
func (r *FSM) updateReservation(notification fsm.Notification) {
204+
func (r *FSM) updateReservation(ctx context.Context,
205+
notification fsm.Notification) {
206+
210207
if r.reservation == nil {
211208
return
212209
}
@@ -229,7 +226,7 @@ func (r *FSM) updateReservation(notification fsm.Notification) {
229226
return
230227
}
231228

232-
err := r.cfg.Store.UpdateReservation(r.ctx, r.reservation)
229+
err := r.cfg.Store.UpdateReservation(ctx, r.reservation)
233230
if err != nil {
234231
r.Errorf("unable to update reservation: %v", err)
235232
}

instantout/reservation/manager.go

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -111,9 +111,7 @@ func (m *Manager) newReservation(ctx context.Context, currentHeight uint32,
111111
// Create the reservation state machine. We need to pass in the runCtx
112112
// of the reservation manager so that the state machine will keep on
113113
// running even if the grpc conte
114-
reservationFSM := NewFSM(
115-
ctx, m.cfg,
116-
)
114+
reservationFSM := NewFSM(m.cfg)
117115

118116
// Add the reservation to the active reservations map.
119117
m.Lock()
@@ -130,7 +128,7 @@ func (m *Manager) newReservation(ctx context.Context, currentHeight uint32,
130128

131129
// Send the init event to the state machine.
132130
go func() {
133-
err = reservationFSM.SendEvent(OnServerRequest, initContext)
131+
err = reservationFSM.SendEvent(m.runCtx, OnServerRequest, initContext)
134132
if err != nil {
135133
log.Errorf("Error sending init event: %v", err)
136134
}
@@ -172,15 +170,15 @@ func (m *Manager) RecoverReservations(ctx context.Context) error {
172170
fsmCtx := context.WithValue(ctx, reservation.ID, nil)
173171

174172
reservationFSM := NewFSMFromReservation(
175-
fsmCtx, m.cfg, reservation,
173+
m.cfg, reservation,
176174
)
177175

178176
m.activeReservations[reservation.ID] = reservationFSM
179177

180178
// As SendEvent can block, we'll start a goroutine to process
181179
// the event.
182180
go func() {
183-
err := reservationFSM.SendEvent(OnRecover, nil)
181+
err := reservationFSM.SendEvent(fsmCtx, OnRecover, nil)
184182
if err != nil {
185183
log.Errorf("FSM %v Error sending recover "+
186184
"event %v, state: %v",
@@ -217,7 +215,7 @@ func (m *Manager) LockReservation(ctx context.Context, id ID) error {
217215
}
218216

219217
// Try to send the lock event to the reservation.
220-
err := reservation.SendEvent(OnLocked, nil)
218+
err := reservation.SendEvent(ctx, OnLocked, nil)
221219
if err != nil {
222220
return err
223221
}
@@ -237,7 +235,7 @@ func (m *Manager) UnlockReservation(ctx context.Context, id ID) error {
237235
}
238236

239237
// Try to send the unlock event to the reservation.
240-
err := reservation.SendEvent(OnUnlocked, nil)
238+
err := reservation.SendEvent(ctx, OnUnlocked, nil)
241239
if err != nil && strings.Contains(err.Error(), "config error") {
242240
// If the error is a config error, we can ignore it, as the
243241
// reservation is already unlocked.

0 commit comments

Comments
 (0)