From 3d90e3c0117411729f1c1f63d4245115af6d2653 Mon Sep 17 00:00:00 2001 From: abettigole Date: Fri, 31 Jul 2026 23:40:11 +0000 Subject: [PATCH] refactor(storage): replace batch dependent partial update Pass complete batch dependent entities through the storage contract while preserving controller-owned version arithmetic and slice immutability on failure. Jira Issues CODEM-204 --- .../storage/batch_dependent_store.go | 4 +- .../mock/batch_dependent_store_mock.go | 12 +-- .../storage/mysql/batch_dependent_store.go | 16 ++-- .../mysql/batch_dependent_store_test.go | 88 ++++++++++++++++--- .../orchestrator/controller/batch/batch.go | 6 +- .../controller/batch/batch_test.go | 84 ++++++++++++++++-- .../submitqueue/extension/storage/suite.go | 43 +++++++++ 7 files changed, 214 insertions(+), 39 deletions(-) diff --git a/submitqueue/extension/storage/batch_dependent_store.go b/submitqueue/extension/storage/batch_dependent_store.go index 01cef446..16014017 100644 --- a/submitqueue/extension/storage/batch_dependent_store.go +++ b/submitqueue/extension/storage/batch_dependent_store.go @@ -37,8 +37,8 @@ type BatchDependentStore interface { // Returns ErrAlreadyExists if the entry already exists for the given batch ID. Create(ctx context.Context, batchDependent entity.BatchDependent) error - // UpdateDependents updates the dependents of a batch dependent and the version to newVersion + // Update replaces the non-key fields of a batch dependent and persists newVersion // if the current persisted version matches oldVersion. If versions do not match, returns ErrVersionMismatch. // Version arithmetic is owned by the caller; the store performs a pure conditional write. - UpdateDependents(ctx context.Context, batchID string, oldVersion, newVersion int32, dependents []string) error + Update(ctx context.Context, batchDependent entity.BatchDependent, oldVersion, newVersion int32) error } diff --git a/submitqueue/extension/storage/mock/batch_dependent_store_mock.go b/submitqueue/extension/storage/mock/batch_dependent_store_mock.go index bac3eb26..73ccfff0 100644 --- a/submitqueue/extension/storage/mock/batch_dependent_store_mock.go +++ b/submitqueue/extension/storage/mock/batch_dependent_store_mock.go @@ -70,16 +70,16 @@ func (mr *MockBatchDependentStoreMockRecorder) Get(ctx, batchID any) *gomock.Cal return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockBatchDependentStore)(nil).Get), ctx, batchID) } -// UpdateDependents mocks base method. -func (m *MockBatchDependentStore) UpdateDependents(ctx context.Context, batchID string, oldVersion, newVersion int32, dependents []string) error { +// Update mocks base method. +func (m *MockBatchDependentStore) Update(ctx context.Context, batchDependent entity.BatchDependent, oldVersion, newVersion int32) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpdateDependents", ctx, batchID, oldVersion, newVersion, dependents) + ret := m.ctrl.Call(m, "Update", ctx, batchDependent, oldVersion, newVersion) ret0, _ := ret[0].(error) return ret0 } -// UpdateDependents indicates an expected call of UpdateDependents. -func (mr *MockBatchDependentStoreMockRecorder) UpdateDependents(ctx, batchID, oldVersion, newVersion, dependents any) *gomock.Call { +// Update indicates an expected call of Update. +func (mr *MockBatchDependentStoreMockRecorder) Update(ctx, batchDependent, oldVersion, newVersion any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateDependents", reflect.TypeOf((*MockBatchDependentStore)(nil).UpdateDependents), ctx, batchID, oldVersion, newVersion, dependents) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Update", reflect.TypeOf((*MockBatchDependentStore)(nil).Update), ctx, batchDependent, oldVersion, newVersion) } diff --git a/submitqueue/extension/storage/mysql/batch_dependent_store.go b/submitqueue/extension/storage/mysql/batch_dependent_store.go index be32bc47..3577fc39 100644 --- a/submitqueue/extension/storage/mysql/batch_dependent_store.go +++ b/submitqueue/extension/storage/mysql/batch_dependent_store.go @@ -91,26 +91,26 @@ func (s *batchDependentStore) Create(ctx context.Context, batchDependent entity. return nil } -// UpdateDependents updates the dependents of a batch dependent and the version to newVersion +// Update replaces the non-key fields of a batch dependent and persists newVersion // if the current persisted version matches oldVersion. If versions do not match, returns ErrVersionMismatch. // Version arithmetic is owned by the caller; this is a pure conditional write. -func (s *batchDependentStore) UpdateDependents(ctx context.Context, batchID string, oldVersion, newVersion int32, dependents []string) (retErr error) { +func (s *batchDependentStore) Update(ctx context.Context, batchDependent entity.BatchDependent, oldVersion, newVersion int32) (retErr error) { op := metrics.Begin(s.scope, "update_dependents", metrics.StorageLatencyBuckets) defer func() { op.Complete(retErr) }() - dependentsJSON, err := json.Marshal(dependents) + dependentsJSON, err := json.Marshal(batchDependent.Dependents) if err != nil { - return fmt.Errorf("failed to marshal dependents batchID=%s for UpdateDependents batch dependent entity: %w", batchID, err) + return fmt.Errorf("failed to marshal dependents batchID=%s for Update batch dependent entity: %w", batchDependent.BatchID, err) } result, err := s.db.ExecContext(ctx, "UPDATE batch_dependent SET dependents = ?, version = ? WHERE batch_id = ? AND version = ?", - dependentsJSON, newVersion, batchID, oldVersion, + dependentsJSON, newVersion, batchDependent.BatchID, oldVersion, ) if err != nil { return fmt.Errorf( "failed to update batch dependent dependents for batchID=%q oldVersion=%d newVersion=%d: %w", - batchID, oldVersion, newVersion, err, + batchDependent.BatchID, oldVersion, newVersion, err, ) } @@ -118,14 +118,14 @@ func (s *batchDependentStore) UpdateDependents(ctx context.Context, batchID stri if err != nil { return fmt.Errorf( "failed to get rows affected from update for batchID=%q oldVersion=%d newVersion=%d: %w", - batchID, oldVersion, newVersion, err, + batchDependent.BatchID, oldVersion, newVersion, err, ) } if rowsAffected != 1 { return fmt.Errorf( "version mismatch for batch dependent update: batchID=%q expected_version=%d: %w", - batchID, oldVersion, storage.ErrVersionMismatch, + batchDependent.BatchID, oldVersion, storage.ErrVersionMismatch, ) } diff --git a/submitqueue/extension/storage/mysql/batch_dependent_store_test.go b/submitqueue/extension/storage/mysql/batch_dependent_store_test.go index 5988d72c..cef3c49a 100644 --- a/submitqueue/extension/storage/mysql/batch_dependent_store_test.go +++ b/submitqueue/extension/storage/mysql/batch_dependent_store_test.go @@ -70,6 +70,37 @@ func TestBatchDependentStore_Get(t *testing.T) { }, want: want, }, + { + name: "found with nil dependents", + batchID: "monorepo/batch/nil", + setup: func(mock sqlmock.Sqlmock) { + rows := sqlmock.NewRows([]string{"batch_id", "dependents", "version"}). + AddRow("monorepo/batch/nil", []byte("null"), int32(2)) + mock.ExpectQuery("SELECT batch_id, dependents, version FROM batch_dependent"). + WithArgs("monorepo/batch/nil"). + WillReturnRows(rows) + }, + want: entity.BatchDependent{ + BatchID: "monorepo/batch/nil", + Version: 2, + }, + }, + { + name: "found with empty dependents", + batchID: "monorepo/batch/empty", + setup: func(mock sqlmock.Sqlmock) { + rows := sqlmock.NewRows([]string{"batch_id", "dependents", "version"}). + AddRow("monorepo/batch/empty", []byte("[]"), int32(3)) + mock.ExpectQuery("SELECT batch_id, dependents, version FROM batch_dependent"). + WithArgs("monorepo/batch/empty"). + WillReturnRows(rows) + }, + want: entity.BatchDependent{ + BatchID: "monorepo/batch/empty", + Dependents: []string{}, + Version: 3, + }, + }, { name: "not found", batchID: "missing", @@ -178,49 +209,82 @@ func TestBatchDependentStore_Create(t *testing.T) { } } -func TestBatchDependentStore_UpdateDependents(t *testing.T) { - const batchID = "monorepo/batch/1" +func TestBatchDependentStore_Update(t *testing.T) { const oldVersion, newVersion = int32(1), int32(2) - dependents := []string{"monorepo/batch/2", "monorepo/batch/3"} + batchDependent := entity.BatchDependent{ + BatchID: "monorepo/batch/1", + Dependents: []string{"monorepo/batch/2", "monorepo/batch/3"}, + Version: oldVersion, + } tests := []struct { name string + entity entity.BatchDependent setup func(mock sqlmock.Sqlmock) wantErr bool wantErrIs error }{ { - name: "success", + name: "success", + entity: batchDependent, + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("UPDATE batch_dependent"). + WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + { + name: "success with nil dependents", + entity: entity.BatchDependent{ + BatchID: "monorepo/batch/nil", + Version: oldVersion, + }, + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("UPDATE batch_dependent"). + WithArgs([]byte("null"), newVersion, "monorepo/batch/nil", oldVersion). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + { + name: "success with empty dependents", + entity: entity.BatchDependent{ + BatchID: "monorepo/batch/empty", + Dependents: []string{}, + Version: oldVersion, + }, setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE batch_dependent"). - WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion). + WithArgs([]byte("[]"), newVersion, "monorepo/batch/empty", oldVersion). WillReturnResult(sqlmock.NewResult(0, 1)) }, }, { - name: "version mismatch", + name: "version mismatch", + entity: batchDependent, setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE batch_dependent"). - WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion). + WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion). WillReturnResult(sqlmock.NewResult(0, 0)) }, wantErr: true, wantErrIs: storage.ErrVersionMismatch, }, { - name: "exec error", + name: "exec error", + entity: batchDependent, setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE batch_dependent"). - WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion). + WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion). WillReturnError(fmt.Errorf("connection reset")) }, wantErr: true, }, { - name: "rows affected error", + name: "rows affected error", + entity: batchDependent, setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE batch_dependent"). - WithArgs(sqlmock.AnyArg(), newVersion, batchID, oldVersion). + WithArgs([]byte(`["monorepo/batch/2","monorepo/batch/3"]`), newVersion, batchDependent.BatchID, oldVersion). WillReturnResult(sqlmock.NewErrorResult(fmt.Errorf("driver error"))) }, wantErr: true, @@ -234,7 +298,7 @@ func TestBatchDependentStore_UpdateDependents(t *testing.T) { tt.setup(mock) - err := store.UpdateDependents(context.Background(), batchID, oldVersion, newVersion, dependents) + err := store.Update(context.Background(), tt.entity, oldVersion, newVersion) if tt.wantErr { require.Error(t, err) if tt.wantErrIs != nil { diff --git a/submitqueue/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index f1b24ff6..ac8de908 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/submitqueue/orchestrator/controller/batch/batch.go @@ -325,9 +325,11 @@ func (c *Controller) populateBatch(ctx context.Context, batch entity.Batch) (ent return entity.Batch{}, fmt.Errorf("failed to get batch dependent for batchID=%s: %w", dependencyID, err) } - dependents := append(existing.Dependents, batch.ID) + updated := existing + updated.Dependents = append([]string(nil), existing.Dependents...) + updated.Dependents = append(updated.Dependents, batch.ID) newVersion := existing.Version + 1 - if err := c.store.GetBatchDependentStore().UpdateDependents(ctx, dependencyID, existing.Version, newVersion, dependents); err != nil { + if err := c.store.GetBatchDependentStore().Update(ctx, updated, existing.Version, newVersion); err != nil { metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) return entity.Batch{}, fmt.Errorf("failed to update batch dependent index for existing batchID=%s and new batchID=%s: %w", dependencyID, batch.ID, err) } diff --git a/submitqueue/orchestrator/controller/batch/batch_test.go b/submitqueue/orchestrator/controller/batch/batch_test.go index 7114b361..e0657556 100644 --- a/submitqueue/orchestrator/controller/batch/batch_test.go +++ b/submitqueue/orchestrator/controller/batch/batch_test.go @@ -357,14 +357,22 @@ func TestController_Process_WithDependencies(t *testing.T) { BatchID: "test-queue/batch/1", Version: 1, }, nil) - mockBatchDependentStore.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), gomock.Any()).Return(nil) + mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{ + BatchID: "test-queue/batch/1", + Dependents: []string{"test-queue/batch/1"}, + Version: 1, + }, int32(1), int32(2)).Return(nil) // batch/2 already has an existing dependent. mockBatchDependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/2").Return(entity.BatchDependent{ BatchID: "test-queue/batch/2", Dependents: []string{"test-queue/batch/99"}, Version: 2, }, nil) - mockBatchDependentStore.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/2", int32(2), int32(3), gomock.Any()).Return(nil) + mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{ + BatchID: "test-queue/batch/2", + Dependents: []string{"test-queue/batch/99", "test-queue/batch/1"}, + Version: 2, + }, int32(2), int32(3)).Return(nil) // Create empty reverse index for the new batch. mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) @@ -418,7 +426,11 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) { BatchID: "test-queue/batch/2", Version: 5, }, nil) - mockBatchDependentStore.EXPECT().UpdateDependents(gomock.Any(), "test-queue/batch/2", int32(5), int32(6), gomock.Any()).Return(nil) + mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{ + BatchID: "test-queue/batch/2", + Dependents: []string{"test-queue/batch/1"}, + Version: 5, + }, int32(5), int32(6)).Return(nil) mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockReqStore := storagemock.NewMockRequestStore(ctrl) @@ -457,6 +469,56 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) { require.NoError(t, err) } +func TestController_Process_BatchDependentUpdateFailureDoesNotMutateFetchedDependents(t *testing.T) { + ctrl := gomock.NewController(t) + + request := testRequest() + activeBatch := entity.Batch{ + ID: "test-queue/batch/99", + Queue: "test-queue", + State: entity.BatchStateCreated, + Version: 1, + } + dependents := make([]string, 1, 2) + dependents[0] = "test-queue/batch/98" + existing := entity.BatchDependent{ + BatchID: activeBatch.ID, + Dependents: dependents, + Version: 4, + } + + mockBatchStore := storagemock.NewMockBatchStore(ctrl) + mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, gomock.Any()).Return([]entity.Batch{activeBatch}, nil) + + mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + mockBatchDependentStore.EXPECT().Get(gomock.Any(), activeBatch.ID).Return(existing, nil) + mockBatchDependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{ + BatchID: activeBatch.ID, + Dependents: []string{"test-queue/batch/98", "test-queue/batch/1"}, + Version: existing.Version, + }, existing.Version, existing.Version+1).Return(errors.New("update failed")) + + mockReqStore := storagemock.NewMockRequestStore(ctrl) + mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + + mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() + mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() + + controller := newTestController(t, ctrl, newSequentialCounter(ctrl), mockStorage, nil, nil) + + msg := entityqueue.NewMessage(request.ID, requestIDPayload(t, request.ID), request.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(1).AnyTimes() + + err := controller.Process(context.Background(), delivery) + require.Error(t, err) + assert.Equal(t, "", dependents[:cap(dependents)][1]) + assert.Equal(t, int32(4), existing.Version) +} + func TestController_Process_AnalyzerFailure(t *testing.T) { ctrl := gomock.NewController(t) @@ -894,9 +956,11 @@ func TestController_PopulateBatch_Errors(t *testing.T) { BatchID: "test-queue/batch/0", Version: 2, }, nil) - dependentStore.EXPECT().UpdateDependents( - gomock.Any(), "test-queue/batch/0", int32(2), int32(3), []string{batch.ID}, - ).Return(storeErr) + dependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{ + BatchID: "test-queue/batch/0", + Dependents: []string{batch.ID}, + Version: 2, + }, int32(2), int32(3)).Return(storeErr) }, errMsg: "failed to update batch dependent index", }, @@ -909,9 +973,11 @@ func TestController_PopulateBatch_Errors(t *testing.T) { Dependents: []string{"test-queue/batch/old"}, Version: 2, }, nil) - dependentStore.EXPECT().UpdateDependents( - gomock.Any(), "test-queue/batch/0", int32(2), int32(3), []string{"test-queue/batch/old", batch.ID}, - ).Return(nil) + dependentStore.EXPECT().Update(gomock.Any(), entity.BatchDependent{ + BatchID: "test-queue/batch/0", + Dependents: []string{"test-queue/batch/old", batch.ID}, + Version: 2, + }, int32(2), int32(3)).Return(nil) batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(storeErr) }, errMsg: "failed to mark batch", diff --git a/test/integration/submitqueue/extension/storage/suite.go b/test/integration/submitqueue/extension/storage/suite.go index 8f9df7e7..c5fc2280 100644 --- a/test/integration/submitqueue/extension/storage/suite.go +++ b/test/integration/submitqueue/extension/storage/suite.go @@ -185,6 +185,49 @@ func (s *StorageContractSuite) TestStorage_OptimisticLocking() { assert.Equal(t, int32(2), retrieved.Version) } +// TestStorage_BatchDependentUpdate verifies full updates, collection encoding, and optimistic locking. +func (s *StorageContractSuite) TestStorage_BatchDependentUpdate() { + t := s.T() + ctx := s.ctx + + batchDependent := entity.BatchDependent{ + BatchID: "test/batch-dependent-update", + Dependents: []string{"test/dependent/1"}, + Version: 1, + } + require.NoError(t, s.storage.GetBatchDependentStore().Create(ctx, batchDependent)) + + nilDependents := batchDependent + nilDependents.Dependents = nil + require.NoError(t, s.storage.GetBatchDependentStore().Update(ctx, nilDependents, 1, 2)) + + retrieved, err := s.storage.GetBatchDependentStore().Get(ctx, batchDependent.BatchID) + require.NoError(t, err) + assert.Nil(t, retrieved.Dependents) + assert.Equal(t, int32(2), retrieved.Version) + + emptyDependents := retrieved + emptyDependents.Dependents = []string{} + require.NoError(t, s.storage.GetBatchDependentStore().Update(ctx, emptyDependents, 2, 3)) + + retrieved, err = s.storage.GetBatchDependentStore().Get(ctx, batchDependent.BatchID) + require.NoError(t, err) + assert.Empty(t, retrieved.Dependents) + assert.NotNil(t, retrieved.Dependents) + assert.Equal(t, int32(3), retrieved.Version) + + staleUpdate := retrieved + staleUpdate.Dependents = []string{"test/dependent/stale"} + err = s.storage.GetBatchDependentStore().Update(ctx, staleUpdate, 2, 4) + assert.ErrorIs(t, err, storage.ErrVersionMismatch) + + retrieved, err = s.storage.GetBatchDependentStore().Get(ctx, batchDependent.BatchID) + require.NoError(t, err) + assert.Empty(t, retrieved.Dependents) + assert.NotNil(t, retrieved.Dependents) + assert.Equal(t, int32(3), retrieved.Version) +} + // TestStorage_NotFound tests getting a non-existent request func (s *StorageContractSuite) TestStorage_NotFound() { t := s.T()