From a52c7a54830561d919a8cbbabfa966b2a19d36dc Mon Sep 17 00:00:00 2001 From: abettigole Date: Sat, 1 Aug 2026 00:08:19 +0000 Subject: [PATCH] refactor(storage): replace request state updates Persist complete Request entities through guarded updates, migrate callers to candidate copies, and document full-entity replacement semantics. Jira Issues CODEM-204 --- submitqueue/core/request/terminate.go | 10 ++- submitqueue/core/request/terminate_test.go | 8 +- submitqueue/entity/request.go | 7 ++ submitqueue/entity/request_test.go | 19 +++++ submitqueue/extension/storage/README.md | 22 +++-- .../storage/mock/request_store_mock.go | 12 +-- .../extension/storage/mysql/request_store.go | 28 ++++--- .../storage/mysql/request_store_test.go | 62 +++++++++++--- .../extension/storage/request_store.go | 7 +- .../orchestrator/controller/batch/batch.go | 9 ++- .../controller/batch/batch_test.go | 34 ++++---- .../orchestrator/controller/cancel/cancel.go | 8 +- .../controller/cancel/cancel_test.go | 51 +++++++----- .../controller/conclude/conclude_test.go | 35 ++++---- .../orchestrator/controller/dlq/batch_test.go | 7 +- .../controller/dlq/buildsignal_test.go | 7 +- .../orchestrator/controller/dlq/dlq_test.go | 49 ++++++----- .../dlq/mergeconflictsignal_test.go | 7 +- .../controller/dlq/mergesignal_test.go | 7 +- .../controller/dlq/request_test.go | 21 ++--- .../mergeconflictsignal.go | 14 ++-- .../mergeconflictsignal_test.go | 32 ++++++-- .../controller/validate/validate_test.go | 6 +- .../submitqueue/extension/storage/suite.go | 81 +++++++++++++++---- 24 files changed, 366 insertions(+), 177 deletions(-) diff --git a/submitqueue/core/request/terminate.go b/submitqueue/core/request/terminate.go index 736cad51..37e2aeca 100644 --- a/submitqueue/core/request/terminate.go +++ b/submitqueue/core/request/terminate.go @@ -114,6 +114,7 @@ func TerminateRequest( // logVersion is the request version reflected in the published terminal log. // It stays at the current version on the idempotent same-state path and // advances to the new version only after a successful reconciling write. + beforeState := request.State logVersion := request.Version outcome := TerminationOutcomeSuccess switch { @@ -131,10 +132,13 @@ func TerminateRequest( }, nil default: newVersion := request.Version + 1 - if err := store.GetRequestStore().UpdateState(ctx, requestID, request.Version, newVersion, targetState); err != nil { + updatedRequest := request.WithState(targetState) + if err := store.GetRequestStore().Update(ctx, updatedRequest, request.Version, newVersion); err != nil { return TerminationResult{}, fmt.Errorf("failed to update request %s state to %s: %w", requestID, targetState, err) } - logVersion = newVersion + updatedRequest.Version = newVersion + request = updatedRequest + logVersion = request.Version } logEntry := entity.NewRequestLog(requestID, status, logVersion, lastError, metadata) @@ -144,7 +148,7 @@ func TerminateRequest( return TerminationResult{ Outcome: outcome, - BeforeState: request.State, + BeforeState: beforeState, AfterState: targetState, }, nil } diff --git a/submitqueue/core/request/terminate_test.go b/submitqueue/core/request/terminate_test.go index a9b1b9b7..1058b01d 100644 --- a/submitqueue/core/request/terminate_test.go +++ b/submitqueue/core/request/terminate_test.go @@ -61,6 +61,7 @@ func TestTerminateRequest(t *testing.T) { const requestID = "q/1" validated := entity.Request{ID: requestID, Queue: "q", State: entity.RequestStateValidated, Version: 3} + originalValidated := validated testCases := map[string]struct { targetState entity.RequestState @@ -101,7 +102,7 @@ func TestTerminateRequest(t *testing.T) { metadata: map[string]string{"source": "validate"}, mockFunc: func(rs *storagemock.MockRequestStore) { rs.EXPECT().Get(gomock.Any(), requestID).Return(validated, nil) - rs.EXPECT().UpdateState(gomock.Any(), requestID, int32(3), int32(4), entity.RequestStateError).Return(nil) + rs.EXPECT().Update(gomock.Any(), validated.WithState(entity.RequestStateError), int32(3), int32(4)).Return(nil) }, wantResult: TerminationResult{ Outcome: TerminationOutcomeSuccess, @@ -157,7 +158,7 @@ func TestTerminateRequest(t *testing.T) { targetState: entity.RequestStateError, mockFunc: func(rs *storagemock.MockRequestStore) { rs.EXPECT().Get(gomock.Any(), requestID).Return(validated, nil) - rs.EXPECT().UpdateState(gomock.Any(), requestID, int32(3), int32(4), entity.RequestStateError).Return(storage.ErrVersionMismatch) + rs.EXPECT().Update(gomock.Any(), validated.WithState(entity.RequestStateError), int32(3), int32(4)).Return(storage.ErrVersionMismatch) }, wantResult: TerminationResult{Outcome: TerminationOutcomeUnknown}, errMsg: "version mismatch", @@ -167,7 +168,7 @@ func TestTerminateRequest(t *testing.T) { targetState: entity.RequestStateError, mockFunc: func(rs *storagemock.MockRequestStore) { rs.EXPECT().Get(gomock.Any(), requestID).Return(validated, nil) - rs.EXPECT().UpdateState(gomock.Any(), requestID, int32(3), int32(4), entity.RequestStateError).Return(nil) + rs.EXPECT().Update(gomock.Any(), validated.WithState(entity.RequestStateError), int32(3), int32(4)).Return(nil) }, publishErr: fmt.Errorf("connection refused"), wantResult: TerminationResult{Outcome: TerminationOutcomeUnknown}, @@ -187,6 +188,7 @@ func TestTerminateRequest(t *testing.T) { res, err := TerminateRequest(context.Background(), store, registry, requestID, tc.targetState, tc.lastError, tc.metadata) + assert.Equal(t, originalValidated, validated) assert.Equal(t, tc.wantResult, res) if tc.errMsg != "" { assert.ErrorContains(t, err, tc.errMsg) diff --git a/submitqueue/entity/request.go b/submitqueue/entity/request.go index 0b88a25b..a96fdda0 100644 --- a/submitqueue/entity/request.go +++ b/submitqueue/entity/request.go @@ -96,6 +96,13 @@ type Request struct { Version int32 `json:"version"` } +// WithState returns a shallow copy of the request with State replaced. +// Slice fields continue to share their backing arrays. +func (r Request) WithState(state RequestState) Request { + r.State = state + return r +} + // ToBytes serializes the Request to JSON bytes for queue message payload. func (r Request) ToBytes() ([]byte, error) { return json.Marshal(r) diff --git a/submitqueue/entity/request_test.go b/submitqueue/entity/request_test.go index 4415fb20..b0322ad1 100644 --- a/submitqueue/entity/request_test.go +++ b/submitqueue/entity/request_test.go @@ -96,6 +96,25 @@ func TestRequestFromBytes_EmptyData(t *testing.T) { assert.Equal(t, int32(0), req.Version) } +func TestRequest_WithState(t *testing.T) { + request := Request{ + ID: "queueA/1", + Queue: "queueA", + State: RequestStateStarted, + Version: 1, + } + + updated := request.WithState(RequestStateValidated) + + assert.Equal(t, RequestStateStarted, request.State) + assert.Equal(t, Request{ + ID: request.ID, + Queue: request.Queue, + State: RequestStateValidated, + Version: request.Version, + }, updated) +} + func TestIsRequestStateTerminal(t *testing.T) { tests := []struct { state RequestState diff --git a/submitqueue/extension/storage/README.md b/submitqueue/extension/storage/README.md index 202a6ee9..2b3a3244 100644 --- a/submitqueue/extension/storage/README.md +++ b/submitqueue/extension/storage/README.md @@ -6,25 +6,31 @@ Pluggable persistence interfaces for SubmitQueue entities (requests, batches, de Entities that support concurrent mutation carry an `int32 Version` field. Updates are conditional on the version: the write only succeeds if the persisted version matches the caller's expected version. On mismatch, the implementation returns `storage.ErrVersionMismatch`, which is declared as a retryable infrastructure error so callers can return it without reclassifying it. -**Version arithmetic is owned by the controller, not the store.** Update methods take both `oldVersion` (the where-clause guard) and `newVersion` (the value to write): +**Updates replace every non-primary-key field.** Callers must pass a complete authoritative entity loaded from storage or constructed with every persisted field; sparse patch entities can clear unrelated columns. The primary key identifies the row and is not rewritten. + +**Version arithmetic is owned by the controller, not the store.** Versioned update methods take a complete entity plus both `oldVersion` (the where-clause guard) and `newVersion` (the value to write): ```go -UpdateState(ctx, id, oldVersion, newVersion int32, newState entity.RequestState) error +Update(ctx, request entity.Request, oldVersion, newVersion int32) error ``` -The store performs a pure conditional write — it does not compute `oldVersion + 1` internally. This keeps the in-memory entity and the persisted row in sync without the storage layer mutating values the caller didn't supply. +The store writes `newVersion` rather than the entity's current `Version` and performs a pure conditional write — it does not compute `oldVersion + 1` internally. This keeps the in-memory entity and the persisted row in sync without the storage layer mutating values the caller didn't supply. ### Caller pattern ```go -newVersion := entity.Version + 1 -if err := store.UpdateState(ctx, entity.ID, entity.Version, newVersion, newState); err != nil { - return err // entity.Version unchanged on failure — safe to retry +oldVersion := request.Version +newVersion := oldVersion + 1 +updated := request +updated.State = newState +if err := store.Update(ctx, updated, oldVersion, newVersion); err != nil { + return err // request remains unchanged on failure — safe to retry } -entity.Version = newVersion // only after the write succeeded +updated.Version = newVersion +request = updated // only after the write succeeded ``` -The post-success assignment matters whenever the entity is read again later in the same flow. Pre-incrementing in memory before the call is a bug pattern: if the call fails and the caller swallows the error, the in-memory version is now ahead of the database and subsequent updates will fail with `ErrVersionMismatch` for non-obvious reasons. +The candidate-copy pattern keeps the caller-owned entity unchanged if the write fails. Clone slice and map fields before changing their contents so the candidate cannot mutate the original through shared backing storage. The post-success assignment matters whenever the entity is read again later in the same flow. Pre-incrementing in memory before the call is a bug pattern: if the call fails and the caller swallows the error, the in-memory version is now ahead of the database and subsequent updates will fail with `ErrVersionMismatch` for non-obvious reasons. ## Read-after-write consistency diff --git a/submitqueue/extension/storage/mock/request_store_mock.go b/submitqueue/extension/storage/mock/request_store_mock.go index f2a19066..1bdc4696 100644 --- a/submitqueue/extension/storage/mock/request_store_mock.go +++ b/submitqueue/extension/storage/mock/request_store_mock.go @@ -70,16 +70,16 @@ func (mr *MockRequestStoreMockRecorder) Get(ctx, id any) *gomock.Call { return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Get", reflect.TypeOf((*MockRequestStore)(nil).Get), ctx, id) } -// UpdateState mocks base method. -func (m *MockRequestStore) UpdateState(ctx context.Context, id string, oldVersion, newVersion int32, newState entity.RequestState) error { +// Update mocks base method. +func (m *MockRequestStore) Update(ctx context.Context, request entity.Request, oldVersion, newVersion int32) error { m.ctrl.T.Helper() - ret := m.ctrl.Call(m, "UpdateState", ctx, id, oldVersion, newVersion, newState) + ret := m.ctrl.Call(m, "Update", ctx, request, oldVersion, newVersion) ret0, _ := ret[0].(error) return ret0 } -// UpdateState indicates an expected call of UpdateState. -func (mr *MockRequestStoreMockRecorder) UpdateState(ctx, id, oldVersion, newVersion, newState any) *gomock.Call { +// Update indicates an expected call of Update. +func (mr *MockRequestStoreMockRecorder) Update(ctx, request, oldVersion, newVersion any) *gomock.Call { mr.mock.ctrl.T.Helper() - return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "UpdateState", reflect.TypeOf((*MockRequestStore)(nil).UpdateState), ctx, id, oldVersion, newVersion, newState) + return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Update", reflect.TypeOf((*MockRequestStore)(nil).Update), ctx, request, oldVersion, newVersion) } diff --git a/submitqueue/extension/storage/mysql/request_store.go b/submitqueue/extension/storage/mysql/request_store.go index 2860246d..01a4487f 100644 --- a/submitqueue/extension/storage/mysql/request_store.go +++ b/submitqueue/extension/storage/mysql/request_store.go @@ -93,36 +93,40 @@ func (r *requestStore) Create(ctx context.Context, request entity.Request) (retE return nil } -// UpdateState updates the state of a land request to newState and the version to 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 (r *requestStore) UpdateState(ctx context.Context, id string, oldVersion, newVersion int32, newState entity.RequestState) (retErr error) { +// Update replaces every non-key field of a land request 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 (r *requestStore) Update(ctx context.Context, request entity.Request, oldVersion, newVersion int32) (retErr error) { op := metrics.Begin(r.scope, "update_state", metrics.StorageLatencyBuckets) defer func() { op.Complete(retErr) }() + changeURIsJSON, err := json.Marshal(request.Change.URIs) + if err != nil { + return fmt.Errorf("failed to marshal change URIs for request id=%s: %w", request.ID, err) + } + result, err := r.db.ExecContext(ctx, - "UPDATE request SET state = ?, version = ? WHERE id = ? AND version = ?", - newState, newVersion, id, oldVersion, + "UPDATE request SET queue = ?, change_uri = ?, land_strategy = ?, state = ?, version = ? WHERE id = ? AND version = ?", + request.Queue, changeURIsJSON, request.LandStrategy, request.State, newVersion, request.ID, oldVersion, ) if err != nil { return fmt.Errorf( - "failed to update request state for id=%q oldVersion=%d newVersion=%d newState=%v: %w", - id, oldVersion, newVersion, newState, err, + "failed to update request for id=%q oldVersion=%d newVersion=%d: %w", + request.ID, oldVersion, newVersion, 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, + "failed to get rows affected from update for id=%q oldVersion=%d newVersion=%d: %w", + request.ID, oldVersion, newVersion, err, ) } if rowsAffected != 1 { return fmt.Errorf( - "version mismatch for request update: id=%q expected_version=%d newState=%v: %w", - id, oldVersion, newState, storage.ErrVersionMismatch, + "version mismatch for request update: id=%q expected_version=%d: %w", + request.ID, oldVersion, storage.ErrVersionMismatch, ) } diff --git a/submitqueue/extension/storage/mysql/request_store_test.go b/submitqueue/extension/storage/mysql/request_store_test.go index c13668b8..1c37bcb1 100644 --- a/submitqueue/extension/storage/mysql/request_store_test.go +++ b/submitqueue/extension/storage/mysql/request_store_test.go @@ -198,53 +198,91 @@ func TestRequestStore_Create(t *testing.T) { } } -func TestRequestStore_UpdateState(t *testing.T) { - const id = "monorepo/1" +func TestRequestStore_Update(t *testing.T) { const oldVersion, newVersion = int32(1), int32(2) - const newState = entity.RequestStateValidated + request := entity.Request{ + ID: "monorepo/1", + Queue: "monorepo-updated", + Change: change.Change{URIs: []string{"github://github.example.com/uber/submitqueue/pull/456/cafebabe"}}, + LandStrategy: mergestrategy.MergeStrategySquashRebase, + State: entity.RequestStateValidated, + Version: oldVersion, + } + changeURIsJSON, err := json.Marshal(request.Change.URIs) + require.NoError(t, err) tests := []struct { name string + request entity.Request setup func(mock sqlmock.Sqlmock) wantErr bool wantErrIs error }{ { - name: "success", + name: "success", + request: request, setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE request"). - WithArgs(newState, newVersion, id, oldVersion). + WithArgs(request.Queue, changeURIsJSON, request.LandStrategy, request.State, newVersion, request.ID, oldVersion). WillReturnResult(sqlmock.NewResult(0, 1)) }, }, { - name: "version mismatch", + name: "version mismatch", + request: request, setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE request"). - WithArgs(newState, newVersion, id, oldVersion). + WithArgs(request.Queue, changeURIsJSON, request.LandStrategy, request.State, newVersion, request.ID, oldVersion). WillReturnResult(sqlmock.NewResult(0, 0)) }, wantErr: true, wantErrIs: storage.ErrVersionMismatch, }, { - name: "exec error", + name: "exec error", + request: request, setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE request"). - WithArgs(newState, newVersion, id, oldVersion). + WithArgs(request.Queue, changeURIsJSON, request.LandStrategy, request.State, newVersion, request.ID, oldVersion). WillReturnError(fmt.Errorf("connection reset")) }, wantErr: true, }, { - name: "rows affected error", + name: "rows affected error", + request: request, setup: func(mock sqlmock.Sqlmock) { mock.ExpectExec("UPDATE request"). - WithArgs(newState, newVersion, id, oldVersion). + WithArgs(request.Queue, changeURIsJSON, request.LandStrategy, request.State, newVersion, request.ID, oldVersion). WillReturnResult(sqlmock.NewErrorResult(fmt.Errorf("driver error"))) }, wantErr: true, }, + { + name: "nil change URIs", + request: entity.Request{ID: request.ID, Queue: request.Queue, LandStrategy: request.LandStrategy, State: request.State, Version: request.Version}, + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("UPDATE request"). + WithArgs(request.Queue, []byte("null"), request.LandStrategy, request.State, newVersion, request.ID, oldVersion). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, + { + name: "empty change URIs", + request: entity.Request{ + ID: request.ID, + Queue: request.Queue, + Change: change.Change{URIs: []string{}}, + LandStrategy: request.LandStrategy, + State: request.State, + Version: request.Version, + }, + setup: func(mock sqlmock.Sqlmock) { + mock.ExpectExec("UPDATE request"). + WithArgs(request.Queue, []byte("[]"), request.LandStrategy, request.State, newVersion, request.ID, oldVersion). + WillReturnResult(sqlmock.NewResult(0, 1)) + }, + }, } for _, tt := range tests { @@ -254,7 +292,7 @@ func TestRequestStore_UpdateState(t *testing.T) { tt.setup(mock) - err := store.UpdateState(context.Background(), id, oldVersion, newVersion, newState) + err := store.Update(context.Background(), tt.request, oldVersion, newVersion) if tt.wantErr { require.Error(t, err) if tt.wantErrIs != nil { diff --git a/submitqueue/extension/storage/request_store.go b/submitqueue/extension/storage/request_store.go index fa3e7b27..00cf1ba9 100644 --- a/submitqueue/extension/storage/request_store.go +++ b/submitqueue/extension/storage/request_store.go @@ -31,8 +31,7 @@ type RequestStore interface { // Returns ErrAlreadyExists if a request with the same ID already exists. Create(ctx context.Context, request entity.Request) error - // UpdateState updates the state of a land request to newState and the version to 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.RequestState) error + // Update replaces every non-key field of a land request 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. + Update(ctx context.Context, request entity.Request, oldVersion, newVersion int32) error } diff --git a/submitqueue/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index 8f7e1392..0ffac598 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/submitqueue/orchestrator/controller/batch/batch.go @@ -190,7 +190,7 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // 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(..., + // The CAS below collapses that window. Whichever of request.Update(..., // 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 @@ -221,7 +221,8 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // (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 { + updatedRequest := request.WithState(entity.RequestStateBatched) + if err := c.store.GetRequestStore().Update(ctx, updatedRequest, request.Version, newRequestVersion); 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. @@ -237,8 +238,8 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er metrics.NamedCounter(c.metricsScope, opName, "request_claim_errors", 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 + updatedRequest.Version = newRequestVersion + request = updatedRequest // Persist the batch before creating references to it. A Creating batch is not eligible for dependency analysis or normal processing. if err := c.store.GetBatchStore().Create(ctx, batch); err != nil { diff --git a/submitqueue/orchestrator/controller/batch/batch_test.go b/submitqueue/orchestrator/controller/batch/batch_test.go index 23b1065a..99e3f71b 100644 --- a/submitqueue/orchestrator/controller/batch/batch_test.go +++ b/submitqueue/orchestrator/controller/batch/batch_test.go @@ -91,7 +91,7 @@ func newTestController(t *testing.T, ctrl *gomock.Controller, cnt *countermock.M mockReqStore := storagemock.NewMockRequestStore(ctrl) req := testRequest() mockReqStore.EXPECT().Get(gomock.Any(), req.ID).Return(req, nil).AnyTimes() - mockReqStore.EXPECT().UpdateState(gomock.Any(), req.ID, req.Version, req.Version+1, entity.RequestStateBatched).Return(nil).AnyTimes() + mockReqStore.EXPECT().Update(gomock.Any(), req.WithState(entity.RequestStateBatched), req.Version, req.Version+1).Return(nil).AnyTimes() mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil).AnyTimes() @@ -184,7 +184,7 @@ func TestController_Process_PublishesBatchedLog(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().Update(gomock.Any(), request.WithState(entity.RequestStateBatched), request.Version, request.Version+1).Return(nil) mockBatchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) mockBatchDependentStore.EXPECT().Create(gomock.Any(), gomock.Any()).Return(nil) @@ -279,7 +279,7 @@ func TestController_Process_RequestBatchStoreFailure(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, request.Version, request.Version+1, entity.RequestStateBatched).Return(nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateBatched), request.Version, request.Version+1).Return(nil) requestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) requestBatchStore.EXPECT().Create(gomock.Any(), entity.RequestBatch{ @@ -392,7 +392,7 @@ 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().Update(gomock.Any(), request.WithState(entity.RequestStateBatched), request.Version, request.Version+1).Return(nil) mockRequestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) mockRequestBatchStore.EXPECT().Create(gomock.Any(), entity.RequestBatch{ @@ -456,7 +456,7 @@ 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().Update(gomock.Any(), request.WithState(entity.RequestStateBatched), request.Version, request.Version+1).Return(nil) mockRequestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) mockRequestBatchStore.EXPECT().Create(gomock.Any(), entity.RequestBatch{ @@ -647,8 +647,8 @@ func TestController_Process_CASLostToCancel(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, + mockReqStore.EXPECT().Update( + gomock.Any(), request.WithState(entity.RequestStateBatched), request.Version, request.Version+1, ).Return(fmt.Errorf("cas: %w", storage.ErrVersionMismatch)) mockStorage := storagemock.NewMockStorage(ctrl) @@ -678,6 +678,8 @@ func TestController_Process_CASLostToCancel(t *testing.T) { delivery.EXPECT().Attempt().Return(1).AnyTimes() require.NoError(t, controller.Process(context.Background(), delivery)) + assert.Equal(t, entity.RequestStateStarted, request.State) + assert.Equal(t, int32(1), request.Version) } // Race-unexpected-error: any CAS failure other than ErrVersionMismatch (e.g. @@ -695,8 +697,8 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { casErr := fmt.Errorf("db connection lost") 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, + mockReqStore.EXPECT().Update( + gomock.Any(), request.WithState(entity.RequestStateBatched), request.Version, request.Version+1, ).Return(casErr) mockStorage := storagemock.NewMockStorage(ctrl) @@ -714,6 +716,8 @@ func TestController_Process_CASUnexpectedErrorPropagates(t *testing.T) { require.Error(t, err) // Cause must be preserved for upstream classification. assert.True(t, errors.Is(err, casErr)) + assert.Equal(t, entity.RequestStateStarted, request.State) + assert.Equal(t, int32(1), request.Version) } // Recovery path: a re-delivered batch message whose prior attempt CAS'd the @@ -747,8 +751,8 @@ func TestController_Process_RecoveryAfterPriorCAS(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, + mockReqStore.EXPECT().Update( + gomock.Any(), request.WithState(entity.RequestStateBatched), request.Version, request.Version+1, ).Return(nil) mockRequestBatchStore := storagemock.NewMockRequestBatchStore(ctrl) @@ -800,7 +804,7 @@ func TestController_Process_ReadiesBatchBeforePublishing(t *testing.T) { batchDependentStore := storagemock.NewMockBatchDependentStore(ctrl) publisher := queuemock.NewMockPublisher(ctrl) gomock.InOrder( - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(1), int32(2), entity.RequestStateBatched).Return(nil), + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateBatched), int32(1), int32(2)).Return(nil), batchStore.EXPECT().Create(gomock.Any(), batch).Return(nil), requestBatchStore.EXPECT().Create(gomock.Any(), entity.RequestBatch{ RequestID: request.ID, @@ -860,9 +864,9 @@ func TestController_Process_RedeliveryMintsFreshBatchID(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), firstRequest.ID).Return(firstRequest, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), firstRequest.ID, int32(1), int32(2), entity.RequestStateBatched).Return(nil) + requestStore.EXPECT().Update(gomock.Any(), firstRequest.WithState(entity.RequestStateBatched), int32(1), int32(2)).Return(nil) requestStore.EXPECT().Get(gomock.Any(), firstRequest.ID).Return(secondRequest, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), firstRequest.ID, int32(2), int32(3), entity.RequestStateBatched).Return(nil) + requestStore.EXPECT().Update(gomock.Any(), secondRequest.WithState(entity.RequestStateBatched), int32(2), int32(3)).Return(nil) var createdIDs []string batchStore := storagemock.NewMockBatchStore(ctrl) @@ -929,7 +933,7 @@ func TestController_Process_InitializationFailure(t *testing.T) { request := testRequest() requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(1), int32(2), entity.RequestStateBatched).Return(nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateBatched), int32(1), int32(2)).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) batchStore.EXPECT().GetByQueueAndStates(gomock.Any(), request.Queue, entity.DependencyBatchStates()).Return(nil, nil) diff --git a/submitqueue/orchestrator/controller/cancel/cancel.go b/submitqueue/orchestrator/controller/cancel/cancel.go index 366c3d22..70af615d 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel.go +++ b/submitqueue/orchestrator/controller/cancel/cancel.go @@ -200,14 +200,14 @@ func (c *Controller) markCancelling(ctx context.Context, request entity.Request) return request, nil } newVersion := request.Version + 1 - if err := c.store.GetRequestStore().UpdateState(ctx, request.ID, request.Version, newVersion, entity.RequestStateCancelling); err != nil { + updatedRequest := request.WithState(entity.RequestStateCancelling) + if err := c.store.GetRequestStore().Update(ctx, updatedRequest, request.Version, newVersion); err != nil { metrics.NamedCounter(c.metricsScope, opName, "request_update_errors", 1) return entity.Request{}, fmt.Errorf("failed to mark request %s as cancelling: %w", request.ID, err) } - request.Version = newVersion - request.State = entity.RequestStateCancelling + updatedRequest.Version = newVersion metrics.NamedCounter(c.metricsScope, opName, "request_cancelling", 1) - return request, nil + return updatedRequest, nil } // findBatches resolves every batch attempt associated with the request. diff --git a/submitqueue/orchestrator/controller/cancel/cancel_test.go b/submitqueue/orchestrator/controller/cancel/cancel_test.go index 9e08a4d8..7b6b12fe 100644 --- a/submitqueue/orchestrator/controller/cancel/cancel_test.go +++ b/submitqueue/orchestrator/controller/cancel/cancel_test.go @@ -162,9 +162,9 @@ func TestProcess_CancelsUnbatchedRequest(t *testing.T) { // the request — the terminal Cancelled CAS. gomock.InOrder( reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(started, nil), - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateCancelling).Return(nil), + reqStore.EXPECT().Update(gomock.Any(), started.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil), reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(cancelling, nil), - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(3), int32(4), entity.RequestStateCancelled).Return(nil), + reqStore.EXPECT().Update(gomock.Any(), cancelling.WithState(entity.RequestStateCancelled), int32(3), int32(4)).Return(nil), ) batchStore := storagemock.NewMockBatchStore(ctrl) @@ -197,7 +197,7 @@ func TestProcess_AlreadyCancelling_SkipsMarkCancelling(t *testing.T) { // see Cancelling (the prior pass already recorded intent). reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(cancelling, nil).Times(2) // Only the terminal CAS — the mark-cancelling step is a no-op. - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(3), int32(4), entity.RequestStateCancelled).Return(nil) + reqStore.EXPECT().Update(gomock.Any(), cancelling.WithState(entity.RequestStateCancelled), int32(3), int32(4)).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) @@ -222,10 +222,11 @@ func TestProcess_MarkCancellingVersionMismatch_Retryable(t *testing.T) { _ = pub reqStore := storagemock.NewMockRequestStore(ctrl) - reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + started := entity.Request{ ID: "q/1", Queue: "q", State: entity.RequestStateStarted, Version: 2, - }, nil) - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateCancelling). + } + reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(started, nil) + reqStore.EXPECT().Update(gomock.Any(), started.WithState(entity.RequestStateCancelling), int32(2), int32(3)). Return(storage.ErrVersionMismatch) store := storagemock.NewMockStorage(ctrl) @@ -235,6 +236,8 @@ func TestProcess_MarkCancellingVersionMismatch_Retryable(t *testing.T) { err := controller.Process(context.Background(), newDelivery(t, ctrl, cancelPayload(t, "q/1", ""), "q/1")) require.Error(t, err) assert.ErrorIs(t, err, storage.ErrVersionMismatch) + assert.Equal(t, entity.RequestStateStarted, started.State) + assert.Equal(t, int32(2), started.Version) } func TestProcess_UnbatchedVersionMismatch_Retryable(t *testing.T) { @@ -247,9 +250,9 @@ func TestProcess_UnbatchedVersionMismatch_Retryable(t *testing.T) { cancelling := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateCancelling, Version: 3} gomock.InOrder( reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(started, nil), - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateCancelling).Return(nil), + reqStore.EXPECT().Update(gomock.Any(), started.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil), reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(cancelling, nil), - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(3), int32(4), entity.RequestStateCancelled). + reqStore.EXPECT().Update(gomock.Any(), cancelling.WithState(entity.RequestStateCancelled), int32(3), int32(4)). Return(storage.ErrVersionMismatch), ) @@ -276,7 +279,7 @@ func TestProcess_UnbatchedRequestDiverged_Acks(t *testing.T) { reqStore := storagemock.NewMockRequestStore(ctrl) gomock.InOrder( reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(started, nil), - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateCancelling).Return(nil), + reqStore.EXPECT().Update(gomock.Any(), started.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil), reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(landed, nil), ) @@ -301,7 +304,7 @@ func TestProcess_UnbatchedRequestDisappears_Retryable(t *testing.T) { reqStore := storagemock.NewMockRequestStore(ctrl) gomock.InOrder( reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(started, nil), - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateCancelling).Return(nil), + reqStore.EXPECT().Update(gomock.Any(), started.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil), reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{}, storage.ErrNotFound), ) @@ -351,7 +354,7 @@ func TestProcess_BatchPath_HandsOffToSpeculate(t *testing.T) { reqStore := storagemock.NewMockRequestStore(ctrl) reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(req, nil) - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateCancelling).Return(nil) + reqStore.EXPECT().Update(gomock.Any(), req.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) // Single batch CAS: intent only. No terminal CAS. @@ -381,7 +384,7 @@ func TestProcess_CancelsEveryApplicableBatch(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil) var operations []string batchStore := storagemock.NewMockBatchStore(ctrl) @@ -430,7 +433,7 @@ func TestProcess_BatchFailureDoesNotPreventLaterCancellation(t *testing.T) { requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) batchStore.EXPECT().Update(gomock.Any(), batch1.WithState(entity.BatchStateCancelling), int32(1), int32(2)).Return(storeErr) @@ -473,7 +476,7 @@ func TestProcess_NonCancellableBatchSuppressesRequestCancellation(t *testing.T) requestStore := storagemock.NewMockRequestStore(ctrl) requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) @@ -494,11 +497,14 @@ func TestProcess_BatchedWithoutMatchCancelsRequest(t *testing.T) { publisher.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).Return(nil) request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateBatched, Version: 2} + cancelling := request.WithState(entity.RequestStateCancelling) + cancelling.Version = 3 requestStore := storagemock.NewMockRequestStore(ctrl) gomock.InOrder( requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil), - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil), - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(3), int32(4), entity.RequestStateCancelled).Return(nil), + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil), + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(cancelling, nil), + requestStore.EXPECT().Update(gomock.Any(), cancelling.WithState(entity.RequestStateCancelled), int32(3), int32(4)).Return(nil), ) batchStore := storagemock.NewMockBatchStore(ctrl) @@ -518,6 +524,8 @@ func TestProcess_CreatingBatchDoesNotSuppressRequestCancellation(t *testing.T) { publisher.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).Return(nil) request := entity.Request{ID: "q/1", Queue: "q", State: entity.RequestStateBatched, Version: 2} + cancelling := request.WithState(entity.RequestStateCancelling) + cancelling.Version = 3 batch := entity.Batch{ ID: "q/batch/1", Queue: request.Queue, @@ -527,9 +535,12 @@ func TestProcess_CreatingBatchDoesNotSuppressRequestCancellation(t *testing.T) { } requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(2), int32(3), entity.RequestStateCancelling).Return(nil) - requestStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(3), int32(4), entity.RequestStateCancelled).Return(nil) + gomock.InOrder( + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(request, nil), + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil), + requestStore.EXPECT().Get(gomock.Any(), request.ID).Return(cancelling, nil), + requestStore.EXPECT().Update(gomock.Any(), cancelling.WithState(entity.RequestStateCancelled), int32(3), int32(4)).Return(nil), + ) batchStore := storagemock.NewMockBatchStore(ctrl) store := storagemock.NewMockStorage(ctrl) @@ -601,7 +612,7 @@ func TestProcess_BatchIntentVersionMismatch_Retryable(t *testing.T) { reqStore := storagemock.NewMockRequestStore(ctrl) reqStore.EXPECT().Get(gomock.Any(), "q/1").Return(req, nil) - reqStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateCancelling).Return(nil) + reqStore.EXPECT().Update(gomock.Any(), req.WithState(entity.RequestStateCancelling), int32(2), int32(3)).Return(nil) batchStore := storagemock.NewMockBatchStore(ctrl) batchStore.EXPECT().Update(gomock.Any(), batch.WithState(entity.BatchStateCancelling), int32(1), int32(2)). diff --git a/submitqueue/orchestrator/controller/conclude/conclude_test.go b/submitqueue/orchestrator/controller/conclude/conclude_test.go index d2d8c0e7..00084f51 100644 --- a/submitqueue/orchestrator/controller/conclude/conclude_test.go +++ b/submitqueue/orchestrator/controller/conclude/conclude_test.go @@ -113,14 +113,16 @@ func TestController_Process(t *testing.T) { }, nil) mockRequestStore := storagemock.NewMockRequestStore(ctrl) - mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/1").Return(entity.Request{ + request1 := entity.Request{ ID: "test-queue/1", Version: 2, State: entity.RequestStateProcessing, - }, nil) - mockRequestStore.EXPECT().UpdateState(gomock.Any(), "test-queue/1", int32(2), int32(3), entity.RequestStateLanded).Return(nil) - mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/2").Return(entity.Request{ + } + mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/1").Return(request1, nil) + mockRequestStore.EXPECT().Update(gomock.Any(), request1.WithState(entity.RequestStateLanded), int32(2), int32(3)).Return(nil) + request2 := entity.Request{ ID: "test-queue/2", Version: 3, State: entity.RequestStateProcessing, - }, nil) - mockRequestStore.EXPECT().UpdateState(gomock.Any(), "test-queue/2", int32(3), int32(4), entity.RequestStateLanded).Return(nil) + } + mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/2").Return(request2, nil) + mockRequestStore.EXPECT().Update(gomock.Any(), request2.WithState(entity.RequestStateLanded), int32(3), int32(4)).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() @@ -149,10 +151,11 @@ func TestController_Process(t *testing.T) { }, nil) mockRequestStore := storagemock.NewMockRequestStore(ctrl) - mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/5").Return(entity.Request{ + request := entity.Request{ ID: "test-queue/5", Version: 1, State: entity.RequestStateProcessing, - }, nil) - mockRequestStore.EXPECT().UpdateState(gomock.Any(), "test-queue/5", int32(1), int32(2), entity.RequestStateError).Return(nil) + } + mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/5").Return(request, nil) + mockRequestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() @@ -181,10 +184,11 @@ func TestController_Process(t *testing.T) { }, nil) mockRequestStore := storagemock.NewMockRequestStore(ctrl) - mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/10").Return(entity.Request{ + request := entity.Request{ ID: "test-queue/10", Version: 4, State: entity.RequestStateProcessing, - }, nil) - mockRequestStore.EXPECT().UpdateState(gomock.Any(), "test-queue/10", int32(4), int32(5), entity.RequestStateCancelled).Return(nil) + } + mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/10").Return(request, nil) + mockRequestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateCancelled), int32(4), int32(5)).Return(nil) mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() @@ -368,10 +372,11 @@ func TestController_Process(t *testing.T) { }, nil) mockRequestStore := storagemock.NewMockRequestStore(ctrl) - mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/1").Return(entity.Request{ + request := entity.Request{ ID: "test-queue/1", Version: 2, State: entity.RequestStateProcessing, - }, nil) - mockRequestStore.EXPECT().UpdateState(gomock.Any(), "test-queue/1", int32(2), int32(3), entity.RequestStateLanded).Return(storage.ErrVersionMismatch) + } + mockRequestStore.EXPECT().Get(gomock.Any(), "test-queue/1").Return(request, nil) + mockRequestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateLanded), int32(2), int32(3)).Return(storage.ErrVersionMismatch) mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() diff --git a/submitqueue/orchestrator/controller/dlq/batch_test.go b/submitqueue/orchestrator/controller/dlq/batch_test.go index d07add99..907ecef4 100644 --- a/submitqueue/orchestrator/controller/dlq/batch_test.go +++ b/submitqueue/orchestrator/controller/dlq/batch_test.go @@ -51,10 +51,11 @@ func TestDLQBatchController_Process_FailsAndFansOut(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batch.WithState(entity.BatchStateFailed), int32(2), int32(3)).Return(nil) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 1, State: entity.RequestStateProcessing, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(1), int32(2), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(entity.RequestLog) error { return nil diff --git a/submitqueue/orchestrator/controller/dlq/buildsignal_test.go b/submitqueue/orchestrator/controller/dlq/buildsignal_test.go index 36ff636c..b45893e4 100644 --- a/submitqueue/orchestrator/controller/dlq/buildsignal_test.go +++ b/submitqueue/orchestrator/controller/dlq/buildsignal_test.go @@ -57,10 +57,11 @@ func TestDLQBuildSignalController_Process_FansOutToBatch(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batch.WithState(entity.BatchStateFailed), int32(3), int32(4)).Return(nil) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 1, State: entity.RequestStateProcessing, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(1), int32(2), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(entity.RequestLog) error { return nil diff --git a/submitqueue/orchestrator/controller/dlq/dlq_test.go b/submitqueue/orchestrator/controller/dlq/dlq_test.go index e7a4ae96..b095a51e 100644 --- a/submitqueue/orchestrator/controller/dlq/dlq_test.go +++ b/submitqueue/orchestrator/controller/dlq/dlq_test.go @@ -79,10 +79,11 @@ func TestFailRequest_CancellingTransitionsToError(t *testing.T) { ctrl := gomock.NewController(t) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 7, State: entity.RequestStateCancelling, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(7), int32(8), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(7), int32(8)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(l entity.RequestLog) error { assert.Equal(t, "q/1", l.RequestID) @@ -102,10 +103,11 @@ func TestFailRequest_TransitionsToError(t *testing.T) { ctrl := gomock.NewController(t) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 3, State: entity.RequestStateValidated, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(3), int32(4), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(3), int32(4)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(l entity.RequestLog) error { assert.Equal(t, "q/1", l.RequestID) @@ -128,10 +130,11 @@ func TestFailRequest_LogPublishErrorPropagates(t *testing.T) { ctrl := gomock.NewController(t) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 3, State: entity.RequestStateValidated, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(3), int32(4), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(3), int32(4)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(entity.RequestLog) error { return fmt.Errorf("publish boom") @@ -185,14 +188,16 @@ func TestFailBatch_TransitionsAndFansOut(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batch.WithState(entity.BatchStateFailed), int32(4), int32(5)).Return(nil) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request1 := entity.Request{ ID: "q/1", Version: 2, State: entity.RequestStateProcessing, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateError).Return(nil) - requestStore.EXPECT().Get(gomock.Any(), "q/2").Return(entity.Request{ + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request1, nil) + requestStore.EXPECT().Update(gomock.Any(), request1.WithState(entity.RequestStateError), int32(2), int32(3)).Return(nil) + request2 := entity.Request{ ID: "q/2", Version: 1, State: entity.RequestStateProcessing, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/2", int32(1), int32(2), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/2").Return(request2, nil) + requestStore.EXPECT().Update(gomock.Any(), request2.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) registry := newTestLogRegistry(t, ctrl, 2, func(entity.RequestLog) error { return nil @@ -217,10 +222,11 @@ func TestFailBatch_FailedFansOutForRepair(t *testing.T) { // no batchStore.Update expected requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 2, State: entity.RequestStateProcessing, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(2), int32(3), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(2), int32(3)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(entity.RequestLog) error { return nil @@ -270,10 +276,11 @@ func TestFailBatch_CancellingTransitionsToFailed(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batch.WithState(entity.BatchStateFailed), int32(6), int32(7)).Return(nil) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 3, State: entity.RequestStateCancelling, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(3), int32(4), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(3), int32(4)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(entity.RequestLog) error { return nil diff --git a/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go b/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go index 87a35514..6d22607b 100644 --- a/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go +++ b/submitqueue/orchestrator/controller/dlq/mergeconflictsignal_test.go @@ -44,10 +44,11 @@ func TestDLQMergeConflictSignalController_Process_ReconcilesRequest(t *testing.T ctrl := gomock.NewController(t) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 1, State: entity.RequestStateProcessing, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(1), int32(2), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(entity.RequestLog) error { return nil diff --git a/submitqueue/orchestrator/controller/dlq/mergesignal_test.go b/submitqueue/orchestrator/controller/dlq/mergesignal_test.go index d4016f0e..5b51dc53 100644 --- a/submitqueue/orchestrator/controller/dlq/mergesignal_test.go +++ b/submitqueue/orchestrator/controller/dlq/mergesignal_test.go @@ -54,10 +54,11 @@ func TestDLQMergeSignalController_Process_ReconcilesBatch(t *testing.T) { batchStore.EXPECT().Update(gomock.Any(), batch.WithState(entity.BatchStateFailed), int32(2), int32(3)).Return(nil) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 1, State: entity.RequestStateProcessing, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(1), int32(2), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(entity.RequestLog) error { return nil diff --git a/submitqueue/orchestrator/controller/dlq/request_test.go b/submitqueue/orchestrator/controller/dlq/request_test.go index 17fc0eab..7efdc852 100644 --- a/submitqueue/orchestrator/controller/dlq/request_test.go +++ b/submitqueue/orchestrator/controller/dlq/request_test.go @@ -45,10 +45,11 @@ func TestDLQRequestController_Process_LandRequestPayload(t *testing.T) { ctrl := gomock.NewController(t) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(entity.Request{ + request := entity.Request{ ID: "q/1", Version: 1, State: entity.RequestStateStarted, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/1", int32(1), int32(2), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/1").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(entity.RequestLog) error { return nil @@ -70,10 +71,11 @@ func TestDLQRequestController_Process_CancelRequestPayload(t *testing.T) { ctrl := gomock.NewController(t) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/7").Return(entity.Request{ + request := entity.Request{ ID: "q/7", Version: 2, State: entity.RequestStateBatched, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/7", int32(2), int32(3), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/7").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(2), int32(3)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(entity.RequestLog) error { return nil @@ -95,10 +97,11 @@ func TestDLQRequestController_Process_RequestIDPayload(t *testing.T) { ctrl := gomock.NewController(t) requestStore := storagemock.NewMockRequestStore(ctrl) - requestStore.EXPECT().Get(gomock.Any(), "q/3").Return(entity.Request{ + request := entity.Request{ ID: "q/3", Version: 1, State: entity.RequestStateValidated, - }, nil) - requestStore.EXPECT().UpdateState(gomock.Any(), "q/3", int32(1), int32(2), entity.RequestStateError).Return(nil) + } + requestStore.EXPECT().Get(gomock.Any(), "q/3").Return(request, nil) + requestStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) registry := newTestLogRegistry(t, ctrl, 1, func(log entity.RequestLog) error { assert.Equal(t, "boom", log.LastError) diff --git a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go index 54b192ae..d5389145 100644 --- a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go +++ b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal.go @@ -127,12 +127,13 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) er // Advance the request to Validated now that the merge-conflict check passed. newVersion := request.Version + 1 - if err := c.store.GetRequestStore().UpdateState(ctx, request.ID, request.Version, newVersion, entity.RequestStateValidated); err != nil { + updatedRequest := request.WithState(entity.RequestStateValidated) + if err := c.store.GetRequestStore().Update(ctx, updatedRequest, request.Version, newVersion); err != nil { metrics.NamedCounter(c.metricsScope, opName, "state_errors", 1) return fmt.Errorf("failed to update request %s state to validated: %w", request.ID, err) } - request.Version = newVersion - request.State = entity.RequestStateValidated + updatedRequest.Version = newVersion + request = updatedRequest logEntry := entity.NewRequestLog(request.ID, entity.RequestStatusValidated, request.Version, "", nil) if err := corerequest.PublishLog(ctx, c.registry, logEntry, request.ID); err != nil { @@ -174,11 +175,12 @@ func (c *Controller) failRequest(ctx context.Context, request entity.Request, re return nil default: newVersion := request.Version + 1 - if err := c.store.GetRequestStore().UpdateState(ctx, request.ID, request.Version, newVersion, entity.RequestStateError); err != nil { + updatedRequest := request.WithState(entity.RequestStateError) + if err := c.store.GetRequestStore().Update(ctx, updatedRequest, request.Version, newVersion); err != nil { return fmt.Errorf("failed to update request %s state to error: %w", request.ID, err) } - request.Version = newVersion - request.State = entity.RequestStateError + updatedRequest.Version = newVersion + request = updatedRequest } logEntry := entity.NewRequestLog(request.ID, entity.RequestStatusError, request.Version, reason, nil) diff --git a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal_test.go b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal_test.go index 56f289d3..e918f2a7 100644 --- a/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal_test.go +++ b/submitqueue/orchestrator/controller/mergeconflictsignal/mergeconflictsignal_test.go @@ -16,6 +16,7 @@ package mergeconflictsignal import ( "context" + "fmt" "testing" "github.com/stretchr/testify/assert" @@ -55,9 +56,9 @@ func TestProcess_MergeablePublishesToBatch(t *testing.T) { ctrl := gomock.NewController(t) reqStore := storagemock.NewMockRequestStore(ctrl) - reqStore.EXPECT().Get(gomock.Any(), testRequestID).Return( - entity.Request{ID: testRequestID, Queue: testQueue, State: entity.RequestStateStarted, Version: 1}, nil) - reqStore.EXPECT().UpdateState(gomock.Any(), testRequestID, int32(1), int32(2), entity.RequestStateValidated).Return(nil) + request := entity.Request{ID: testRequestID, Queue: testQueue, State: entity.RequestStateStarted, Version: 1} + reqStore.EXPECT().Get(gomock.Any(), testRequestID).Return(request, nil) + reqStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateValidated), int32(1), int32(2)).Return(nil) store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() @@ -109,10 +110,10 @@ func TestProcess_NotMergeableMarksRequestError(t *testing.T) { ctrl := gomock.NewController(t) reqStore := storagemock.NewMockRequestStore(ctrl) - reqStore.EXPECT().Get(gomock.Any(), testRequestID).Return( - entity.Request{ID: testRequestID, Queue: testQueue, State: entity.RequestStateStarted, Version: 1}, nil) + request := entity.Request{ID: testRequestID, Queue: testQueue, State: entity.RequestStateStarted, Version: 1} + reqStore.EXPECT().Get(gomock.Any(), testRequestID).Return(request, nil) // The request is driven to terminal Error inline (version 1 -> 2). - reqStore.EXPECT().UpdateState(gomock.Any(), testRequestID, int32(1), int32(2), entity.RequestStateError).Return(nil) + reqStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetRequestStore().Return(reqStore).AnyTimes() @@ -156,6 +157,25 @@ func TestProcess_NotMergeableMarksRequestError(t *testing.T) { assert.Equal(t, "conflict in foo.go", logEntry.LastError) } +func TestFailRequest_UpdateFailureLeavesRequestUnchanged(t *testing.T) { + ctrl := gomock.NewController(t) + + request := entity.Request{ID: testRequestID, Queue: testQueue, State: entity.RequestStateStarted, Version: 1} + reqStore := storagemock.NewMockRequestStore(ctrl) + reqStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(fmt.Errorf("db down")) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetRequestStore().Return(reqStore) + + controller := NewController(zaptest.NewLogger(t).Sugar(), tally.NoopScope, store, consumer.TopicRegistry{}, + runwaymq.TopicKeyMergeConflictCheckSignal, "orchestrator-mergeconflictsignal") + + err := controller.failRequest(context.Background(), request, "conflict") + require.Error(t, err) + assert.Equal(t, entity.RequestStateStarted, request.State) + assert.Equal(t, int32(1), request.Version) +} + func TestProcess_HaltedRequestSkips(t *testing.T) { ctrl := gomock.NewController(t) diff --git a/submitqueue/orchestrator/controller/validate/validate_test.go b/submitqueue/orchestrator/controller/validate/validate_test.go index 615638a0..1f304626 100644 --- a/submitqueue/orchestrator/controller/validate/validate_test.go +++ b/submitqueue/orchestrator/controller/validate/validate_test.go @@ -471,7 +471,7 @@ func TestController_Process_DuplicateDetection(t *testing.T) { if tt.wantRejected { // A detected duplicate is terminated (Started → Error) rather than // dead-lettered, then the delivery is acked. - mockReqStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(1), int32(2), entity.RequestStateError).Return(nil) + mockReqStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) } store := storagemock.NewMockStorage(ctrl) store.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() @@ -637,7 +637,7 @@ func TestController_Process_CustomValidatorFails(t *testing.T) { store, mockReqStore := newMockStorage(ctrl, request) store.EXPECT().GetChangeStore().Return(newMockChangeStore(ctrl)).AnyTimes() // A validator rejection terminates the request (Started → Error) rather than dead-lettering it. - mockReqStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(1), int32(2), entity.RequestStateError).Return(nil) + mockReqStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) logger := zaptest.NewLogger(t).Sugar() @@ -700,7 +700,7 @@ func TestController_Process_CustomValidatorFailure_TerminationPublishFails(t *te } store, mockReqStore := newMockStorage(ctrl, request) store.EXPECT().GetChangeStore().Return(newMockChangeStore(ctrl)).AnyTimes() - mockReqStore.EXPECT().UpdateState(gomock.Any(), request.ID, int32(1), int32(2), entity.RequestStateError).Return(nil) + mockReqStore.EXPECT().Update(gomock.Any(), request.WithState(entity.RequestStateError), int32(1), int32(2)).Return(nil) mockPub := queuemock.NewMockPublisher(ctrl) mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).Return(fmt.Errorf("publish boom")).AnyTimes() diff --git a/test/integration/submitqueue/extension/storage/suite.go b/test/integration/submitqueue/extension/storage/suite.go index 24616b12..2bcd3dce 100644 --- a/test/integration/submitqueue/extension/storage/suite.go +++ b/test/integration/submitqueue/extension/storage/suite.go @@ -16,7 +16,9 @@ package storage import ( "context" + "fmt" "sort" + "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -124,14 +126,15 @@ func (s *StorageContractSuite) TestStorage_CreateAndGet_StackedPRs() { assert.Equal(t, request.LandStrategy, retrieved.LandStrategy) } -// TestStorage_UpdateState tests updating request state -func (s *StorageContractSuite) TestStorage_UpdateState() { +// TestStorage_Update tests replacing all non-key request fields. +func (s *StorageContractSuite) TestStorage_Update() { t := s.T() ctx := s.ctx request := entity.Request{ ID: "test/update", Queue: "test-queue", + Change: change.Change{URIs: []string{"github://github.example.com/uber/monorepo/pull/1/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}}, State: entity.RequestStateStarted, LandStrategy: mergestrategy.MergeStrategyMerge, Version: 1, @@ -141,15 +144,19 @@ func (s *StorageContractSuite) TestStorage_UpdateState() { err := s.storage.GetRequestStore().Create(ctx, request) require.NoError(t, err) - // Update state - err = s.storage.GetRequestStore().UpdateState(ctx, request.ID, request.Version, request.Version+1, entity.RequestStateProcessing) - require.NoError(t, err, "failed to update request state") + updated := request + updated.Queue = "updated-queue" + updated.Change.URIs = []string{"github://github.example.com/uber/monorepo/pull/2/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"} + updated.LandStrategy = mergestrategy.MergeStrategySquashRebase + updated.State = entity.RequestStateProcessing + err = s.storage.GetRequestStore().Update(ctx, updated, request.Version, request.Version+1) + require.NoError(t, err, "failed to update request") // Verify update retrieved, err := s.storage.GetRequestStore().Get(ctx, request.ID) require.NoError(t, err) - assert.Equal(t, entity.RequestStateProcessing, retrieved.State) - assert.Equal(t, int32(2), retrieved.Version, "version should increment after update") + updated.Version = request.Version + 1 + assert.Equal(t, updated, retrieved) } // TestStorage_OptimisticLocking tests version-based optimistic locking @@ -169,20 +176,66 @@ func (s *StorageContractSuite) TestStorage_OptimisticLocking() { err := s.storage.GetRequestStore().Create(ctx, request) require.NoError(t, err) - // Update with correct version - err = s.storage.GetRequestStore().UpdateState(ctx, request.ID, 1, 2, entity.RequestStateProcessing) + // Update with correct version. + updated := request + updated.Queue = "updated-queue" + updated.Change.URIs = []string{"github://github.example.com/uber/monorepo/pull/2/bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"} + updated.LandStrategy = mergestrategy.MergeStrategySquashRebase + updated.State = entity.RequestStateProcessing + err = s.storage.GetRequestStore().Update(ctx, updated, 1, 2) require.NoError(t, err, "update with correct version should succeed") - // Try to update with stale version (should fail) - err = s.storage.GetRequestStore().UpdateState(ctx, request.ID, 1, 2, entity.RequestStateLanded) + // Try to replace every field with a stale version. + stale := request + stale.Queue = "stale-queue" + stale.Change.URIs = []string{"github://github.example.com/uber/monorepo/pull/3/cccccccccccccccccccccccccccccccccccccccc"} + stale.LandStrategy = mergestrategy.MergeStrategyRebase + stale.State = entity.RequestStateLanded + err = s.storage.GetRequestStore().Update(ctx, stale, 1, 3) assert.Error(t, err, "update with stale version should fail") assert.ErrorIs(t, err, storage.ErrVersionMismatch, "should return ErrVersionMismatch") - // Verify state wasn't changed by stale update + // Verify no field was changed by the stale update. retrieved, err := s.storage.GetRequestStore().Get(ctx, request.ID) require.NoError(t, err) - assert.Equal(t, entity.RequestStateProcessing, retrieved.State, "stale update should not modify state") - assert.Equal(t, int32(2), retrieved.Version) + updated.Version = 2 + assert.Equal(t, updated, retrieved) +} + +// TestStorage_UpdateChangeURIs tests nil and empty URI replacement semantics. +func (s *StorageContractSuite) TestStorage_UpdateChangeURIs() { + t := s.T() + ctx := s.ctx + + tests := []struct { + name string + uris []string + }{ + {name: "nil", uris: nil}, + {name: "empty", uris: []string{}}, + } + + for i, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + request := entity.Request{ + ID: fmt.Sprintf("test/update-change-uris-%d", i), + Queue: "test-queue", + Change: change.Change{URIs: []string{"github://github.example.com/uber/monorepo/pull/1/aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}}, + State: entity.RequestStateStarted, + LandStrategy: mergestrategy.MergeStrategyMerge, + Version: 1, + } + require.NoError(t, s.storage.GetRequestStore().Create(ctx, request)) + + updated := request + updated.Change.URIs = tt.uris + require.NoError(t, s.storage.GetRequestStore().Update(ctx, updated, request.Version, request.Version+1)) + + retrieved, err := s.storage.GetRequestStore().Get(ctx, request.ID) + require.NoError(t, err) + assert.Equal(t, tt.uris, retrieved.Change.URIs) + }) + } } // TestStorage_BatchDependentUpdate verifies full updates, collection encoding, and optimistic locking.