Skip to content

Commit c751df0

Browse files
ubettigoleubettigole
committed
refactor(storage): replace batch state updates
Persist complete Batch entities through guarded updates and migrate controllers to failure-safe candidate copies. Jira Issues CODEM-204
1 parent e5625cd commit c751df0

18 files changed

Lines changed: 306 additions & 99 deletions

File tree

‎submitqueue/extension/storage/batch_store.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,10 @@ type BatchStore interface {
3131
// Returns ErrAlreadyExists if a batch with the same ID already exists.
3232
Create(ctx context.Context, batch entity.Batch) error
3333

34-
// UpdateState updates the state of a batch to newState and the version to newVersion
34+
// Update replaces every non-key field of a batch and writes newVersion
3535
// if the current persisted version matches oldVersion. If versions do not match, returns ErrVersionMismatch.
3636
// Version arithmetic is owned by the caller; the store performs a pure conditional write.
37-
UpdateState(ctx context.Context, id string, oldVersion, newVersion int32, newState entity.BatchState) error
37+
Update(ctx context.Context, batch entity.Batch, oldVersion, newVersion int32) error
3838

3939
// GetByQueueAndStates retrieves all batches that belong to the given queue and are in the given states.
4040
GetByQueueAndStates(ctx context.Context, queue string, states []entity.BatchState) ([]entity.Batch, error)

‎submitqueue/extension/storage/mock/batch_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_store.go‎

Lines changed: 18 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -102,36 +102,46 @@ func (s *batchStore) Create(ctx context.Context, batch entity.Batch) (retErr err
102102
return nil
103103
}
104104

105-
// UpdateState updates the state of a batch to newState and the version to newVersion
105+
// Update replaces every non-key field of a batch and writes newVersion
106106
// if the current persisted version matches oldVersion. If versions do not match, returns ErrVersionMismatch.
107107
// Version arithmetic is owned by the caller; this is a pure conditional write.
108-
func (s *batchStore) UpdateState(ctx context.Context, id string, oldVersion, newVersion int32, newState entity.BatchState) (retErr error) {
108+
func (s *batchStore) Update(ctx context.Context, batch entity.Batch, oldVersion, newVersion int32) (retErr error) {
109109
op := metrics.Begin(s.scope, "update_state", metrics.StorageLatencyBuckets)
110110
defer func() { op.Complete(retErr) }()
111111

112+
containsJSON, err := json.Marshal(batch.Contains)
113+
if err != nil {
114+
return fmt.Errorf("failed to marshal contains=%v id=%s for Update batch entity: %w", batch.Contains, batch.ID, err)
115+
}
116+
117+
dependenciesJSON, err := json.Marshal(batch.Dependencies)
118+
if err != nil {
119+
return fmt.Errorf("failed to marshal dependencies=%v id=%s for Update batch entity: %w", batch.Dependencies, batch.ID, err)
120+
}
121+
112122
result, err := s.db.ExecContext(ctx,
113-
"UPDATE batch SET state = ?, version = ? WHERE id = ? AND version = ?",
114-
newState, newVersion, id, oldVersion,
123+
"UPDATE batch SET queue = ?, contains = ?, dependencies = ?, state = ?, version = ? WHERE id = ? AND version = ?",
124+
batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion,
115125
)
116126
if err != nil {
117127
return fmt.Errorf(
118-
"failed to update batch state for id=%q oldVersion=%d newVersion=%d newState=%v: %w",
119-
id, oldVersion, newVersion, newState, err,
128+
"failed to update batch for id=%q oldVersion=%d newVersion=%d newState=%v: %w",
129+
batch.ID, oldVersion, newVersion, batch.State, err,
120130
)
121131
}
122132

123133
rowsAffected, err := result.RowsAffected()
124134
if err != nil {
125135
return fmt.Errorf(
126136
"failed to get rows affected from update for id=%q oldVersion=%d newVersion=%d newState=%v: %w",
127-
id, oldVersion, newVersion, newState, err,
137+
batch.ID, oldVersion, newVersion, batch.State, err,
128138
)
129139
}
130140

131141
if rowsAffected != 1 {
132142
return fmt.Errorf(
133143
"version mismatch for batch update: id=%q expected_version=%d newState=%v: %w",
134-
id, oldVersion, newState, storage.ErrVersionMismatch,
144+
batch.ID, oldVersion, batch.State, storage.ErrVersionMismatch,
135145
)
136146
}
137147

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

Lines changed: 57 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -198,53 +198,98 @@ func TestBatchStore_Create(t *testing.T) {
198198
}
199199
}
200200

201-
func TestBatchStore_UpdateState(t *testing.T) {
202-
const id = "monorepo/batch/1"
201+
func TestBatchStore_Update(t *testing.T) {
203202
const oldVersion, newVersion = int32(1), int32(2)
204-
const newState = entity.BatchStateMerging
203+
batch := entity.Batch{
204+
ID: "monorepo/batch/1",
205+
Queue: "monorepo-updated",
206+
Contains: []string{"monorepo/3", "monorepo/4"},
207+
Dependencies: []string{"monorepo/batch/1", "monorepo/batch/2"},
208+
State: entity.BatchStateMerging,
209+
Version: oldVersion,
210+
}
211+
containsJSON, err := json.Marshal(batch.Contains)
212+
require.NoError(t, err)
213+
dependenciesJSON, err := json.Marshal(batch.Dependencies)
214+
require.NoError(t, err)
205215

206216
tests := []struct {
207217
name string
218+
batch entity.Batch
208219
setup func(mock sqlmock.Sqlmock)
209220
wantErr bool
210221
wantErrIs error
211222
}{
212223
{
213-
name: "success",
224+
name: "success",
225+
batch: batch,
214226
setup: func(mock sqlmock.Sqlmock) {
215227
mock.ExpectExec("UPDATE batch").
216-
WithArgs(newState, newVersion, id, oldVersion).
228+
WithArgs(batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion).
217229
WillReturnResult(sqlmock.NewResult(0, 1))
218230
},
219231
},
220232
{
221-
name: "version mismatch",
233+
name: "version mismatch",
234+
batch: batch,
222235
setup: func(mock sqlmock.Sqlmock) {
223236
mock.ExpectExec("UPDATE batch").
224-
WithArgs(newState, newVersion, id, oldVersion).
237+
WithArgs(batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion).
225238
WillReturnResult(sqlmock.NewResult(0, 0))
226239
},
227240
wantErr: true,
228241
wantErrIs: storage.ErrVersionMismatch,
229242
},
230243
{
231-
name: "exec error",
244+
name: "exec error",
245+
batch: batch,
232246
setup: func(mock sqlmock.Sqlmock) {
233247
mock.ExpectExec("UPDATE batch").
234-
WithArgs(newState, newVersion, id, oldVersion).
248+
WithArgs(batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion).
235249
WillReturnError(fmt.Errorf("connection reset"))
236250
},
237251
wantErr: true,
238252
},
239253
{
240-
name: "rows affected error",
254+
name: "rows affected error",
255+
batch: batch,
241256
setup: func(mock sqlmock.Sqlmock) {
242257
mock.ExpectExec("UPDATE batch").
243-
WithArgs(newState, newVersion, id, oldVersion).
258+
WithArgs(batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion).
244259
WillReturnResult(sqlmock.NewErrorResult(fmt.Errorf("driver error")))
245260
},
246261
wantErr: true,
247262
},
263+
{
264+
name: "nil collections",
265+
batch: entity.Batch{
266+
ID: batch.ID,
267+
Queue: batch.Queue,
268+
State: batch.State,
269+
Version: batch.Version,
270+
},
271+
setup: func(mock sqlmock.Sqlmock) {
272+
mock.ExpectExec("UPDATE batch").
273+
WithArgs(batch.Queue, []byte("null"), []byte("null"), batch.State, newVersion, batch.ID, oldVersion).
274+
WillReturnResult(sqlmock.NewResult(0, 1))
275+
},
276+
},
277+
{
278+
name: "empty collections",
279+
batch: entity.Batch{
280+
ID: batch.ID,
281+
Queue: batch.Queue,
282+
Contains: []string{},
283+
Dependencies: []string{},
284+
State: batch.State,
285+
Version: batch.Version,
286+
},
287+
setup: func(mock sqlmock.Sqlmock) {
288+
mock.ExpectExec("UPDATE batch").
289+
WithArgs(batch.Queue, []byte("[]"), []byte("[]"), batch.State, newVersion, batch.ID, oldVersion).
290+
WillReturnResult(sqlmock.NewResult(0, 1))
291+
},
292+
},
248293
}
249294

250295
for _, tt := range tests {
@@ -254,7 +299,7 @@ func TestBatchStore_UpdateState(t *testing.T) {
254299

255300
tt.setup(mock)
256301

257-
err := store.UpdateState(context.Background(), id, oldVersion, newVersion, newState)
302+
err := store.Update(context.Background(), tt.batch, oldVersion, newVersion)
258303
if tt.wantErr {
259304
require.Error(t, err)
260305
if tt.wantErrIs != nil {

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -338,12 +338,12 @@ func (c *Controller) populateBatch(ctx context.Context, batch entity.Batch) (ent
338338
// The batch's own reverse-index row now exists and every dependency lists this batch as a dependent.
339339
// Structural initialization is complete, so transition Creating → Created to make the batch ready for processing once published to speculate.
340340
newVersion := batch.Version + 1
341-
if err := c.store.GetBatchStore().UpdateState(ctx, batch.ID, batch.Version, newVersion, entity.BatchStateCreated); err != nil {
341+
batch.State = entity.BatchStateCreated
342+
if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil {
342343
metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1)
343344
return entity.Batch{}, fmt.Errorf("failed to mark batch %s created: %w", batch.ID, err)
344345
}
345346
batch.Version = newVersion
346-
batch.State = entity.BatchStateCreated
347347
return batch, nil
348348
}
349349

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

Lines changed: 56 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,11 @@ import (
4141
"go.uber.org/zap/zaptest"
4242
)
4343

44+
func batchWithState(batch entity.Batch, state entity.BatchState) entity.Batch {
45+
batch.State = state
46+
return batch
47+
}
48+
4449
// requestIDPayload serializes a RequestID to JSON bytes for test message payloads.
4550
func requestIDPayload(t *testing.T, id string) []byte {
4651
payload, err := entity.RequestID{ID: id}.ToBytes()
@@ -86,7 +91,7 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M
8691
mockBatchStore := storagemock.NewMockBatchStore(ctrl)
8792
mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
8893
mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
89-
mockBatchStore.EXPECT().UpdateState(gomock.Any(), gomock.Any(), int32(1), int32(2), entity.BatchStateCreated).Return(nil).AnyTimes()
94+
mockBatchStore.EXPECT().Update(gomock.Any(), gomock.Any(), int32(1), int32(2)).Return(nil).AnyTimes()
9095

9196
mockReqStore := storagemock.NewMockRequestStore(ctrl)
9297
req := testRequest()
@@ -173,7 +178,14 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) {
173178
mockBatchStore := storagemock.NewMockBatchStore(ctrl)
174179
mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
175180
mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
176-
mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
181+
mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{
182+
ID: "test-queue/batch/1",
183+
Queue: request.Queue,
184+
Contains: []string{request.ID},
185+
Dependencies: []string{},
186+
State: entity.BatchStateCreated,
187+
Version: 1,
188+
}, int32(1), int32(2)).Return(nil)
177189

178190
mockReqStore := storagemock.NewMockRequestStore(ctrl)
179191
mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil)
@@ -349,7 +361,14 @@ func TestController_Process_WithDependencies(t *testing.T) {
349361
mockBatchStore := storagemock.NewMockBatchStore(ctrl)
350362
mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(activeBatches, nil)
351363
mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
352-
mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
364+
mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{
365+
ID: "test-queue/batch/1",
366+
Queue: request.Queue,
367+
Contains: []string{request.ID},
368+
Dependencies: []string{"test-queue/batch/1", "test-queue/batch/2"},
369+
State: entity.BatchStateCreated,
370+
Version: 1,
371+
}, int32(1), int32(2)).Return(nil)
353372

354373
mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
355374
// batch/1 has no existing dependents.
@@ -418,7 +437,14 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) {
418437
mockBatchStore := storagemock.NewMockBatchStore(ctrl)
419438
mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(activeBatches, nil)
420439
mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
421-
mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
440+
mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{
441+
ID: "test-queue/batch/1",
442+
Queue: request.Queue,
443+
Contains: []string{request.ID},
444+
Dependencies: []string{"test-queue/batch/2"},
445+
State: entity.BatchStateCreated,
446+
Version: 1,
447+
}, int32(1), int32(2)).Return(nil)
422448

423449
mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
424450
// Only batch/2 is selected by the analyzer, so only it gets a reverse-index update.
@@ -712,7 +738,14 @@ func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) {
712738
mockBatchStore := storagemock.NewMockBatchStore(ctrl)
713739
mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil)
714740
mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
715-
mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
741+
mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{
742+
ID: "test-queue/batch/1",
743+
Queue: request.Queue,
744+
Contains: []string{request.ID},
745+
Dependencies: []string{},
746+
State: entity.BatchStateCreated,
747+
Version: 1,
748+
}, int32(1), int32(2)).Return(nil)
716749

717750
mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
718751
mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
@@ -784,7 +817,7 @@ func TestController_Process_ReadiesBatchBeforePublishing(t *testing.T) {
784817
Dependents: []string{},
785818
Version: 1,
786819
}).Return(nil),
787-
batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(nil),
820+
batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateCreated), int32(1), int32(2)).Return(nil),
788821
publisher.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).Return(nil),
789822
publisher.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).Return(nil),
790823
)
@@ -846,8 +879,22 @@ func TestController_Process_RedeliveryMintsFreshBatchID(t *testing.T) {
846879
return nil
847880
},
848881
).Times(2)
849-
batchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
850-
batchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/2", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
882+
batchStore.EXPECT().Update(gomock.Any(), entity.Batch{
883+
ID: "test-queue/batch/1",
884+
Queue: firstRequest.Queue,
885+
Contains: []string{firstRequest.ID},
886+
Dependencies: []string{},
887+
State: entity.BatchStateCreated,
888+
Version: 1,
889+
}, int32(1), int32(2)).Return(nil)
890+
batchStore.EXPECT().Update(gomock.Any(), entity.Batch{
891+
ID: "test-queue/batch/2",
892+
Queue: firstRequest.Queue,
893+
Contains: []string{firstRequest.ID},
894+
Dependencies: []string{},
895+
State: entity.BatchStateCreated,
896+
Version: 1,
897+
}, int32(1), int32(2)).Return(nil)
851898

852899
batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
853900
batchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).Times(2)
@@ -978,7 +1025,7 @@ func TestController_PopulateBatch_Errors(t *testing.T) {
9781025
Dependents: []string{"test-queue/batch/old", batch.ID},
9791026
Version: 2,
9801027
}, int32(2), int32(3)).Return(nil)
981-
batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(storeErr)
1028+
batchStore.EXPECT().Update(gomock.Any(), batchWithState(batch, entity.BatchStateCreated), int32(1), int32(2)).Return(storeErr)
9821029
},
9831030
errMsg: "failed to mark batch",
9841031
},

‎submitqueue/orchestrator/controller/cancel/cancel.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -312,7 +312,8 @@ func (c *Controller) cancelBatch(ctx context.Context, batch entity.Batch) error
312312

313313
if batch.State != entity.BatchStateCancelling {
314314
newVersion := batch.Version + 1
315-
if err := c.store.GetBatchStore().UpdateState(ctx, batch.ID, batch.Version, newVersion, entity.BatchStateCancelling); err != nil {
315+
batch.State = entity.BatchStateCancelling
316+
if err := c.store.GetBatchStore().Update(ctx, batch, batch.Version, newVersion); err != nil {
316317
metrics.NamedCounter(c.metricsScope, opName, "batch_update_errors", 1)
317318
// storage.ErrVersionMismatch here means the batch advanced concurrently
318319
// (e.g. speculate / merge progressed). Returned as-is because the
@@ -322,7 +323,6 @@ func (c *Controller) cancelBatch(ctx context.Context, batch entity.Batch) error
322323
return fmt.Errorf("failed to mark batch %s as cancelling: %w", batch.ID, err)
323324
}
324325
batch.Version = newVersion
325-
batch.State = entity.BatchStateCancelling
326326
metrics.NamedCounter(c.metricsScope, opName, "batch_cancelling", 1)
327327
} else {
328328
metrics.NamedCounter(c.metricsScope, opName, "batch_already_cancelling", 1)

0 commit comments

Comments
 (0)