refactor(storage): replace batch state updates

submitqueue/entity/batch.go

@@ -156,6 +156,13 @@ type Batch struct {
 	Version int32
 }
 
+// WithState returns a shallow copy of the batch with State replaced.
+// Slice fields continue to share their backing arrays.
+func (b Batch) WithState(state BatchState) Batch {
+	b.State = state
+	return b
+}
+
 // ToBytes serializes the Batch to JSON bytes for queue message payload.
 func (b Batch) ToBytes() ([]byte, error) {
 	return json.Marshal(b)

submitqueue/entity/batch_test.go

@@ -66,6 +66,25 @@ func TestDependencyBatchStates_ExcludesCreating(t *testing.T) {
 	assert.NotContains(t, DependencyBatchStates(), BatchStateCreating)
 }
 
+func TestBatch_WithState(t *testing.T) {
+	batch := Batch{
+		ID:      "queueA/batch/1",
+		Queue:   "queueA",
+		State:   BatchStateCreated,
+		Version: 1,
+	}
+
+	updated := batch.WithState(BatchStateSpeculating)
+
+	assert.Equal(t, BatchStateCreated, batch.State)
+	assert.Equal(t, Batch{
+		ID:      batch.ID,
+		Queue:   batch.Queue,
+		State:   BatchStateSpeculating,
+		Version: batch.Version,
+	}, updated)
+}
+
 func TestBatch_SerializationRoundTrip(t *testing.T) {
 	tests := []struct {
 		name  string

submitqueue/extension/storage/batch_store.go

@@ -31,10 +31,10 @@ type BatchStore interface {
 	// Returns ErrAlreadyExists if a batch with the same ID already exists.
 	Create(ctx context.Context, batch entity.Batch) error
 
-	// UpdateState updates the state of a batch to newState and the version to newVersion
+	// Update replaces every non-key field of a batch and writes 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.
-	UpdateState(ctx context.Context, id string, oldVersion, newVersion int32, newState entity.BatchState) error
+	Update(ctx context.Context, batch entity.Batch, oldVersion, newVersion int32) error
 
 	// GetByQueueAndStates retrieves all batches that belong to the given queue and are in the given states.
 	GetByQueueAndStates(ctx context.Context, queue string, states []entity.BatchState) ([]entity.Batch, error)

submitqueue/extension/storage/mysql/batch_store.go

@@ -102,36 +102,46 @@ func (s *batchStore) Create(ctx context.Context, batch entity.Batch) (retErr err
 	return nil
 }
 
-// UpdateState updates the state of a batch to newState and the version to newVersion
+// Update replaces every non-key field of a batch and writes 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 *batchStore) UpdateState(ctx context.Context, id string, oldVersion, newVersion int32, newState entity.BatchState) (retErr error) {
+func (s *batchStore) Update(ctx context.Context, batch entity.Batch, oldVersion, newVersion int32) (retErr error) {
 	op := metrics.Begin(s.scope, "update_state", metrics.StorageLatencyBuckets)
 	defer func() { op.Complete(retErr) }()
 
+	containsJSON, err := json.Marshal(batch.Contains)
+	if err != nil {
+		return fmt.Errorf("failed to marshal contains=%v id=%s for Update batch entity: %w", batch.Contains, batch.ID, err)
+	}
+
+	dependenciesJSON, err := json.Marshal(batch.Dependencies)
+	if err != nil {
+		return fmt.Errorf("failed to marshal dependencies=%v id=%s for Update batch entity: %w", batch.Dependencies, batch.ID, err)
+	}
+
 	result, err := s.db.ExecContext(ctx,
-		"UPDATE batch SET state = ?, version = ? WHERE id = ? AND version = ?",
-		newState, newVersion, id, oldVersion,
+		"UPDATE batch SET queue = ?, contains = ?, dependencies = ?, state = ?, version = ? WHERE id = ? AND version = ?",
+		batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion,
 	)
 	if err != nil {
 		return fmt.Errorf(
-			"failed to update batch state for id=%q oldVersion=%d newVersion=%d newState=%v: %w",
-			id, oldVersion, newVersion, newState, err,
+			"failed to update batch for id=%q oldVersion=%d newVersion=%d newState=%v: %w",
+			batch.ID, oldVersion, newVersion, batch.State, err,
 		)
 	}
 
 	rowsAffected, err := result.RowsAffected()
 	if err != nil {
 		return fmt.Errorf(
 			"failed to get rows affected from update for id=%q oldVersion=%d newVersion=%d newState=%v: %w",
-			id, oldVersion, newVersion, newState, err,
+			batch.ID, oldVersion, newVersion, batch.State, err,
 		)
 	}
 
 	if rowsAffected != 1 {
 		return fmt.Errorf(
 			"version mismatch for batch update: id=%q expected_version=%d newState=%v: %w",
-			id, oldVersion, newState, storage.ErrVersionMismatch,
+			batch.ID, oldVersion, batch.State, storage.ErrVersionMismatch,
 		)
 	}
 

submitqueue/extension/storage/mysql/batch_store_test.go

@@ -198,53 +198,98 @@ func TestBatchStore_Create(t *testing.T) {
 	}
 }
 
-func TestBatchStore_UpdateState(t *testing.T) {
-	const id = "monorepo/batch/1"
+func TestBatchStore_Update(t *testing.T) {
 	const oldVersion, newVersion = int32(1), int32(2)
-	const newState = entity.BatchStateMerging
+	batch := entity.Batch{
+		ID:           "monorepo/batch/1",
+		Queue:        "monorepo-updated",
+		Contains:     []string{"monorepo/3", "monorepo/4"},
+		Dependencies: []string{"monorepo/batch/1", "monorepo/batch/2"},
+		State:        entity.BatchStateMerging,
+		Version:      oldVersion,
+	}
+	containsJSON, err := json.Marshal(batch.Contains)
+	require.NoError(t, err)
+	dependenciesJSON, err := json.Marshal(batch.Dependencies)
+	require.NoError(t, err)
 
 	tests := []struct {
 		name      string
+		batch     entity.Batch
 		setup     func(mock sqlmock.Sqlmock)
 		wantErr   bool
 		wantErrIs error
 	}{
 		{
-			name: "success",
+			name:  "success",
+			batch: batch,
 			setup: func(mock sqlmock.Sqlmock) {
 				mock.ExpectExec("UPDATE batch").
-					WithArgs(newState, newVersion, id, oldVersion).
+					WithArgs(batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion).
 					WillReturnResult(sqlmock.NewResult(0, 1))
 			},
 		},
 		{
-			name: "version mismatch",
+			name:  "version mismatch",
+			batch: batch,
 			setup: func(mock sqlmock.Sqlmock) {
 				mock.ExpectExec("UPDATE batch").
-					WithArgs(newState, newVersion, id, oldVersion).
+					WithArgs(batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion).
 					WillReturnResult(sqlmock.NewResult(0, 0))
 			},
 			wantErr:   true,
 			wantErrIs: storage.ErrVersionMismatch,
 		},
 		{
-			name: "exec error",
+			name:  "exec error",
+			batch: batch,
 			setup: func(mock sqlmock.Sqlmock) {
 				mock.ExpectExec("UPDATE batch").
-					WithArgs(newState, newVersion, id, oldVersion).
+					WithArgs(batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion).
 					WillReturnError(fmt.Errorf("connection reset"))
 			},
 			wantErr: true,
 		},
 		{
-			name: "rows affected error",
+			name:  "rows affected error",
+			batch: batch,
 			setup: func(mock sqlmock.Sqlmock) {
 				mock.ExpectExec("UPDATE batch").
-					WithArgs(newState, newVersion, id, oldVersion).
+					WithArgs(batch.Queue, containsJSON, dependenciesJSON, batch.State, newVersion, batch.ID, oldVersion).
 					WillReturnResult(sqlmock.NewErrorResult(fmt.Errorf("driver error")))
 			},
 			wantErr: true,
 		},
+		{
+			name: "nil collections",
+			batch: entity.Batch{
+				ID:      batch.ID,
+				Queue:   batch.Queue,
+				State:   batch.State,
+				Version: batch.Version,
+			},
+			setup: func(mock sqlmock.Sqlmock) {
+				mock.ExpectExec("UPDATE batch").
+					WithArgs(batch.Queue, []byte("null"), []byte("null"), batch.State, newVersion, batch.ID, oldVersion).
+					WillReturnResult(sqlmock.NewResult(0, 1))
+			},
+		},
+		{
+			name: "empty collections",
+			batch: entity.Batch{
+				ID:           batch.ID,
+				Queue:        batch.Queue,
+				Contains:     []string{},
+				Dependencies: []string{},
+				State:        batch.State,
+				Version:      batch.Version,
+			},
+			setup: func(mock sqlmock.Sqlmock) {
+				mock.ExpectExec("UPDATE batch").
+					WithArgs(batch.Queue, []byte("[]"), []byte("[]"), batch.State, newVersion, batch.ID, oldVersion).
+					WillReturnResult(sqlmock.NewResult(0, 1))
+			},
+		},
 	}
 
 	for _, tt := range tests {
@@ -254,7 +299,7 @@ func TestBatchStore_UpdateState(t *testing.T) {
 
 			tt.setup(mock)
 
-			err := store.UpdateState(context.Background(), id, oldVersion, newVersion, newState)
+			err := store.Update(context.Background(), tt.batch, oldVersion, newVersion)
 			if tt.wantErr {
 				require.Error(t, err)
 				if tt.wantErrIs != nil {

submitqueue/orchestrator/controller/batch/batch.go

@@ -338,13 +338,13 @@ func (c *Controller) populateBatch(ctx context.Context, batch entity.Batch) (ent
 	// The batch's own reverse-index row now exists and every dependency lists this batch as a dependent.
 	// Structural initialization is complete, so transition Creating → Created to make the batch ready for processing once published to speculate.
 	newVersion := batch.Version + 1
-	if err := c.store.GetBatchStore().UpdateState(ctx, batch.ID, batch.Version, newVersion, entity.BatchStateCreated); err != nil {
+	updated := batch.WithState(entity.BatchStateCreated)
+	if err := c.store.GetBatchStore().Update(ctx, updated, batch.Version, newVersion); err != nil {
 		metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1)
 		return entity.Batch{}, fmt.Errorf("failed to mark batch %s created: %w", batch.ID, err)
 	}
-	batch.Version = newVersion
-	batch.State = entity.BatchStateCreated
-	return batch, nil
+	updated.Version = newVersion
+	return updated, nil
 }
 
 // publish publishes a batch ID to the specified topic key.

submitqueue/orchestrator/controller/batch/batch_test.go

@@ -86,7 +86,7 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M
 		mockBatchStore := storagemock.NewMockBatchStore(ctrl)
 		mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 		mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).AnyTimes()
-		mockBatchStore.EXPECT().UpdateState(gomock.Any(), gomock.Any(), int32(1), int32(2), entity.BatchStateCreated).Return(nil).AnyTimes()
+		mockBatchStore.EXPECT().Update(gomock.Any(), gomock.Any(), int32(1), int32(2)).Return(nil).AnyTimes()
 
 		mockReqStore := storagemock.NewMockRequestStore(ctrl)
 		req := testRequest()
@@ -173,7 +173,14 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) {
 	mockBatchStore := storagemock.NewMockBatchStore(ctrl)
 	mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).AnyTimes()
 	mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
-	mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
+	mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{
+		ID:           "test-queue/batch/1",
+		Queue:        request.Queue,
+		Contains:     []string{request.ID},
+		Dependencies: []string{},
+		State:        entity.BatchStateCreated,
+		Version:      1,
+	}, int32(1), int32(2)).Return(nil)
 
 	mockReqStore := storagemock.NewMockRequestStore(ctrl)
 	mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil)
@@ -349,7 +356,14 @@ func TestController_Process_WithDependencies(t *testing.T) {
 	mockBatchStore := storagemock.NewMockBatchStore(ctrl)
 	mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(activeBatches, nil)
 	mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
-	mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
+	mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{
+		ID:           "test-queue/batch/1",
+		Queue:        request.Queue,
+		Contains:     []string{request.ID},
+		Dependencies: []string{"test-queue/batch/1", "test-queue/batch/2"},
+		State:        entity.BatchStateCreated,
+		Version:      1,
+	}, int32(1), int32(2)).Return(nil)
 
 	mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
 	// batch/1 has no existing dependents.
@@ -418,7 +432,14 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) {
 	mockBatchStore := storagemock.NewMockBatchStore(ctrl)
 	mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(activeBatches, nil)
 	mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
-	mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
+	mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{
+		ID:           "test-queue/batch/1",
+		Queue:        request.Queue,
+		Contains:     []string{request.ID},
+		Dependencies: []string{"test-queue/batch/2"},
+		State:        entity.BatchStateCreated,
+		Version:      1,
+	}, int32(1), int32(2)).Return(nil)
 
 	mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
 	// Only batch/2 is selected by the analyzer, so only it gets a reverse-index update.
@@ -712,7 +733,14 @@ func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) {
 	mockBatchStore := storagemock.NewMockBatchStore(ctrl)
 	mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil)
 	mockBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
-	mockBatchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
+	mockBatchStore.EXPECT().Update(gomock.Any(), entity.Batch{
+		ID:           "test-queue/batch/1",
+		Queue:        request.Queue,
+		Contains:     []string{request.ID},
+		Dependencies: []string{},
+		State:        entity.BatchStateCreated,
+		Version:      1,
+	}, int32(1), int32(2)).Return(nil)
 
 	mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
 	mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil)
@@ -784,7 +812,7 @@ func TestController_Process_ReadiesBatchBeforePublishing(t *testing.T) {
 			Dependents: []string{},
 			Version:    1,
 		}).Return(nil),
-		batchStore.EXPECT().UpdateState(gomock.Any(), batch.ID, int32(1), int32(2), entity.BatchStateCreated).Return(nil),
+		batchStore.EXPECT().Update(gomock.Any(), batch.WithState(entity.BatchStateCreated), int32(1), int32(2)).Return(nil),
 		publisher.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).Return(nil),
 		publisher.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).Return(nil),
 	)
@@ -846,8 +874,22 @@ func TestController_Process_RedeliveryMintsFreshBatchID(t *testing.T) {
 			return nil
 		},
 	).Times(2)
-	batchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/1", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
-	batchStore.EXPECT().UpdateState(gomock.Any(), "test-queue/batch/2", int32(1), int32(2), entity.BatchStateCreated).Return(nil)
+	batchStore.EXPECT().Update(gomock.Any(), entity.Batch{
+		ID:           "test-queue/batch/1",
+		Queue:        firstRequest.Queue,
+		Contains:     []string{firstRequest.ID},
+		Dependencies: []string{},
+		State:        entity.BatchStateCreated,
+		Version:      1,
+	}, int32(1), int32(2)).Return(nil)
+	batchStore.EXPECT().Update(gomock.Any(), entity.Batch{
+		ID:           "test-queue/batch/2",
+		Queue:        firstRequest.Queue,
+		Contains:     []string{firstRequest.ID},
+		Dependencies: []string{},
+		State:        entity.BatchStateCreated,
+		Version:      1,
+	}, int32(1), int32(2)).Return(nil)
 
 	batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl)
 	batchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).Times(2)
@@ -978,7 +1020,7 @@ func TestController_PopulateBatch_Errors(t *testing.T) {
 					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)
+				batchStore.EXPECT().Update(gomock.Any(), batch.WithState(entity.BatchStateCreated), int32(1), int32(2)).Return(storeErr)
 			},
 			errMsg: "failed to mark batch",
 		},