Skip to content

Commit 4430627

Browse files
sputn1cksputn1ck
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 4430627

File tree

4 files changed

+39
-47
lines changed

4 files changed

+39
-47
lines changed

‎instantout/reservation/actions.go

Lines changed: 18 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -22,16 +22,16 @@ 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,
26+
eventCtx fsm.EventContext) 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

32-
keyRes, err := f.cfg.Wallet.DeriveNextKey(
33-
f.ctx, KeyFamily,
34-
)
34+
keyRes, err := f.cfg.Wallet.DeriveNextKey(ctx, KeyFamily)
3535
if err != nil {
3636
return f.HandleError(err)
3737
}
@@ -45,7 +45,7 @@ func (f *FSM) InitAction(eventCtx fsm.EventContext) fsm.EventType {
4545
ClientKey: keyRes.PubKey.SerializeCompressed(),
4646
}
4747

48-
_, err = f.cfg.ReservationClient.OpenReservation(f.ctx, request)
48+
_, err = f.cfg.ReservationClient.OpenReservation(ctx, request)
4949
if err != nil {
5050
return f.HandleError(err)
5151
}
@@ -66,7 +66,7 @@ func (f *FSM) InitAction(eventCtx fsm.EventContext) fsm.EventType {
6666
f.reservation = reservation
6767

6868
// Create the reservation in the database.
69-
err = f.cfg.Store.CreateReservation(f.ctx, reservation)
69+
err = f.cfg.Store.CreateReservation(ctx, reservation)
7070
if err != nil {
7171
return f.HandleError(err)
7272
}
@@ -77,13 +77,15 @@ func (f *FSM) InitAction(eventCtx fsm.EventContext) fsm.EventType {
7777
// SubscribeToConfirmationAction is the action that is executed when the
7878
// reservation is waiting for confirmation. It subscribes to the confirmation
7979
// of the reservation transaction.
80-
func (f *FSM) SubscribeToConfirmationAction(_ fsm.EventContext) fsm.EventType {
80+
func (f *FSM) SubscribeToConfirmationAction(ctx context.Context,
81+
_ fsm.EventContext) fsm.EventType {
82+
8183
pkscript, err := f.reservation.GetPkScript()
8284
if err != nil {
8385
return f.HandleError(err)
8486
}
8587

86-
callCtx, cancel := context.WithCancel(f.ctx)
88+
callCtx, cancel := context.WithCancel(ctx)
8789
defer cancel()
8890

8991
// Subscribe to the confirmation of the reservation transaction.
@@ -141,7 +143,7 @@ func (f *FSM) SubscribeToConfirmationAction(_ fsm.EventContext) fsm.EventType {
141143
return OnTimedOut
142144
}
143145

144-
case <-f.ctx.Done():
146+
case <-ctx.Done():
145147
return fsm.NoOp
146148
}
147149
}
@@ -150,10 +152,10 @@ func (f *FSM) SubscribeToConfirmationAction(_ fsm.EventContext) fsm.EventType {
150152
// AsyncWaitForExpiredOrSweptAction waits for the reservation to be either
151153
// expired or swept. This is non-blocking and can be used to wait for the
152154
// reservation to expire while expecting other events.
153-
func (f *FSM) AsyncWaitForExpiredOrSweptAction(_ fsm.EventContext,
154-
) fsm.EventType {
155+
func (f *FSM) AsyncWaitForExpiredOrSweptAction(ctx context.Context,
156+
_ fsm.EventContext) fsm.EventType {
155157

156-
notifCtx, cancel := context.WithCancel(f.ctx)
158+
notifCtx, cancel := context.WithCancel(ctx)
157159

158160
blockHeightChan, errEpochChan, err := f.cfg.ChainNotifier.
159161
RegisterBlockEpochNtfn(notifCtx)
@@ -184,13 +186,13 @@ func (f *FSM) AsyncWaitForExpiredOrSweptAction(_ fsm.EventContext,
184186
errSpendChan,
185187
)
186188
if err != nil {
187-
f.handleAsyncError(err)
189+
f.handleAsyncError(ctx, err)
188190
return
189191
}
190192
if op == fsm.NoOp {
191193
return
192194
}
193-
err = f.SendEvent(op, nil)
195+
err = f.SendEvent(ctx, op, nil)
194196
if err != nil {
195197
f.Errorf("Error sending %s event: %v", op, err)
196198
}
@@ -229,10 +231,10 @@ func (f *FSM) handleSubcriptions(ctx context.Context,
229231
}
230232
}
231233

232-
func (f *FSM) handleAsyncError(err error) {
234+
func (f *FSM) handleAsyncError(ctx context.Context, err error) {
233235
f.LastActionError = err
234236
f.Errorf("Error on async action: %v", err)
235-
err2 := f.SendEvent(fsm.OnError, err)
237+
err2 := f.SendEvent(ctx, fsm.OnError, err)
236238
if err2 != nil {
237239
f.Errorf("Error sending event: %v", err2)
238240
}

‎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 & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,6 @@ type Manager struct {
2222
// activeReservations contains all the active reservationsFSMs.
2323
activeReservations map[ID]*FSM
2424

25-
runCtx context.Context
26-
2725
sync.Mutex
2826
}
2927

@@ -42,7 +40,6 @@ func (m *Manager) Run(ctx context.Context, height int32) error {
4240
runCtx, cancel := context.WithCancel(ctx)
4341
defer cancel()
4442

45-
m.runCtx = runCtx
4643
currentHeight := height
4744

4845
err := m.RecoverReservations(runCtx)
@@ -111,9 +108,7 @@ func (m *Manager) newReservation(ctx context.Context, currentHeight uint32,
111108
// Create the reservation state machine. We need to pass in the runCtx
112109
// of the reservation manager so that the state machine will keep on
113110
// running even if the grpc conte
114-
reservationFSM := NewFSM(
115-
ctx, m.cfg,
116-
)
111+
reservationFSM := NewFSM(m.cfg)
117112

118113
// Add the reservation to the active reservations map.
119114
m.Lock()
@@ -130,7 +125,7 @@ func (m *Manager) newReservation(ctx context.Context, currentHeight uint32,
130125

131126
// Send the init event to the state machine.
132127
go func() {
133-
err = reservationFSM.SendEvent(OnServerRequest, initContext)
128+
err = reservationFSM.SendEvent(ctx, OnServerRequest, initContext)
134129
if err != nil {
135130
log.Errorf("Error sending init event: %v", err)
136131
}
@@ -171,16 +166,14 @@ func (m *Manager) RecoverReservations(ctx context.Context) error {
171166

172167
fsmCtx := context.WithValue(ctx, reservation.ID, nil)
173168

174-
reservationFSM := NewFSMFromReservation(
175-
fsmCtx, m.cfg, reservation,
176-
)
169+
reservationFSM := NewFSMFromReservation(m.cfg, reservation)
177170

178171
m.activeReservations[reservation.ID] = reservationFSM
179172

180173
// As SendEvent can block, we'll start a goroutine to process
181174
// the event.
182175
go func() {
183-
err := reservationFSM.SendEvent(OnRecover, nil)
176+
err := reservationFSM.SendEvent(fsmCtx, OnRecover, nil)
184177
if err != nil {
185178
log.Errorf("FSM %v Error sending recover "+
186179
"event %v, state: %v",
@@ -217,7 +210,7 @@ func (m *Manager) LockReservation(ctx context.Context, id ID) error {
217210
}
218211

219212
// Try to send the lock event to the reservation.
220-
err := reservation.SendEvent(OnLocked, nil)
213+
err := reservation.SendEvent(ctx, OnLocked, nil)
221214
if err != nil {
222215
return err
223216
}
@@ -237,7 +230,7 @@ func (m *Manager) UnlockReservation(ctx context.Context, id ID) error {
237230
}
238231

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

0 commit comments

Comments
 (0)