From 7f18ea0bb61b891a00035bcc734e61f365348187 Mon Sep 17 00:00:00 2001 From: Albert Wu Date: Mon, 20 Jul 2026 11:36:16 -0700 Subject: [PATCH] fix(orchestrator): reconcile batch redelivery --- submitqueue/entity/BUILD.bazel | 1 + submitqueue/entity/request_batch.go | 27 ++ submitqueue/extension/storage/BUILD.bazel | 1 + .../extension/storage/mock/BUILD.bazel | 1 + .../storage/mock/request_batch_store_mock.go | 71 +++ .../extension/storage/mock/storage_mock.go | 42 +- .../extension/storage/mysql/BUILD.bazel | 2 + .../storage/mysql/request_batch_store.go | 76 +++ .../storage/mysql/request_batch_store_test.go | 164 +++++++ .../extension/storage/mysql/storage.go | 7 + .../extension/storage/mysql/storage_test.go | 1 + .../extension/storage/request_batch_store.go | 34 ++ submitqueue/extension/storage/storage.go | 3 + .../orchestrator/controller/batch/batch.go | 439 +++++++++++------- .../controller/batch/batch_test.go | 351 +++++++++++--- 15 files changed, 983 insertions(+), 237 deletions(-) create mode 100644 submitqueue/entity/request_batch.go create mode 100644 submitqueue/extension/storage/mock/request_batch_store_mock.go create mode 100644 submitqueue/extension/storage/mysql/request_batch_store.go create mode 100644 submitqueue/extension/storage/mysql/request_batch_store_test.go create mode 100644 submitqueue/extension/storage/request_batch_store.go diff --git a/submitqueue/entity/BUILD.bazel b/submitqueue/entity/BUILD.bazel index b39a26ef..e87203e9 100644 --- a/submitqueue/entity/BUILD.bazel +++ b/submitqueue/entity/BUILD.bazel @@ -18,6 +18,7 @@ go_library( "queue.go", "queue_config.go", "request.go", + "request_batch.go", "request_history.go", "request_log.go", "request_summary.go", diff --git a/submitqueue/entity/request_batch.go b/submitqueue/entity/request_batch.go new file mode 100644 index 00000000..4b6bf6a8 --- /dev/null +++ b/submitqueue/entity/request_batch.go @@ -0,0 +1,27 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package entity + +// RequestBatch is the immutable assignment of one request to one batch. +// RequestID is unique, so retries resolve the same BatchID instead of creating +// another logical batch for the request. +type RequestBatch struct { + // RequestID is the globally unique request identifier. + RequestID string + // BatchID is the globally unique batch identifier assigned to the request. + BatchID string + // Version is the version of the assignment. Immutable assignments start at 1. + Version int32 +} diff --git a/submitqueue/extension/storage/BUILD.bazel b/submitqueue/extension/storage/BUILD.bazel index d2081a00..2daa36fa 100644 --- a/submitqueue/extension/storage/BUILD.bazel +++ b/submitqueue/extension/storage/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "batch_store.go", "build_store.go", "change_store.go", + "request_batch_store.go", "request_log_store.go", "request_queue_summary_store.go", "request_store.go", diff --git a/submitqueue/extension/storage/mock/BUILD.bazel b/submitqueue/extension/storage/mock/BUILD.bazel index 4b2f4f51..11b10ad1 100644 --- a/submitqueue/extension/storage/mock/BUILD.bazel +++ b/submitqueue/extension/storage/mock/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "batch_store_mock.go", "build_store_mock.go", "change_store_mock.go", + "request_batch_store_mock.go", "request_log_store_mock.go", "request_queue_summary_store_mock.go", "request_store_mock.go", diff --git a/submitqueue/extension/storage/mock/request_batch_store_mock.go b/submitqueue/extension/storage/mock/request_batch_store_mock.go new file mode 100644 index 00000000..5b6cd71f --- /dev/null +++ b/submitqueue/extension/storage/mock/request_batch_store_mock.go @@ -0,0 +1,71 @@ +// Code generated by MockGen. DO NOT EDIT. +// Source: request_batch_store.go +// +// Generated by this command: +// +// mockgen -source=request_batch_store.go -destination=mock/request_batch_store_mock.go -package=mock +// + +// Package mock is a generated GoMock package. +package mock + +import ( + context "context" + reflect "reflect" + + entity "github.com/uber/submitqueue/submitqueue/entity" + gomock "go.uber.org/mock/gomock" +) + +// MockRequestBatchStore is a mock of RequestBatchStore interface. +type MockRequestBatchStore struct { + ctrl *gomock.Controller + recorder *MockRequestBatchStoreMockRecorder + isgomock struct{} +} + +// MockRequestBatchStoreMockRecorder is the mock recorder for MockRequestBatchStore. +type MockRequestBatchStoreMockRecorder struct { + mock *MockRequestBatchStore +} + +// NewMockRequestBatchStore creates a new mock instance. +func NewMockRequestBatchStore(ctrl *gomock.Controller) *MockRequestBatchStore { + mock := &MockRequestBatchStore{ctrl: ctrl} + mock.recorder = &MockRequestBatchStoreMockRecorder{mock} + return mock +} + +// EXPECT returns an object that allows the caller to indicate expected use. +func (m *MockRequestBatchStore) EXPECT() *MockRequestBatchStoreMockRecorder { + return m.recorder +} + +// Create mocks base method. +func (m *MockRequestBatchStore) Create(ctx context.Context, requestBatch entity.RequestBatch) error { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Create", ctx, requestBatch) + ret0, _ := ret[0].(error) + return ret0 +} + +// Create indicates an expected call of Create. +func (mr *MockRequestBatchStoreMockRecorder) Create(ctx, requestBatch any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Create", reflect.TypeOf((*MockRequestBatchStore)(nil).Create), ctx, requestBatch) +} + +// Get mocks base method. +func (m *MockRequestBatchStore) Get(ctx context.Context, requestID string) (entity.RequestBatch, error) { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "Get", ctx, requestID) + ret0, _ := ret[0].(entity.RequestBatch) + ret1, _ := ret[1].(error) + return ret0, ret1 +} + +// Get indicates an expected call of Get. +func (mr *MockRequestBatchStoreMockRecorder) Get(ctx, requestID any) *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockRequestBatchStore)(nil).Get), ctx, requestID) +} diff --git a/submitqueue/extension/storage/mock/storage_mock.go b/submitqueue/extension/storage/mock/storage_mock.go index 6655716f..6267c7c6 100644 --- a/submitqueue/extension/storage/mock/storage_mock.go +++ b/submitqueue/extension/storage/mock/storage_mock.go @@ -110,6 +110,20 @@ func (mr *MockStorageMockRecorder) GetChangeStore() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetChangeStore", reflect.TypeOf((*MockStorage)(nil).GetChangeStore)) } +// GetRequestBatchStore mocks base method. +func (m *MockStorage) GetRequestBatchStore() storage.RequestBatchStore { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetRequestBatchStore") + ret0, _ := ret[0].(storage.RequestBatchStore) + return ret0 +} + +// GetRequestBatchStore indicates an expected call of GetRequestBatchStore. +func (mr *MockStorageMockRecorder) GetRequestBatchStore() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRequestBatchStore", reflect.TypeOf((*MockStorage)(nil).GetRequestBatchStore)) +} + // GetRequestLogStore mocks base method. func (m *MockStorage) GetRequestLogStore() storage.RequestLogStore { m.ctrl.T.Helper() @@ -152,20 +166,6 @@ func (mr *MockStorageMockRecorder) GetRequestStore() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRequestStore", reflect.TypeOf((*MockStorage)(nil).GetRequestStore)) } -// GetSpeculationPathBuildStore mocks base method. -func (m *MockStorage) GetSpeculationPathBuildStore() storage.SpeculationPathBuildStore { - m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "GetSpeculationPathBuildStore") - ret0, _ := ret[0].(storage.SpeculationPathBuildStore) - return ret0 -} - -// GetSpeculationPathBuildStore indicates an expected call of GetSpeculationPathBuildStore. -func (mr *MockStorageMockRecorder) GetSpeculationPathBuildStore() *gomock.Call { - mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSpeculationPathBuildStore", reflect.TypeOf((*MockStorage)(nil).GetSpeculationPathBuildStore)) -} - // GetRequestSummaryStore mocks base method. func (m *MockStorage) GetRequestSummaryStore() storage.RequestSummaryStore { m.ctrl.T.Helper() @@ -194,6 +194,20 @@ func (mr *MockStorageMockRecorder) GetRequestURIStore() *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetRequestURIStore", reflect.TypeOf((*MockStorage)(nil).GetRequestURIStore)) } +// GetSpeculationPathBuildStore mocks base method. +func (m *MockStorage) GetSpeculationPathBuildStore() storage.SpeculationPathBuildStore { + m.ctrl.T.Helper() + ret := m.ctrl.Call(m, "GetSpeculationPathBuildStore") + ret0, _ := ret[0].(storage.SpeculationPathBuildStore) + return ret0 +} + +// GetSpeculationPathBuildStore indicates an expected call of GetSpeculationPathBuildStore. +func (mr *MockStorageMockRecorder) GetSpeculationPathBuildStore() *gomock.Call { + mr.mock.ctrl.T.Helper() + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GetSpeculationPathBuildStore", reflect.TypeOf((*MockStorage)(nil).GetSpeculationPathBuildStore)) +} + // GetSpeculationTreeStore mocks base method. func (m *MockStorage) GetSpeculationTreeStore() storage.SpeculationTreeStore { m.ctrl.T.Helper() diff --git a/submitqueue/extension/storage/mysql/BUILD.bazel b/submitqueue/extension/storage/mysql/BUILD.bazel index 1cbbf031..6bd72749 100644 --- a/submitqueue/extension/storage/mysql/BUILD.bazel +++ b/submitqueue/extension/storage/mysql/BUILD.bazel @@ -7,6 +7,7 @@ go_library( "batch_store.go", "build_store.go", "change_store.go", + "request_batch_store.go", "request_log_store.go", "request_queue_summary_store.go", "request_store.go", @@ -34,6 +35,7 @@ go_test( "batch_store_test.go", "build_store_test.go", "change_store_test.go", + "request_batch_store_test.go", "request_log_store_test.go", "request_queue_summary_store_test.go", "request_store_test.go", diff --git a/submitqueue/extension/storage/mysql/request_batch_store.go b/submitqueue/extension/storage/mysql/request_batch_store.go new file mode 100644 index 00000000..a5833212 --- /dev/null +++ b/submitqueue/extension/storage/mysql/request_batch_store.go @@ -0,0 +1,76 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mysql + +import ( + "context" + "database/sql" + "errors" + "fmt" + + "github.com/go-sql-driver/mysql" + "github.com/uber-go/tally" + + "github.com/uber/submitqueue/platform/metrics" + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +type requestBatchStore struct { + db *sql.DB + scope tally.Scope +} + +// NewRequestBatchStore creates a new MySQL-backed RequestBatchStore. +func NewRequestBatchStore(db *sql.DB, scope tally.Scope) storage.RequestBatchStore { + return &requestBatchStore{db: db, scope: scope} +} + +// Get retrieves the immutable batch assignment for a request. +func (s *requestBatchStore) Get(ctx context.Context, requestID string) (ret entity.RequestBatch, retErr error) { + op := metrics.Begin(s.scope, "get") + defer func() { op.Complete(retErr) }() + + err := s.db.QueryRowContext(ctx, + "SELECT request_id, batch_id, version FROM request_batch WHERE request_id = ?", + requestID, + ).Scan(&ret.RequestID, &ret.BatchID, &ret.Version) + if errors.Is(err, sql.ErrNoRows) { + return entity.RequestBatch{}, storage.WrapNotFound(err) + } + if err != nil { + return entity.RequestBatch{}, fmt.Errorf("failed to get request batch assignment requestID=%s: %w", requestID, err) + } + return ret, nil +} + +// Create reserves an immutable batch assignment for a request. +func (s *requestBatchStore) Create(ctx context.Context, requestBatch entity.RequestBatch) (retErr error) { + op := metrics.Begin(s.scope, "create") + defer func() { op.Complete(retErr) }() + + _, err := s.db.ExecContext(ctx, + "INSERT INTO request_batch (request_id, batch_id, version) VALUES (?, ?, ?)", + requestBatch.RequestID, requestBatch.BatchID, requestBatch.Version, + ) + if err != nil { + var mysqlErr *mysql.MySQLError + if errors.As(err, &mysqlErr) && mysqlErr.Number == mysqlErrDuplicateEntry { + return fmt.Errorf("request batch assignment requestID=%s: %w", requestBatch.RequestID, storage.ErrAlreadyExists) + } + return fmt.Errorf("failed to create request batch assignment requestID=%s batchID=%s: %w", requestBatch.RequestID, requestBatch.BatchID, err) + } + return nil +} diff --git a/submitqueue/extension/storage/mysql/request_batch_store_test.go b/submitqueue/extension/storage/mysql/request_batch_store_test.go new file mode 100644 index 00000000..d50a4379 --- /dev/null +++ b/submitqueue/extension/storage/mysql/request_batch_store_test.go @@ -0,0 +1,164 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package mysql + +import ( + "context" + "database/sql" + "fmt" + "testing" + + "github.com/DATA-DOG/go-sqlmock" + "github.com/go-sql-driver/mysql" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "github.com/uber/submitqueue/submitqueue/entity" + "github.com/uber/submitqueue/submitqueue/extension/storage" +) + +func setupRequestBatchStoreTest(t *testing.T) (*sql.DB, sqlmock.Sqlmock, storage.RequestBatchStore) { + t.Helper() + + db, mock, err := sqlmock.New() + require.NoError(t, err) + return db, mock, NewRequestBatchStore(db, testMetrics()) +} + +func TestRequestBatchStore_Get(t *testing.T) { + want := entity.RequestBatch{ + RequestID: "monorepo/1", + BatchID: "monorepo/batch/7", + Version: 1, + } + + tests := []struct { + name string + setup func(sqlmock.Sqlmock) + wantErr bool + wantErrIs error + }{ + { + name: "found", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT request_id, batch_id, version FROM request_batch"). + WithArgs(want.RequestID). + WillReturnRows(sqlmock.NewRows([]string{"request_id", "batch_id", "version"}). + AddRow(want.RequestID, want.BatchID, want.Version)) + }, + }, + { + name: "not found", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT request_id, batch_id, version FROM request_batch"). + WithArgs(want.RequestID). + WillReturnError(sql.ErrNoRows) + }, + wantErr: true, + wantErrIs: storage.ErrNotFound, + }, + { + name: "query error", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectQuery("SELECT request_id, batch_id, version FROM request_batch"). + WithArgs(want.RequestID). + WillReturnError(fmt.Errorf("connection reset")) + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, store := setupRequestBatchStoreTest(t) + defer db.Close() + tt.setup(mock) + + got, err := store.Get(context.Background(), want.RequestID) + if tt.wantErr { + require.Error(t, err) + if tt.wantErrIs != nil { + assert.ErrorIs(t, err, tt.wantErrIs) + } + } else { + require.NoError(t, err) + assert.Equal(t, want, got) + } + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} + +func TestRequestBatchStore_Create(t *testing.T) { + requestBatch := entity.RequestBatch{ + RequestID: "monorepo/1", + BatchID: "monorepo/batch/7", + Version: 1, + } + + tests := []struct { + name string + setup func(sqlmock.Sqlmock) + wantErr bool + wantErrIs error + }{ + { + name: "success", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT INTO request_batch"). + WithArgs(requestBatch.RequestID, requestBatch.BatchID, requestBatch.Version). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + { + name: "duplicate request returns ErrAlreadyExists", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT INTO request_batch"). + WithArgs(requestBatch.RequestID, requestBatch.BatchID, requestBatch.Version). + WillReturnError(&mysql.MySQLError{Number: mysqlErrDuplicateEntry}) + }, + wantErr: true, + wantErrIs: storage.ErrAlreadyExists, + }, + { + name: "exec error", + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("INSERT INTO request_batch"). + WithArgs(requestBatch.RequestID, requestBatch.BatchID, requestBatch.Version). + WillReturnError(fmt.Errorf("connection reset")) + }, + wantErr: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + db, mock, store := setupRequestBatchStoreTest(t) + defer db.Close() + tt.setup(mock) + + err := store.Create(context.Background(), requestBatch) + if tt.wantErr { + require.Error(t, err) + if tt.wantErrIs != nil { + assert.ErrorIs(t, err, tt.wantErrIs) + } + } else { + require.NoError(t, err) + } + require.NoError(t, mock.ExpectationsWereMet()) + }) + } +} diff --git a/submitqueue/extension/storage/mysql/storage.go b/submitqueue/extension/storage/mysql/storage.go index ffcf4d02..3990a1ee 100644 --- a/submitqueue/extension/storage/mysql/storage.go +++ b/submitqueue/extension/storage/mysql/storage.go @@ -32,6 +32,7 @@ type mysqlStorage struct { requestStore storage.RequestStore changeStore storage.ChangeStore batchStore storage.BatchStore + requestBatchStore storage.RequestBatchStore batchDependentStore storage.BatchDependentStore buildStore storage.BuildStore speculationPathBuildStore storage.SpeculationPathBuildStore @@ -49,6 +50,7 @@ func NewStorage(db *sql.DB, scope tally.Scope) (storage.Storage, error) { requestStore: NewRequestStore(db, scope.SubScope("request_store")), changeStore: NewChangeStore(db, scope.SubScope("change_store")), batchStore: NewBatchStore(db, scope.SubScope("batch_store")), + requestBatchStore: NewRequestBatchStore(db, scope.SubScope("request_batch_store")), batchDependentStore: NewBatchDependentStore(db, scope.SubScope("batch_dependent_store")), buildStore: NewBuildStore(db, scope.SubScope("build_store")), speculationPathBuildStore: NewSpeculationPathBuildStore(db, scope.SubScope("speculation_path_build_store")), @@ -75,6 +77,11 @@ func (f *mysqlStorage) GetBatchStore() storage.BatchStore { return f.batchStore } +// GetRequestBatchStore returns the MySQL-backed RequestBatchStore. +func (f *mysqlStorage) GetRequestBatchStore() storage.RequestBatchStore { + return f.requestBatchStore +} + // GetBatchDependentStore returns the MySQL-backed BatchDependentStore. func (f *mysqlStorage) GetBatchDependentStore() storage.BatchDependentStore { return f.batchDependentStore diff --git a/submitqueue/extension/storage/mysql/storage_test.go b/submitqueue/extension/storage/mysql/storage_test.go index bf0cb68b..05bc6491 100644 --- a/submitqueue/extension/storage/mysql/storage_test.go +++ b/submitqueue/extension/storage/mysql/storage_test.go @@ -39,6 +39,7 @@ func TestNewStorage(t *testing.T) { assert.NotNil(t, s.GetRequestStore()) assert.NotNil(t, s.GetChangeStore()) assert.NotNil(t, s.GetBatchStore()) + assert.NotNil(t, s.GetRequestBatchStore()) assert.NotNil(t, s.GetBatchDependentStore()) assert.NotNil(t, s.GetBuildStore()) assert.NotNil(t, s.GetSpeculationPathBuildStore()) diff --git a/submitqueue/extension/storage/request_batch_store.go b/submitqueue/extension/storage/request_batch_store.go new file mode 100644 index 00000000..618b63cd --- /dev/null +++ b/submitqueue/extension/storage/request_batch_store.go @@ -0,0 +1,34 @@ +// Copyright (c) 2025 Uber Technologies, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package storage + +//go:generate mockgen -source=request_batch_store.go -destination=mock/request_batch_store_mock.go -package=mock + +import ( + "context" + + "github.com/uber/submitqueue/submitqueue/entity" +) + +// RequestBatchStore manages immutable request-to-batch assignments. +type RequestBatchStore interface { + // Get retrieves the assignment keyed by request ID. + // Returns ErrNotFound if the request has not been assigned. + Get(ctx context.Context, requestID string) (entity.RequestBatch, error) + + // Create reserves the batch assignment for a request. + // Returns ErrAlreadyExists if the request already has an assignment. + Create(ctx context.Context, requestBatch entity.RequestBatch) error +} diff --git a/submitqueue/extension/storage/storage.go b/submitqueue/extension/storage/storage.go index b54736c3..c7a4eab5 100644 --- a/submitqueue/extension/storage/storage.go +++ b/submitqueue/extension/storage/storage.go @@ -53,6 +53,9 @@ type Storage interface { // GetBatchStore returns the BatchStore instance. GetBatchStore() BatchStore + // GetRequestBatchStore returns the RequestBatchStore instance. + GetRequestBatchStore() RequestBatchStore + // GetBatchDependentStore returns the BatchDependentStore instance. GetBatchDependentStore() BatchDependentStore diff --git a/submitqueue/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index 9bb13e85..1fed3382 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/submitqueue/orchestrator/controller/batch/batch.go @@ -90,7 +90,6 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r return fmt.Errorf("failed to deserialize request ID: %w", err) } - // Fetch request from storage request, err := c.store.GetRequestStore().Get(ctx, rid.ID) if err != nil { metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) @@ -106,216 +105,332 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r "partition_key", msg.PartitionKey, ) - // Short-circuit if the request has been halted — either it already reached a - // terminal state, or the cancel controller has recorded a cancellation intent - // (RequestStateCancelling). A halted request must never spawn a new batch. - if entity.IsRequestStateHalted(request.State) { - c.metricsScope.Counter("skipped_halted").Inc(1) - c.logger.Infow("skipping batch for halted request", - "request_id", request.ID, - "state", string(request.State), - ) + progress, err := classifyRequest(request.State) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "unexpected_state_errors", 1) + return fmt.Errorf("cannot batch request %s: %w", request.ID, err) + } + switch progress { + case requestDownstreamProgressed: + metrics.NamedCounter(c.metricsScope, opName, "downstream_progressed", 1) + return nil + case requestSuperseded: + metrics.NamedCounter(c.metricsScope, opName, "skipped_halted", 1) return nil } - // TODO: if capacity is full, wait here for other requests to accumulate to batch them together, or include a request into an existing batch if it's not too late. + requestBatch, assignmentExisted, err := c.resolveRequestBatch(ctx, request) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "request_batch_store_errors", 1) + return err + } + + if progress == requestNeedsTransition { + newVersion := request.Version + 1 + if err := c.store.GetRequestStore().UpdateState(ctx, request.ID, request.Version, newVersion, entity.RequestStateBatched); err != nil { + if errors.Is(err, storage.ErrVersionMismatch) { + metrics.NamedCounter(c.metricsScope, opName, "request_claim_conflicts", 1) + } + metrics.NamedCounter(c.metricsScope, opName, "request_claim_errors", 1) + return fmt.Errorf("failed to claim request %s for batch %s: %w", request.ID, requestBatch.BatchID, err) + } + request.Version = newVersion + request.State = entity.RequestStateBatched + } + + batch, err := c.reconcileBatch(ctx, request, requestBatch, assignmentExisted) + if err != nil { + return err + } + + switch batch.State { + case entity.BatchStateCreated: + // This controller owns the Created batch fanout. Replay all of it + // because an earlier delivery may have stopped after any write or + // publish. + case entity.BatchStateScored, entity.BatchStateSpeculating, entity.BatchStateMerging: + metrics.NamedCounter(c.metricsScope, opName, "batch_downstream_progressed", 1) + return nil + case entity.BatchStateCancelling, entity.BatchStateSucceeded, entity.BatchStateFailed, entity.BatchStateCancelled: + metrics.NamedCounter(c.metricsScope, opName, "batch_superseded", 1) + return nil + default: + metrics.NamedCounter(c.metricsScope, opName, "unexpected_batch_state_errors", 1) + return fmt.Errorf("batch %s has invalid state %q for batch reconciliation", batch.ID, batch.State) + } + + if err := c.reconcileBatchDependents(ctx, batch); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) + return err + } + + logEntry := entity.NewRequestLog(request.ID, entity.RequestStatusBatched, request.Version, "", map[string]string{ + "batch_id": batch.ID, + }) + if err := corerequest.PublishLog(ctx, c.registry, logEntry, request.ID); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "request_log_errors", 1) + return fmt.Errorf("failed to publish request log for request %s: %w", request.ID, err) + } + + if err := c.publish(ctx, topickey.TopicKeyScore, batch.ID, batch.Queue); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) + return fmt.Errorf("failed to publish batch ID to score topic: %w", err) + } + + c.logger.Infow("published batch to score topic", + "batch_id", batch.ID, + "topic_key", topickey.TopicKeyScore, + ) + return nil +} + +type requestProgress int + +const ( + requestNeedsTransition requestProgress = iota + requestTransitionApplied + requestDownstreamProgressed + requestSuperseded +) + +func classifyRequest(state entity.RequestState) (requestProgress, error) { + switch state { + case entity.RequestStateStarted, entity.RequestStateValidated: + return requestNeedsTransition, nil + case entity.RequestStateBatched: + return requestTransitionApplied, nil + case entity.RequestStateProcessing: + return requestDownstreamProgressed, nil + case entity.RequestStateCancelling, entity.RequestStateLanded, entity.RequestStateError, entity.RequestStateCancelled: + return requestSuperseded, nil + default: + return requestNeedsTransition, fmt.Errorf("unexpected request state %q", state) + } +} + +// resolveRequestBatch returns the durable assignment for this logical batch +// operation. A new assignment is written before the request CAS so every retry +// can reconstruct the same batch ID. +func (c *Controller) resolveRequestBatch(ctx context.Context, request entity.Request) (entity.RequestBatch, bool, error) { + store := c.store.GetRequestBatchStore() + requestBatch, err := store.Get(ctx, request.ID) + if err == nil { + if err := validateRequestBatch(request, requestBatch); err != nil { + return entity.RequestBatch{}, false, err + } + return requestBatch, true, nil + } + if !errors.Is(err, storage.ErrNotFound) { + return entity.RequestBatch{}, false, fmt.Errorf("failed to get batch assignment for request %s: %w", request.ID, err) + } + + // request_batch predates this controller contract. Recover an active legacy + // batch once rather than assigning a second batch to an already-batched + // request. + if request.State == entity.RequestStateBatched { + batch, found, err := c.findLegacyBatch(ctx, request) + if err != nil { + return entity.RequestBatch{}, false, err + } + if !found { + return entity.RequestBatch{}, false, fmt.Errorf("batched request %s has no durable batch assignment or active batch", request.ID) + } + requestBatch = entity.RequestBatch{RequestID: request.ID, BatchID: batch.ID, Version: 1} + if err := store.Create(ctx, requestBatch); err != nil { + if !errors.Is(err, storage.ErrAlreadyExists) { + return entity.RequestBatch{}, false, fmt.Errorf("failed to recover batch assignment for request %s: %w", request.ID, err) + } + requestBatch, err = store.Get(ctx, request.ID) + if err != nil { + return entity.RequestBatch{}, false, fmt.Errorf("failed to reload concurrent batch assignment for request %s: %w", request.ID, err) + } + if err := validateRequestBatch(request, requestBatch); err != nil { + return entity.RequestBatch{}, false, err + } + } + return requestBatch, true, nil + } - // Generate a globally unique batch ID. seq, err := c.counter.Next(ctx, "batch/"+request.Queue) if err != nil { - metrics.NamedCounter(c.metricsScope, opName, "counter_errors", 1) - return fmt.Errorf("failed to generate batch ID for queue=%s: %w", request.Queue, err) + return entity.RequestBatch{}, false, fmt.Errorf("failed to generate batch ID for queue=%s: %w", request.Queue, err) + } + requestBatch = entity.RequestBatch{ + RequestID: request.ID, + BatchID: fmt.Sprintf("%s/batch/%d", request.Queue, seq), + Version: 1, + } + if err := store.Create(ctx, requestBatch); err == nil { + return requestBatch, false, nil + } else if !errors.Is(err, storage.ErrAlreadyExists) { + return entity.RequestBatch{}, false, fmt.Errorf("failed to reserve batch assignment for request %s: %w", request.ID, err) + } + + requestBatch, err = store.Get(ctx, request.ID) + if err != nil { + return entity.RequestBatch{}, false, fmt.Errorf("failed to reload concurrent batch assignment for request %s: %w", request.ID, err) + } + if err := validateRequestBatch(request, requestBatch); err != nil { + return entity.RequestBatch{}, false, err + } + return requestBatch, true, nil +} + +func validateRequestBatch(request entity.Request, requestBatch entity.RequestBatch) error { + if requestBatch.RequestID != request.ID { + return fmt.Errorf("batch assignment key mismatch: requested %s, got %s", request.ID, requestBatch.RequestID) + } + if requestBatch.BatchID == "" { + return fmt.Errorf("batch assignment for request %s has an empty batch ID", request.ID) + } + return nil +} + +func (c *Controller) findLegacyBatch(ctx context.Context, request entity.Request) (entity.Batch, bool, error) { + batches, err := c.store.GetBatchStore().GetByQueueAndStates(ctx, request.Queue, entity.ActiveBatchStates()) + if err != nil { + return entity.Batch{}, false, fmt.Errorf("failed to find legacy batch for request %s: %w", request.ID, err) + } + + var found entity.Batch + for _, batch := range batches { + if !contains(batch.Contains, request.ID) { + continue + } + if found.ID != "" { + return entity.Batch{}, false, fmt.Errorf("request %s belongs to multiple active batches: %s and %s", request.ID, found.ID, batch.ID) + } + found = batch + } + return found, found.ID != "", nil +} + +func (c *Controller) reconcileBatch(ctx context.Context, request entity.Request, requestBatch entity.RequestBatch, assignmentExisted bool) (entity.Batch, error) { + if assignmentExisted { + batch, err := c.store.GetBatchStore().Get(ctx, requestBatch.BatchID) + if err == nil { + if err := validateBatch(request, requestBatch, batch); err != nil { + return entity.Batch{}, err + } + return batch, nil + } + if !errors.Is(err, storage.ErrNotFound) { + return entity.Batch{}, fmt.Errorf("failed to get assigned batch %s: %w", requestBatch.BatchID, err) + } } batch := entity.Batch{ - ID: fmt.Sprintf("%s/batch/%d", request.Queue, seq), + ID: requestBatch.BatchID, Queue: request.Queue, Contains: []string{request.ID}, State: entity.BatchStateCreated, Version: 1, } - // Get active batches for this queue and ask the conflict analyzer which - // of them the new batch must serialize behind. The dependency set drives - // the speculation graph downstream. activeBatches, err := c.store.GetBatchStore().GetByQueueAndStates(ctx, request.Queue, entity.DependencyBatchStates()) if err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1) - return fmt.Errorf("failed to get active batches for queue=%s: %w", request.Queue, err) + return entity.Batch{}, fmt.Errorf("failed to get active batches for queue=%s: %w", request.Queue, err) } - - // Dedupe by batch ID since a single (analyzed, in-flight) pair may be - // reported with multiple Conflict entries when different conflict types - // apply; the dependency graph only tracks the relation. analyzer, err := c.analyzers.For(conflict.Config{QueueName: batch.Queue}) if err != nil { - metrics.NamedCounter(c.metricsScope, opName, "conflict_analyzer_errors", 1) - return fmt.Errorf("failed to build conflict analyzer for queue=%s: %w", batch.Queue, err) + return entity.Batch{}, fmt.Errorf("failed to build conflict analyzer for queue=%s: %w", batch.Queue, err) } conflicts, err := analyzer.Analyze(ctx, batch, activeBatches) if err != nil { - metrics.NamedCounter(c.metricsScope, opName, "conflict_analyzer_errors", 1) - return fmt.Errorf("failed to analyze conflicts for batchID=%s: %w", batch.ID, err) + return entity.Batch{}, fmt.Errorf("failed to analyze conflicts for batchID=%s: %w", batch.ID, err) } seen := make(map[string]struct{}, len(conflicts)) - conflictingIDs := make([]string, 0, len(conflicts)) - for _, cf := range conflicts { - if _, ok := seen[cf.BatchID]; ok { + for _, conflict := range conflicts { + if _, ok := seen[conflict.BatchID]; ok { continue } - seen[cf.BatchID] = struct{}{} - conflictingIDs = append(conflictingIDs, cf.BatchID) + seen[conflict.BatchID] = struct{}{} + batch.Dependencies = append(batch.Dependencies, conflict.BatchID) } - batch.Dependencies = conflictingIDs - - // Update reverse index for each conflicting batch (BatchDependent = - // "batches that depend on me"). One UpdateDependents call per conflict. - for _, depID := range conflictingIDs { - existing, err := c.store.GetBatchDependentStore().Get(ctx, depID) + if err := c.store.GetBatchStore().Create(ctx, batch); err != nil { + if !errors.Is(err, storage.ErrAlreadyExists) { + return entity.Batch{}, fmt.Errorf("failed to create batch %s: %w", batch.ID, err) + } + batch, err = c.store.GetBatchStore().Get(ctx, batch.ID) if err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) - return fmt.Errorf("failed to get batch dependent for batchID=%s: %w", depID, err) + return entity.Batch{}, fmt.Errorf("failed to reload concurrent batch %s: %w", requestBatch.BatchID, err) } + if err := validateBatch(request, requestBatch, batch); err != nil { + return entity.Batch{}, err + } + } - dependents := append(existing.Dependents, batch.ID) + c.logger.Infow("batch reconciled", + "batch_id", batch.ID, + "request_id", request.ID, + "queue", request.Queue, + "dependency_count", len(batch.Dependencies), + ) + return batch, nil +} - newVersion := existing.Version + 1 - if err := c.store.GetBatchDependentStore().UpdateDependents(ctx, depID, existing.Version, newVersion, dependents); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) - return fmt.Errorf("failed to update batch dependent index for existing batchID=%s and new batchID=%s: %w", depID, batch.ID, err) - } +func validateBatch(request entity.Request, requestBatch entity.RequestBatch, batch entity.Batch) error { + if batch.ID != requestBatch.BatchID { + return fmt.Errorf("assigned batch ID mismatch for request %s: expected %s, got %s", request.ID, requestBatch.BatchID, batch.ID) + } + if batch.Queue != request.Queue { + return fmt.Errorf("assigned batch %s has queue %s, expected %s", batch.ID, batch.Queue, request.Queue) } + if !contains(batch.Contains, request.ID) { + return fmt.Errorf("assigned batch %s does not contain request %s", batch.ID, request.ID) + } + return nil +} - // Create new reverse index entry for the new batch. It would be empty for now, but will be updated as new batches are created that conflict with this batch. - bd := entity.BatchDependent{ +func (c *Controller) reconcileBatchDependents(ctx context.Context, batch entity.Batch) error { + store := c.store.GetBatchDependentStore() + own := entity.BatchDependent{ BatchID: batch.ID, Dependents: []string{}, Version: 1, } + if err := store.Create(ctx, own); err != nil && !errors.Is(err, storage.ErrAlreadyExists) { + return fmt.Errorf("failed to create dependent index for batch %s: %w", batch.ID, err) + } - if err := c.store.GetBatchDependentStore().Create(ctx, bd); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_dependent_store_errors", 1) - return fmt.Errorf("failed to create batch dependent index for new batchID=%s: %w", batch.ID, err) - } - - // Claim the request for this batch with a CAS-write that transitions the - // request to RequestStateBatched. This CAS is the serialization point - // between the batch controller and the cancel controller — without it, the - // two would race over an empty interleaving and produce an orphan batch - // containing a cancelled request. - // - // Concrete race that this CAS closes (T1..T7 are wall-clock orderings of - // independent batch- and cancel-controller goroutines): - // - // T1 batch.Get(R) → R{State: Validated, Version: 1} - // T2 cancel.Get(R) → R{State: Validated, Version: 1} - // T3 cancel.markCancelling CAS 1→2 → R{State: Cancelling, Version: 2} - // T4 cancel.findActiveBatch(R) → none (batch has not been Created yet) - // T5 cancel.cancelRequest CAS 2→3 → R{State: Cancelled, Version: 3} - // T6 batch.IsRequestStateHalted(R) → false (stale in-memory copy from T1) - // T7 batch.BatchStore.Create(B{[R]}) → orphan batch containing a cancelled R - // - // After T7 the orphan batch flows through score → speculate → merge → conclude; - // conclude does NOT gate on the source request state when writing the terminal - // state, so it would CAS the request from Cancelled back to Landed, silently - // undoing the user's cancel. - // - // The CAS below collapses that window. Whichever of batch.UpdateState(..., - // RequestStateBatched) and cancel.markCancelling(... RequestStateCancelling) - // reaches storage first wins; the loser sees storage.ErrVersionMismatch: - // - If cancel won: this CAS fails. We ack the message (cancel will drive R - // to its terminal state on its own; no batch is needed). The reverse-index - // entry above becomes a dangling BatchDependent — tolerated per the - // "downstream should handle stale entries" contract on this store. - // - If batch won: cancel.markCancelling will fail with ErrVersionMismatch - // on its next attempt, re-fetch R, observe RequestStateBatched, and take - // the batch-cancellation branch (which terminates the whole batch). - // - // Note on re-delivery: a retry of a batch message that already CAS'd R to - // Batched but failed before/after BatchStore.Create lands in this code with - // R already in RequestStateBatched. The top-level IsRequestStateHalted check - // does NOT include Batched (Batched is forward-progress, not halted), so we - // reach here and re-CAS Batched → Batched (a version-only bump). The bump - // keeps the same serialization invariant on every attempt — if cancel sneaks - // in between our Get and this CAS, our version is stale and we abandon, just - // like the first-delivery case. The cost is an extra batch (the previous - // attempt may have already created one) which is tolerated per the comment - // on BatchStore.Create below. - // - // Residual window: a thin race remains between this CAS and BatchStore.Create. - // During that window cancel.findActiveBatch can still observe R in Batched - // with no batch yet persisted, and take the request-only cancel path — which - // then leaves R in Cancelled and the batch we are about to create orphaned. - // Fully closing this requires cancel-side wait/retry when its pre-CAS - // observation was RequestStateBatched; deferred to a follow-up since the - // window is narrow (one storage round-trip) and the user-visible outcome - // (request cancelled) is still correct — the orphan batch just gets - // reconciled by conclude as if it had no requests to act on. - newRequestVersion := request.Version + 1 - if err := c.store.GetRequestStore().UpdateState(ctx, request.ID, request.Version, newRequestVersion, entity.RequestStateBatched); err != nil { - // ErrVersionMismatch == cancel (or another writer) advanced R first. Ack - // the message: there is nothing for us to do, and retrying would not help - // since the new state of R is now visible to the cancel pipeline. - if errors.Is(err, storage.ErrVersionMismatch) { - c.metricsScope.Counter("request_claim_lost_race").Inc(1) - c.logger.Infow("abandoning batch creation; request advanced concurrently (likely cancel)", - "request_id", request.ID, - "request_version", request.Version, - "unused_batch_id", batch.ID, - ) - return nil + for _, dependencyID := range batch.Dependencies { + if err := c.ensureDependent(ctx, dependencyID, batch.ID); err != nil { + return err } - c.metricsScope.Counter("request_claim_errors").Inc(1) - return fmt.Errorf("failed to claim request %s for batch %s: %w", request.ID, batch.ID, err) } - request.Version = newRequestVersion - request.State = entity.RequestStateBatched + return nil +} - // Persist batch to storage. - // This is the final operation that concludes the batch creation process. If it fails, BatchDependents will be pointing to a batch id that does not exist. - // We do not reuse batch ids, a retry of this operation will create a new batch with a new ID. The downstream logic that operates on BatchDependent should be able to handle stale entries. - if err := c.store.GetBatchStore().Create(ctx, batch); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "batch_store_errors", 1) - return fmt.Errorf("failed to create batch in batch store: %w", err) +func (c *Controller) ensureDependent(ctx context.Context, dependencyID, dependentID string) error { + store := c.store.GetBatchDependentStore() + existing, err := store.Get(ctx, dependencyID) + if err != nil { + return fmt.Errorf("failed to get dependent index for batch %s: %w", dependencyID, err) } - - c.logger.Infow("batch created", - "batch_id", batch.ID, - "request_id", request.ID, - "queue", request.Queue, - "dependency_count", len(batch.Dependencies), - ) - - // Record the "batched" status in the request log. This status corresponds to - // the RequestStateBatched transition CAS'd above, so it carries the request - // version for reconciliation (unlike the batch-level "scored" status). The - // message ID is scoped to (requestID, status), so a redelivery that creates a - // fresh batch re-emits "batched" with a different batch_id but is deduped to - // the first entry — acceptable, the request is batched either way. - logEntry := entity.NewRequestLog(request.ID, entity.RequestStatusBatched, request.Version, "", map[string]string{ - "batch_id": batch.ID, - }) - if err := corerequest.PublishLog(ctx, c.registry, logEntry, request.ID); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "request_log_errors", 1) - return fmt.Errorf("failed to publish request log for request %s: %w", request.ID, err) + if contains(existing.Dependents, dependentID) { + return nil } - // Publish to score topic for further processing. - // If it fails and the controller retries, a new batch will be created with the new batch ID but the same request ID. - // The downstream logic should be able to handle stale entries by looking at the state of the batch. - if err := c.publish(ctx, topickey.TopicKeyScore, batch.ID, batch.Queue); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "publish_errors", 1) - return fmt.Errorf("failed to publish batch ID to score topic: %w", err) + dependents := append(append([]string(nil), existing.Dependents...), dependentID) + newVersion := existing.Version + 1 + if err := store.UpdateDependents(ctx, dependencyID, existing.Version, newVersion, dependents); err != nil { + if errors.Is(err, storage.ErrVersionMismatch) { + c.metricsScope.Counter("dependent_version_conflicts").Inc(1) + } + return fmt.Errorf("failed to add dependent %s to batch %s: %w", dependentID, dependencyID, err) } + return nil +} - c.logger.Infow("published batch to score topic", - "batch_id", batch.ID, - "topic_key", topickey.TopicKeyScore, - ) - - return nil // Success - message will be acked +func contains(values []string, value string) bool { + for _, candidate := range values { + if candidate == value { + return true + } + } + return false } // publish publishes a batch ID to the specified topic key. diff --git a/submitqueue/orchestrator/controller/batch/batch_test.go b/submitqueue/orchestrator/controller/batch/batch_test.go index fe02f62f..2b86738a 100644 --- a/submitqueue/orchestrator/controller/batch/batch_test.go +++ b/submitqueue/orchestrator/controller/batch/batch_test.go @@ -101,6 +101,27 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() } + var requestBatch entity.RequestBatch + mockRequestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + mockRequestBatchStore.EXPECT().Get(gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, requestID string) (entity.RequestBatch, error) { + if requestBatch.RequestID == "" { + return entity.RequestBatch{}, fmt.Errorf("request batch not found: %w", storage.ErrNotFound) + } + return requestBatch, nil + }, + ).AnyTimes() + mockRequestBatchStore.EXPECT().Create(gomock.Any(), gomock.Any()).DoAndReturn( + func(ctx context.Context, assignment entity.RequestBatch) error { + if requestBatch.RequestID != "" { + return fmt.Errorf("request batch exists: %w", storage.ErrAlreadyExists) + } + requestBatch = assignment + return nil + }, + ).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() + if analyzer == nil { analyzer = all.New() } @@ -176,9 +197,20 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) { mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) + mockRequestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + mockRequestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return( + entity.RequestBatch{}, fmt.Errorf("missing: %w", storage.ErrNotFound), + ) + mockRequestBatchStore.EXPECT().Create(gomock.Any(), entity.RequestBatch{ + RequestID: request.ID, + BatchID: "test-queue/batch/1", + Version: 1, + }).Return(nil) + mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() // Capture messages published to the log topic. @@ -246,19 +278,139 @@ func TestController_Process_StorageFailure(t *testing.T) { assert.Error(t, err) } -func TestController_Process_PublishFailure(t *testing.T) { +func TestController_Process_RedeliveryReusesBatchAndReplaysFanout(t *testing.T) { ctrl := gomock.NewController(t) - controller := newTestController(t, ctrl, newSequentialCounter(ctrl), nil, nil, fmt.Errorf("publish failed")) - request := testRequest() + batchedRequest := request + batchedRequest.State = entity.RequestStateBatched + batchedRequest.Version++ + requestBatch := entity.RequestBatch{ + RequestID: request.ID, + BatchID: "test-queue/batch/1", + Version: 1, + } + batch := entity.Batch{ + ID: requestBatch.BatchID, + Queue: request.Queue, + Contains: []string{request.ID}, + State: entity.BatchStateCreated, + Version: 1, + } + + mockReqStore := storagemock.NewMockRequestStore(ctrl) + gomock.InOrder( + mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil), + mockReqStore.EXPECT().UpdateState( + gomock.Any(), request.ID, request.Version, batchedRequest.Version, entity.RequestStateBatched, + ).Return(nil), + mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(batchedRequest, nil), + ) + + mockRequestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + gomock.InOrder( + mockRequestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return( + entity.RequestBatch{}, fmt.Errorf("missing: %w", storage.ErrNotFound), + ), + mockRequestBatchStore.EXPECT().Create(gomock.Any(), requestBatch).Return(nil), + mockRequestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return(requestBatch, nil), + ) + + mockBatchStore := storagemock.NewMockBatchStore(ctrl) + gomock.InOrder( + mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.DependencyBatchStates()).Return(nil, nil), + mockBatchStore.EXPECT().Create(gomock.Any(), batch).Return(nil), + mockBatchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil), + ) + + mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + gomock.InOrder( + mockBatchDependentStore.EXPECT().Create(gomock.Any(), entity.BatchDependent{ + BatchID: batch.ID, + Dependents: []string{}, + Version: 1, + }).Return(nil), + mockBatchDependentStore.EXPECT().Create(gomock.Any(), entity.BatchDependent{ + BatchID: batch.ID, + Dependents: []string{}, + Version: 1, + }).Return(fmt.Errorf("exists: %w", storage.ErrAlreadyExists)), + ) + + mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() + mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() + mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + + var logMessages []entityqueue.Message + var scoreMessages []entityqueue.Message + publishErr := fmt.Errorf("publish failed") + mockPub := queuemock.NewMockPublisher(ctrl) + gomock.InOrder( + mockPub.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).DoAndReturn( + func(ctx context.Context, topic string, msg entityqueue.Message) error { + logMessages = append(logMessages, msg) + return nil + }, + ), + mockPub.EXPECT().Publish(gomock.Any(), "score", gomock.Any()).DoAndReturn( + func(ctx context.Context, topic string, msg entityqueue.Message) error { + scoreMessages = append(scoreMessages, msg) + return publishErr + }, + ), + mockPub.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).DoAndReturn( + func(ctx context.Context, topic string, msg entityqueue.Message) error { + logMessages = append(logMessages, msg) + return nil + }, + ), + mockPub.EXPECT().Publish(gomock.Any(), "score", gomock.Any()).DoAndReturn( + func(ctx context.Context, topic string, msg entityqueue.Message) error { + scoreMessages = append(scoreMessages, msg) + return nil + }, + ), + ) + mockQ := queuemock.NewMockQueue(ctrl) + mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeyScore, Name: "score", Queue: mockQ}, + {Key: topickey.TopicKeyLog, Name: "log", Queue: mockQ}, + }) + require.NoError(t, err) + + analyzerFactory := conflictmock.NewMockFactory(ctrl) + analyzerFactory.EXPECT().For(gomock.Any()).Return(all.New(), nil) + controller := NewController( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + registry, + newSequentialCounter(ctrl), + mockStorage, + analyzerFactory, + topickey.TopicKeyBatch, + "orchestrator-batch", + ) + 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) - assert.Error(t, err) + require.ErrorIs(t, controller.Process(context.Background(), delivery), publishErr) + require.NoError(t, controller.Process(context.Background(), delivery)) + + require.Len(t, logMessages, 2) + assert.Equal(t, logMessages[0].ID, logMessages[1].ID) + assert.Equal(t, logMessages[0].Payload, logMessages[1].Payload) + assert.Equal(t, logMessages[0].PartitionKey, logMessages[1].PartitionKey) + + require.Len(t, scoreMessages, 2) + assert.Equal(t, scoreMessages[0].ID, scoreMessages[1].ID) + assert.Equal(t, scoreMessages[0].Payload, scoreMessages[1].Payload) + assert.Equal(t, scoreMessages[0].PartitionKey, scoreMessages[1].PartitionKey) } func TestController_Process_CounterFailure(t *testing.T) { @@ -319,7 +471,9 @@ func TestController_Process_WithDependencies(t *testing.T) { mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - mockReqStore.EXPECT().UpdateState(gomock.Any(), request.ID, request.Version, request.Version+1, entity.RequestStateBatched).Return(nil) + mockReqStore.EXPECT().UpdateState( + gomock.Any(), request.ID, request.Version, request.Version+1, entity.RequestStateBatched, + ).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() @@ -337,6 +491,90 @@ func TestController_Process_WithDependencies(t *testing.T) { require.NoError(t, err) } +func TestController_Process_RedeliveryDoesNotDuplicateDependent(t *testing.T) { + ctrl := gomock.NewController(t) + + request := testRequest() + request.State = entity.RequestStateBatched + request.Version = 2 + requestBatch := entity.RequestBatch{ + RequestID: request.ID, + BatchID: "test-queue/batch/3", + Version: 1, + } + batch := entity.Batch{ + ID: requestBatch.BatchID, + Queue: request.Queue, + Contains: []string{request.ID}, + Dependencies: []string{"test-queue/batch/2"}, + State: entity.BatchStateCreated, + Version: 1, + } + + mockReqStore := storagemock.NewMockRequestStore(ctrl) + mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + + mockRequestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + mockRequestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return(requestBatch, nil) + + mockBatchStore := storagemock.NewMockBatchStore(ctrl) + mockBatchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + + mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) + mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return( + fmt.Errorf("exists: %w", storage.ErrAlreadyExists), + ) + mockBatchDependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/2").Return(entity.BatchDependent{ + BatchID: "test-queue/batch/2", + Dependents: []string{batch.ID}, + Version: 4, + }, nil) + // No UpdateDependents expectation. The existing edge is already the + // reconciled postcondition. + + mockStorage := storagemock.NewMockStorage(ctrl) + mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() + mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() + mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).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(2).AnyTimes() + + require.NoError(t, controller.Process(context.Background(), delivery)) +} + +func TestController_EnsureDependentVersionConflictReturnsForRedelivery(t *testing.T) { + ctrl := gomock.NewController(t) + + dependentStore := storagemock.NewMockBatchDependentStore(ctrl) + dependentStore.EXPECT().Get(gomock.Any(), "test-queue/batch/1").Return(entity.BatchDependent{ + BatchID: "test-queue/batch/1", + Version: 3, + }, nil) + dependentStore.EXPECT().UpdateDependents( + gomock.Any(), + "test-queue/batch/1", + int32(3), + int32(4), + []string{"test-queue/batch/2"}, + ).Return(fmt.Errorf("cas: %w", storage.ErrVersionMismatch)) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchDependentStore().Return(dependentStore) + + controller := &Controller{ + metricsScope: tally.NoopScope, + store: store, + } + + err := controller.ensureDependent(context.Background(), "test-queue/batch/1", "test-queue/batch/2") + require.ErrorIs(t, err, storage.ErrVersionMismatch) +} + func TestController_Process_AnalyzerSelectsSubset(t *testing.T) { ctrl := gomock.NewController(t) @@ -363,7 +601,9 @@ func TestController_Process_AnalyzerSelectsSubset(t *testing.T) { mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - mockReqStore.EXPECT().UpdateState(gomock.Any(), request.ID, request.Version, request.Version+1, entity.RequestStateBatched).Return(nil) + mockReqStore.EXPECT().UpdateState( + gomock.Any(), request.ID, request.Version, request.Version+1, entity.RequestStateBatched, + ).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() @@ -399,6 +639,9 @@ func TestController_Process_AnalyzerFailure(t *testing.T) { mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) + mockReqStore.EXPECT().UpdateState( + gomock.Any(), request.ID, request.Version, request.Version+1, entity.RequestStateBatched, + ).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() @@ -450,17 +693,17 @@ func TestController_Process_HaltedShortCircuit(t *testing.T) { request.State = state request.Version = 7 - // Batch store with no EXPECTs — must not be queried. + // The batch store has no expectations and must not be queried. mockBatchStore := storagemock.NewMockBatchStore(ctrl) mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - // No UpdateState expected — gomock fails if called. + // No UpdateState call is expected; gomock fails if it is called. mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() - // Counter with no EXPECTs — must not be called. + // The counter has no expectations and must not be called. cnt := countermock.NewMockCounter(ctrl) controller := newTestController(t, ctrl, cnt, mockStorage, nil, fmt.Errorf("should not publish")) @@ -475,42 +718,39 @@ func TestController_Process_HaltedShortCircuit(t *testing.T) { } } -// Race-lost path: the cancel controller's markCancelling CAS landed first, -// so the batch controller's request-claim CAS (Validated → Batched) fails -// with storage.ErrVersionMismatch. The controller must ack the message (the -// cancel pipeline now owns the request) and must NOT call BatchStore.Create -// or publish to the score topic. -// -// This test exercises the race where the halted check at the top of Process -// passed against a stale in-memory copy from the initial Get (the cancel -// controller's CAS landed between our Get and our UpdateState). The CAS -// failure is the safety net that prevents an orphan batch in that window. -func TestController_Process_CASLostToCancel(t *testing.T) { +// A request CAS conflict must return the version mismatch so consumer +// redelivery can reload and classify the latest durable state. +func TestController_Process_VersionConflictReturnsForRedelivery(t *testing.T) { ctrl := gomock.NewController(t) request := testRequest() mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil) - // Create must NOT be called — gomock fails if it is. - - mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) - // The reverse-index Create still runs because it precedes the CAS; this is - // tolerated per the "downstream handles stale entries" contract. - mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockReqStore := storagemock.NewMockRequestStore(ctrl) - mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - mockReqStore.EXPECT().UpdateState( - gomock.Any(), request.ID, request.Version, request.Version+1, entity.RequestStateBatched, - ).Return(fmt.Errorf("cas: %w", storage.ErrVersionMismatch)) + gomock.InOrder( + mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil), + mockReqStore.EXPECT().UpdateState( + gomock.Any(), request.ID, request.Version, request.Version+1, entity.RequestStateBatched, + ).Return(fmt.Errorf("cas: %w", storage.ErrVersionMismatch)), + ) + + mockRequestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) + mockRequestBatchStore.EXPECT().Get(gomock.Any(), request.ID).Return( + entity.RequestBatch{}, fmt.Errorf("missing: %w", storage.ErrNotFound), + ) + mockRequestBatchStore.EXPECT().Create(gomock.Any(), entity.RequestBatch{ + RequestID: request.ID, + BatchID: "test-queue/batch/1", + Version: 1, + }).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() - mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() + mockStorage.EXPECT().GetRequestBatchStore().Return(mockRequestBatchStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() - // Publisher with no EXPECTs — must not be called. + // Publisher with no expectations must not be called. mockPub := queuemock.NewMockPublisher(ctrl) mockQ := queuemock.NewMockQueue(ctrl) mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() @@ -521,7 +761,6 @@ func TestController_Process_CASLostToCancel(t *testing.T) { require.NoError(t, err) analyzerFactory := conflictmock.NewMockFactory(ctrl) - analyzerFactory.EXPECT().For(gomock.Any()).Return(all.New(), nil).AnyTimes() controller := NewController( zaptest.NewLogger(t).Sugar(), tally.NoopScope, registry, newSequentialCounter(ctrl), mockStorage, analyzerFactory, topickey.TopicKeyBatch, "orchestrator-batch", @@ -532,25 +771,18 @@ func TestController_Process_CASLostToCancel(t *testing.T) { delivery.EXPECT().Message().Return(msg).AnyTimes() delivery.EXPECT().Attempt().Return(1).AnyTimes() - require.NoError(t, controller.Process(context.Background(), delivery)) + err = controller.Process(context.Background(), delivery) + require.ErrorIs(t, err, storage.ErrVersionMismatch) } -// Race-unexpected-error: any CAS failure other than ErrVersionMismatch (e.g. -// transient storage error) must surface as an error so the message is nacked -// for retry. We must NOT call BatchStore.Create on the way out. +// A non-conflict CAS error must be preserved so the consumer can classify it. func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { ctrl := gomock.NewController(t) request := testRequest() + casErr := fmt.Errorf("db connection lost") mockBatchStore := storagemock.NewMockBatchStore(ctrl) - mockBatchStore.EXPECT().GetByQueueAndStates(gomock.Any(), "test-queue", gomock.Any()).Return(nil, nil) - // Create must NOT be called — gomock fails if it is. - - mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) - mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) - - casErr := fmt.Errorf("db connection lost") mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) mockReqStore.EXPECT().UpdateState( @@ -559,7 +791,6 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { 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) @@ -571,36 +802,34 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { err := controller.Process(context.Background(), delivery) require.Error(t, err) - // Cause must be preserved for upstream classification. assert.True(t, errors.Is(err, casErr)) } -// Recovery path: a re-delivered batch message whose prior attempt CAS'd the -// request to RequestStateBatched but failed before BatchStore.Create. The -// halted check at the top of Process does NOT include Batched (Batched is -// forward-progress, not halted), so we reach the CAS again and re-bump the -// version on the request (Batched → Batched, version+1). The batch is then -// re-created with a new batch ID, which is tolerated per the existing -// duplicate-handling comment on BatchStore.Create. +// A batched request without a request_batch row can recover the one active +// legacy batch that already contains it. The request version is not bumped. func TestController_Process_RecoveryAfterPriorCAS(t *testing.T) { ctrl := gomock.NewController(t) request := testRequest() request.State = entity.RequestStateBatched - request.Version = 2 // prior attempt bumped from 1 → 2 + request.Version = 2 + batch := entity.Batch{ + ID: "test-queue/batch/7", + Queue: request.Queue, + Contains: []string{request.ID}, + State: entity.BatchStateCreated, + Version: 1, + } 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().GetByQueueAndStates(gomock.Any(), request.Queue, entity.ActiveBatchStates()).Return([]entity.Batch{batch}, nil) + mockBatchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) mockReqStore := storagemock.NewMockRequestStore(ctrl) mockReqStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - mockReqStore.EXPECT().UpdateState( - gomock.Any(), request.ID, request.Version, request.Version+1, entity.RequestStateBatched, - ).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes()