Skip to content

Commit efcd134

Browse files
ubettigoleubettigole
authored
refactor(storage): replace batch dependent partial update (#494)
## Summary Pass complete batch dependent entities through the storage contract while preserving controller-owned version arithmetic and slice immutability on failure. ## Test Plan Unit tested ## Issues
1 parent 09cfd80 commit efcd134

7 files changed

Lines changed: 214 additions & 39 deletions

File tree

‎submitqueue/extension/storage/batch_dependent_store.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -37,8 +37,8 @@ type BatchDependentStore interface {
3737
// Returns ErrAlreadyExists if the entry already exists for the given batch ID.
3838
Create(ctx context.Context, batchDependent entity.BatchDependent) error
3939

40-
// UpdateDependents updates the dependents of a batch dependent and the version to newVersion
40+
// Update replaces the non-key fields of a batch dependent and persists newVersion
4141
// if the current persisted version matches oldVersion. If versions do not match, returns ErrVersionMismatch.
4242
// Version arithmetic is owned by the caller; the store performs a pure conditional write.
43-
UpdateDependents(ctx context.Context, batchID string, oldVersion, newVersion int32, dependents []string) error
43+
Update(ctx context.Context, batchDependent entity.BatchDependent, oldVersion, newVersion int32) error
4444
}

‎submitqueue/extension/storage/mock/batch_dependent_store_mock.go‎

Lines changed: 6 additions & 6 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎submitqueue/extension/storage/mysql/batch_dependent_store.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -91,41 +91,41 @@ func (s *batchDependentStore) Create(ctx context.Context, batchDependent entity.
9191
return nil
9292
}
9393

94-
// UpdateDependents updates the dependents of a batch dependent and the version to newVersion
94+
// Update replaces the non-key fields of a batch dependent and persists newVersion
9595
// if the current persisted version matches oldVersion. If versions do not match, returns ErrVersionMismatch.
9696
// Version arithmetic is owned by the caller; this is a pure conditional write.
97-
func (s *batchDependentStore) UpdateDependents(ctx context.Context, batchID string, oldVersion, newVersion int32, dependents []string) (retErr error) {
97+
func (s *batchDependentStore) Update(ctx context.Context, batchDependent entity.BatchDependent, oldVersion, newVersion int32) (retErr error) {
9898
op := metrics.Begin(s.scope, "update_dependents", metrics.StorageLatencyBuckets)
9999
defer func() { op.Complete(retErr) }()
100100

101-
dependentsJSON, err := json.Marshal(dependents)
101+
dependentsJSON, err := json.Marshal(batchDependent.Dependents)
102102
if err != nil {
103-
return fmt.Errorf("failed to marshal dependents batchID=%s for UpdateDependents batch dependent entity: %w", batchID, err)
103+
return fmt.Errorf("failed to marshal dependents batchID=%s for Update batch dependent entity: %w", batchDependent.BatchID, err)
104104
}
105105

106106
result, err := s.db.ExecContext(ctx,
107107
"UPDATE batch_dependent SET dependents = ?, version = ? WHERE batch_id = ? AND version = ?",
108-
dependentsJSON, newVersion, batchID, oldVersion,
108+
dependentsJSON, newVersion, batchDependent.BatchID, oldVersion,
109109
)
110110
if err != nil {
111111
return fmt.Errorf(
112112
"failed to update batch dependent dependents for batchID=%q oldVersion=%d newVersion=%d: %w",
113-
batchID, oldVersion, newVersion, err,
113+
batchDependent.BatchID, oldVersion, newVersion, err,
114114
)
115115
}
116116

117117
rowsAffected, err := result.RowsAffected()
118118
if err != nil {
119119
return fmt.Errorf(
120120
"failed to get rows affected from update for batchID=%q oldVersion=%d newVersion=%d: %w",
121-
batchID, oldVersion, newVersion, err,
121+
batchDependent.BatchID, oldVersion, newVersion, err,
122122
)
123123
}
124124

125125
if rowsAffected != 1 {
126126
return fmt.Errorf(
127127
"version mismatch for batch dependent update: batchID=%q expected_version=%d: %w",
128-
batchID, oldVersion, storage.ErrVersionMismatch,
128+
batchDependent.BatchID, oldVersion, storage.ErrVersionMismatch,
129129
)
130130
}
131131

‎submitqueue/extension/storage/mysql/batch_dependent_store_test.go‎

Lines changed: 76 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,37 @@ func TestBatchDependentStore_Get(t *testing.T) {
7070
},
7171
want: want,
7272
},
73+
{
74+
name: "found with nil dependents",
75+
batchID: "monorepo/batch/nil",
76+
setup: func(mock sqlmock.Sqlmock) {
77+
rows := sqlmock.NewRows([]string{"batch_id", "dependents", "version"}).
78+
AddRow("monorepo/batch/nil", []byte("null"), int32(2))
79+
mock.ExpectQuery("SELECT batch_id, dependents, version FROM batch_dependent").
80+
WithArgs("monorepo/batch/nil").
81+
WillReturnRows(rows)
82+
},
83+
want: entity.BatchDependent{
84+
BatchID: "monorepo/batch/nil",
85+
Version: 2,
86+
},
87+
},
88+
{
89+
name: "found with empty dependents",
90+
batchID: "monorepo/batch/empty",
91+
setup: func(mock sqlmock.Sqlmock) {
92+
rows := sqlmock.NewRows([]string{"batch_id", "dependents", "version"}).
93+
AddRow("monorepo/batch/empty", []byte("[]"), int32(3))
94+
mock.ExpectQuery("SELECT batch_id, dependents, version FROM batch_dependent").
95+
WithArgs("monorepo/batch/empty").
96+
WillReturnRows(rows)
97+
},
98+
want: entity.BatchDependent{
99+
BatchID: "monorepo/batch/empty",
100+
Dependents: []string{},
101+
Version: 3,
102+
},
103+
},
73104
{
74105
name: "not found",
75106
batchID: "missing",
@@ -178,49 +209,82 @@ func TestBatchDependentStore_Create(t *testing.T) {
178209
}
179210
}
180211

181-
func TestBatchDependentStore_UpdateDependents(t *testing.T) {
182-
const batchID = "monorepo/batch/1"
212+
func TestBatchDependentStore_Update(t *testing.T) {
183213
const oldVersion, newVersion = int32(1), int32(2)
184-
dependents := []string{"monorepo/batch/2", "monorepo/batch/3"}
214+
batchDependent := entity.BatchDependent{
215+
BatchID: "monorepo/batch/1",
216+
Dependents: []string{"monorepo/batch/2", "monorepo/batch/3"},
217+
Version: oldVersion,
218+
}
185219

186220
tests := []struct {
187221
name string
222+
entity entity.BatchDependent
188223
setup func(mock sqlmock.Sqlmock)
189224
wantErr bool
190225
wantErrIs error
191226
}{
192227
{
193-
name: "success",
228+
name: "success",
229+
entity: batchDependent,
230+
setup: func(mock sqlmock.Sqlmock) {
231+
mock.ExpectExec("UPDATE batch_dependent").
232+
WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion).
233+
WillReturnResult(sqlmock.NewResult(0, 1))
234+
},
235+
},
236+
{
237+
name: "success with nil dependents",
238+
entity: entity.BatchDependent{
239+
BatchID: "monorepo/batch/nil",
240+
Version: oldVersion,
241+
},
242+
setup: func(mock sqlmock.Sqlmock) {
243+
mock.ExpectExec("UPDATE batch_dependent").
244+
WithArgs([]byte("null"), newVersion, "monorepo/batch/nil", oldVersion).
245+
WillReturnResult(sqlmock.NewResult(0, 1))
246+
},
247+
},
248+
{
249+
name: "success with empty dependents",
250+
entity: entity.BatchDependent{
251+
BatchID: "monorepo/batch/empty",
252+
Dependents: []string{},
253+
Version: oldVersion,
254+
},
194255
setup: func(mock sqlmock.Sqlmock) {
195256
mock.ExpectExec("UPDATE batch_dependent").
196-
WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion).
257+
WithArgs([]byte("[]"), newVersion, "monorepo/batch/empty", oldVersion).
197258
WillReturnResult(sqlmock.NewResult(0, 1))
198259
},
199260
},
200261
{
201-
name: "version mismatch",
262+
name: "version mismatch",
263+
entity: batchDependent,
202264
setup: func(mock sqlmock.Sqlmock) {
203265
mock.ExpectExec("UPDATE batch_dependent").
204-
WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion).
266+
WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion).
205267
WillReturnResult(sqlmock.NewResult(0, 0))
206268
},
207269
wantErr: true,
208270
wantErrIs: storage.ErrVersionMismatch,
209271
},
210272
{
211-
name: "exec error",
273+
name: "exec error",
274+
entity: batchDependent,
212275
setup: func(mock sqlmock.Sqlmock) {
213276
mock.ExpectExec("UPDATE batch_dependent").
214-
WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion).
277+
WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion).
215278
WillReturnError(fmt.Errorf("connection reset"))
216279
},
217280
wantErr: true,
218281
},
219282
{
220-
name: "rows affected error",
283+
name: "rows affected error",
284+
entity: batchDependent,
221285
setup: func(mock sqlmock.Sqlmock) {
222286
mock.ExpectExec("UPDATE batch_dependent").
223-
WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion).
287+
WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion).
224288
WillReturnResult(sqlmock.NewErrorResult(fmt.Errorf("driver error")))
225289
},
226290
wantErr: true,
@@ -234,7 +298,7 @@ func TestBatchDependentStore_UpdateDependents(t *testing.T) {
234298

235299
tt.setup(mock)
236300

237-
err := store.UpdateDependents(context.Background(), batchID, oldVersion, newVersion, dependents)
301+
err := store.Update(context.Background(), tt.entity, oldVersion, newVersion)
238302
if tt.wantErr {
239303
require.Error(t, err)
240304
if tt.wantErrIs != nil {

‎submitqueue/orchestrator/controller/batch/batch.go‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -325,9 +325,11 @@ func (c *Controller) populateBatch(ctx context.Context, batch entity.Batch) (ent
325325
return entity.Batch{}, fmt.Errorf("failed to get batch dependent for batchID=%s: %w", dependencyID, err)
326326
}
327327

328-
dependents := append(existing.Dependents, batch.ID)
328+
updated := existing
329+
updated.Dependents = append([]string(nil), existing.Dependents...)
330+
updated.Dependents = append(updated.Dependents, batch.ID)
329331
newVersion := existing.Version + 1
330-
if err := c.store.GetBatchDependentStore().UpdateDependents(ctx, dependencyID, existing.Version, newVersion, dependents); err != nil {
332+
if err := c.store.GetBatchDependentStore().Update(ctx, updated, existing.Version, newVersion); err != nil {
331333
metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1)
332334
return entity.Batch{}, fmt.Errorf("failed to update batch dependent index for existing batchID=%s and new batchID=%s: %w", dependencyID, batch.ID, err)
333335
}

‎submitqueue/orchestrator/controller/batch/batch_test.go‎

Lines changed: 75 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -357,14 +357,22 @@ func TestController_Process_WithDependencies(t *testing.T) {
357357
BatchID: "test-queue/batch/1",
358358
Version: 1,
359359
}, nil)
360-
mockBatchDependentStore.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), gomock.Any()).Return(nil)
360+
mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
361+
BatchID: "test-queue/batch/1",
362+
Dependents: []string{"test-queue/batch/1"},
363+
Version: 1,
364+
}, int32(1), int32(2)).Return(nil)
361365
// batch/2 already has an existing dependent.
362366
mockBatchDependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/2").Return(entity.BatchDependent{
363367
BatchID: "test-queue/batch/2",
364368
Dependents: []string{"test-queue/batch/99"},
365369
Version: 2,
366370
}, nil)
367-
mockBatchDependentStore.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/2", int32(2), int32(3), gomock.Any()).Return(nil)
371+
mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
372+
BatchID: "test-queue/batch/2",
373+
Dependents: []string{"test-queue/batch/99", "test-queue/batch/1"},
374+
Version: 2,
375+
}, int32(2), int32(3)).Return(nil)
368376
// Create empty reverse index for the new batch.
369377
mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
370378

@@ -418,7 +426,11 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) {
418426
BatchID: "test-queue/batch/2",
419427
Version: 5,
420428
}, nil)
421-
mockBatchDependentStore.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/2", int32(5), int32(6), gomock.Any()).Return(nil)
429+
mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
430+
BatchID: "test-queue/batch/2",
431+
Dependents: []string{"test-queue/batch/1"},
432+
Version: 5,
433+
}, int32(5), int32(6)).Return(nil)
422434
mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
423435

424436
mockReqStore := storagemock.NewMockRequestStore(ctrl)
@@ -457,6 +469,56 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) {
457469
require.NoError(t, err)
458470
}
459471

472+
func TestController_Process_BatchDependentUpdateFailureDoesNotMutateFetchedDependents(t *testing.T) {
473+
ctrl := gomock.NewController(t)
474+
475+
request := testRequest()
476+
activeBatch := entity.Batch{
477+
ID: "test-queue/batch/99",
478+
Queue: "test-queue",
479+
State: entity.BatchStateCreated,
480+
Version: 1,
481+
}
482+
dependents := make([]string, 1, 2)
483+
dependents[0] = "test-queue/batch/98"
484+
existing := entity.BatchDependent{
485+
BatchID: activeBatch.ID,
486+
Dependents: dependents,
487+
Version: 4,
488+
}
489+
490+
mockBatchStore := storagemock.NewMockBatchStore(ctrl)
491+
mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, gomock.Any()).Return([]entity.Batch{activeBatch}, nil)
492+
493+
mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
494+
mockBatchDependentStore.EXPECT().Get(gomock.Any(), activeBatch.ID).Return(existing, nil)
495+
mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
496+
BatchID: activeBatch.ID,
497+
Dependents: []string{"test-queue/batch/98", "test-queue/batch/1"},
498+
Version: existing.Version,
499+
}, existing.Version, existing.Version+1).Return(errors.New("update failed"))
500+
501+
mockReqStore := storagemock.NewMockRequestStore(ctrl)
502+
mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil)
503+
504+
mockStorage := storagemock.NewMockStorage(ctrl)
505+
mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes()
506+
mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes()
507+
mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes()
508+
509+
controller := newTestController(t, ctrl, newSequentialCounter(ctrl), mockStorage, nil, nil)
510+
511+
msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil)
512+
delivery := queuemock.NewMockDelivery(ctrl)
513+
delivery.EXPECT().Message().Return(msg).AnyTimes()
514+
delivery.EXPECT().Attempt().Return(1).AnyTimes()
515+
516+
err := controller.Process(context.Background(), delivery)
517+
require.Error(t, err)
518+
assert.Equal(t, "", dependents[:cap(dependents)][1])
519+
assert.Equal(t, int32(4), existing.Version)
520+
}
521+
460522
func TestController_Process_AnalyzerFailure(t *testing.T) {
461523
ctrl := gomock.NewController(t)
462524

@@ -894,9 +956,11 @@ func TestController_PopulateBatch_Errors(t *testing.T) {
894956
BatchID: "test-queue/batch/0",
895957
Version: 2,
896958
}, nil)
897-
dependentStore.EXPECT().UpdateDependents(
898-
gomock.Any(), "test-queue/batch/0", int32(2), int32(3), []string{batch.ID},
899-
).Return(storeErr)
959+
dependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
960+
BatchID: "test-queue/batch/0",
961+
Dependents: []string{batch.ID},
962+
Version: 2,
963+
}, int32(2), int32(3)).Return(storeErr)
900964
},
901965
errMsg: "failed to update batch dependent index",
902966
},
@@ -909,9 +973,11 @@ func TestController_PopulateBatch_Errors(t *testing.T) {
909973
Dependents: []string{"test-queue/batch/old"},
910974
Version: 2,
911975
}, nil)
912-
dependentStore.EXPECT().UpdateDependents(
913-
gomock.Any(), "test-queue/batch/0", int32(2), int32(3), []string{"test-queue/batch/old", batch.ID},
914-
).Return(nil)
976+
dependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{
977+
BatchID: "test-queue/batch/0",
978+
Dependents: []string{"test-queue/batch/old", batch.ID},
979+
Version: 2,
980+
}, int32(2), int32(3)).Return(nil)
915981
batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(storeErr)
916982
},
917983
errMsg: "failed to mark batch",

0 commit comments

Comments
 (0)