From d002c4dc3eae74a2f4ff752d11a3908706e6ac5e Mon Sep 17 00:00:00 2001 From: Albert Wu Date: Mon, 20 Jul 2026 16:07:35 -0700 Subject: [PATCH] fix(orchestrator): reconcile score redelivery --- .../orchestrator/controller/score/BUILD.bazel | 2 + .../orchestrator/controller/score/score.go | 87 ++++++---- .../controller/score/score_test.go | 152 +++++++++++++++++- 3 files changed, 208 insertions(+), 33 deletions(-) diff --git a/submitqueue/orchestrator/controller/score/BUILD.bazel b/submitqueue/orchestrator/controller/score/BUILD.bazel index f244c286..bcfe9deb 100644 --- a/submitqueue/orchestrator/controller/score/BUILD.bazel +++ b/submitqueue/orchestrator/controller/score/BUILD.bazel @@ -8,6 +8,7 @@ go_library( deps = [ "//platform/base/messagequeue:go_default_library", "//platform/consumer:go_default_library", + "//platform/errs:go_default_library", "//platform/metrics:go_default_library", "//submitqueue/core/request:go_default_library", "//submitqueue/core/topickey:go_default_library", @@ -32,6 +33,7 @@ go_test( "//submitqueue/core/topickey:go_default_library", "//submitqueue/entity:go_default_library", "//submitqueue/extension/scorer/mock:go_default_library", + "//submitqueue/extension/storage:go_default_library", "//submitqueue/extension/storage/mock:go_default_library", "@com_github_stretchr_testify//assert:go_default_library", "@com_github_stretchr_testify//require:go_default_library", diff --git a/submitqueue/orchestrator/controller/score/score.go b/submitqueue/orchestrator/controller/score/score.go index 257ae86d..2cc53e9d 100644 --- a/submitqueue/orchestrator/controller/score/score.go +++ b/submitqueue/orchestrator/controller/score/score.go @@ -16,11 +16,13 @@ package score import ( "context" + "errors" "fmt" "github.com/uber-go/tally" entityqueue "github.com/uber/submitqueue/platform/base/messagequeue" "github.com/uber/submitqueue/platform/consumer" + "github.com/uber/submitqueue/platform/errs" "github.com/uber/submitqueue/platform/metrics" corerequest "github.com/uber/submitqueue/submitqueue/core/request" "github.com/uber/submitqueue/submitqueue/core/topickey" @@ -104,10 +106,10 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r "partition_key", msg.PartitionKey, ) - // Short-circuit when the batch is in BatchStateCancelling — the cancel - // controller has handed the batch off to speculate, which owns the terminal - // write to Cancelled and the downstream dependent / conclude publishes. We - // must not race it to conclude (conclude requires terminal). Silently ack. + var batchScore float64 + // Short-circuit when the batch is in BatchStateCancelling. The cancel + // controller has handed the batch off to speculate, which owns the + // terminal write and downstream fanout. if batch.State == entity.BatchStateCancelling { c.metricsScope.Counter("skipped_cancelling").Inc(1) c.logger.Infow("skipping score for cancelling batch", @@ -116,12 +118,8 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r return nil } - // Short-circuit if the batch is already terminal. Score never writes a - // terminal state, so it owns no recovery here: whichever controller wrote - // the terminal state (speculate.cancelBatch / failOnDependency, or merge) - // already published to conclude, and speculate's terminal self-heal - // republishes conclude on every redelivery of a terminal batch. Silently - // ack — same pattern as build / buildsignal on halted. + // Score owns no terminal-state recovery. The controller that wrote the + // terminal state owns its remaining fanout. if batch.State.IsTerminal() { c.metricsScope.Counter("skipped_terminal").Inc(1) c.logger.Infow("skipping score for terminal batch", @@ -131,25 +129,60 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r return nil } - // Score the batch. The scorer resolves the batch's changes itself. - batchScore, err := c.scoreBatch(ctx, batch) - if err != nil { - metrics.NamedCounter(c.metricsScope, opName, "scorer_errors", 1) - return fmt.Errorf("failed to score batch %s: %w", batch.ID, err) - } + switch batch.State { + case entity.BatchStateCreated: + // Score the batch. The scorer resolves the batch's changes itself. + batchScore, err = c.scoreBatch(ctx, batch) + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "scorer_errors", 1) + return fmt.Errorf("failed to score batch %s: %w", batch.ID, err) + } + + newVersion := batch.Version + 1 + err = c.store.GetBatchStore().UpdateScoreAndState( + ctx, + batch.ID, + batch.Version, + newVersion, + batchScore, + entity.BatchStateScored, + ) + if errors.Is(err, storage.ErrVersionMismatch) { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return errs.NewRetryableError( + fmt.Errorf("failed to update score for batch %s: %w", batch.ID, err), + ) + } + if err != nil { + metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) + return fmt.Errorf("failed to update score for batch %s: %w", batch.ID, err) + } + + batch.Version = newVersion + batch.Score = batchScore + batch.State = entity.BatchStateScored + c.logger.Infow("scored batch", + "batch_id", batch.ID, + "score", batchScore, + ) - // Atomically update score and state to "scored" in the database - newVersion := batch.Version + 1 - if err := c.store.GetBatchStore().UpdateScoreAndState(ctx, batch.ID, batch.Version, newVersion, batchScore, entity.BatchStateScored); err != nil { - metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1) - return fmt.Errorf("failed to update score for batch %s: %w", batch.ID, err) - } - batch.Version = newVersion + case entity.BatchStateScored: + // The durable transition already happened, but its fanout may be + // incomplete. Preserve the committed score and replay all outputs. + batchScore = batch.Score + + case entity.BatchStateSpeculating, entity.BatchStateMerging: + // Normal under at-least-once delivery: a prior score attempt may have + // published to speculate before its acknowledgement was recorded. + // Downstream processing has already advanced the batch, so this stale + // delivery is satisfied and must not regress the batch to Scored. + c.metricsScope.Counter("skipped_downstream").Inc(1) + return nil - c.logger.Infow("scored batch", - "batch_id", batch.ID, - "score", batchScore, - ) + default: + c.metricsScope.Counter("unexpected_state").Inc(1) + return fmt.Errorf("unexpected batch state %q for batch %s", batch.State, batch.ID) + } // Publish request log entries for all requests in the batch if err := corerequest.PublishBatchLogs(ctx, c.registry, batch.Contains, entity.RequestStatusScored, map[string]string{ diff --git a/submitqueue/orchestrator/controller/score/score_test.go b/submitqueue/orchestrator/controller/score/score_test.go index 87325dcc..776f6e22 100644 --- a/submitqueue/orchestrator/controller/score/score_test.go +++ b/submitqueue/orchestrator/controller/score/score_test.go @@ -30,6 +30,7 @@ import ( "github.com/uber/submitqueue/submitqueue/core/topickey" "github.com/uber/submitqueue/submitqueue/entity" scorermock "github.com/uber/submitqueue/submitqueue/extension/scorer/mock" + "github.com/uber/submitqueue/submitqueue/extension/storage" storagemock "github.com/uber/submitqueue/submitqueue/extension/storage/mock" "go.uber.org/mock/gomock" "go.uber.org/zap/zaptest" @@ -329,7 +330,7 @@ func TestController_InterfaceImplementation(t *testing.T) { // A batch already in a terminal state (e.g. cancelled while the score message // was in flight) must be short-circuited: no scoring, no UpdateScoreAndState, -// and no fan-out — score owns no terminal write, so it owns no recovery; the +// and no fan-out. Score owns no terminal write, so it owns no recovery; the // controller that wrote the terminal state already published to conclude, and // speculate's terminal self-heal republishes on redelivery. func TestController_Process_TerminalShortCircuit(t *testing.T) { @@ -351,13 +352,13 @@ func TestController_Process_TerminalShortCircuit(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() - // Scorer with no EXPECTs — must not be called. + // The scorer has no expectations and must not be called. mockScorer := scorermock.NewMockScorer(ctrl) logger := zaptest.NewLogger(t).Sugar() scope := tally.NoopScope - // Publisher with no EXPECTs — must not be called (no fan-out on terminal). + // The publisher has no expectations and must not be called for a terminal batch. mockPub := queuemock.NewMockPublisher(ctrl) mockQ := queuemock.NewMockQueue(ctrl) @@ -389,7 +390,7 @@ func TestController_Process_TerminalShortCircuit(t *testing.T) { // A batch in BatchStateCancelling must be silently acked: no scoring, no // UpdateScoreAndState, and crucially NO conclude publish (speculate owns the // terminal write to Cancelled and the downstream dependent / conclude -// publishes — conclude would also error on a non-terminal Cancelling batch). +// publishes because conclude would also error on a non-terminal Cancelling batch). func TestController_Process_CancellingShortCircuit(t *testing.T) { ctrl := gomock.NewController(t) @@ -402,13 +403,13 @@ func TestController_Process_CancellingShortCircuit(t *testing.T) { mockStorage := storagemock.NewMockStorage(ctrl) mockStorage.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() - // Scorer with no EXPECTs — must not be called. + // The scorer has no expectations and must not be called. mockScorer := scorermock.NewMockScorer(ctrl) logger := zaptest.NewLogger(t).Sugar() scope := tally.NoopScope - // Publisher with no EXPECTs — must not be called (no fan-out for Cancelling). + // The publisher has no expectations and must not be called for a cancelling batch. mockPub := queuemock.NewMockPublisher(ctrl) mockQ := queuemock.NewMockQueue(ctrl) mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() @@ -433,3 +434,142 @@ func TestController_Process_CancellingShortCircuit(t *testing.T) { require.NoError(t, controller.Process(context.Background(), delivery)) } + +func TestController_Process_DownstreamStateDoesNotRegress(t *testing.T) { + for _, state := range []entity.BatchState{ + entity.BatchStateSpeculating, + entity.BatchStateMerging, + } { + t.Run(string(state), func(t *testing.T) { + ctrl := gomock.NewController(t) + + batch := testBatch() + batch.State = state + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + // No scorer, state update, or publisher calls are expected. A score + // delivery observed after downstream progress is only a stale poke. + scorerFactory := scorermock.NewMockFactory(ctrl) + controller := NewController( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + store, + scorerFactory, + consumer.TopicRegistry{}, + topickey.TopicKeyScore, + "orchestrator-score", + ) + + msg := entityqueue.NewMessage(batch.ID, batchIDPayload(t, batch.ID), batch.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(1).AnyTimes() + + require.NoError(t, controller.Process(context.Background(), delivery)) + }) + } +} + +func TestController_Process_VersionConflictReturnsForRedelivery(t *testing.T) { + ctrl := gomock.NewController(t) + + batch := testBatch() + const losingScore = 0.25 + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + batchStore.EXPECT().UpdateScoreAndState( + gomock.Any(), + batch.ID, + batch.Version, + batch.Version+1, + losingScore, + entity.BatchStateScored, + ).Return(storage.ErrVersionMismatch) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + mockScorer := scorermock.NewMockScorer(ctrl) + mockScorer.EXPECT().Score(gomock.Any(), batch).Return(losingScore, nil) + + scorerFactory := scorermock.NewMockFactory(ctrl) + scorerFactory.EXPECT().For(gomock.Any()).Return(mockScorer, nil) + + controller := NewController( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + store, + scorerFactory, + consumer.TopicRegistry{}, + topickey.TopicKeyScore, + "orchestrator-score", + ) + + msg := entityqueue.NewMessage(batch.ID, batchIDPayload(t, batch.ID), batch.Queue, nil) + delivery := queuemock.NewMockDelivery(ctrl) + delivery.EXPECT().Message().Return(msg).AnyTimes() + delivery.EXPECT().Attempt().Return(1).AnyTimes() + + err := controller.Process(context.Background(), delivery) + require.ErrorIs(t, err, storage.ErrVersionMismatch) + assert.True(t, errs.IsRetryable(err)) +} + +func TestController_Process_ScoredReplaysFanout(t *testing.T) { + ctrl := gomock.NewController(t) + + batch := testBatch() + batch.State = entity.BatchStateScored + batch.Score = 0.9 + + batchStore := storagemock.NewMockBatchStore(ctrl) + batchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes() + + mockPub := queuemock.NewMockPublisher(ctrl) + gomock.InOrder( + mockPub.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).DoAndReturn( + func(_ context.Context, _ string, msg entityqueue.Message) error { + logEntry, err := entity.RequestLogFromBytes(msg.Payload) + require.NoError(t, err) + assert.Equal(t, "0.9000", logEntry.Metadata["score"]) + return nil + }, + ), + mockPub.EXPECT().Publish(gomock.Any(), "speculate", gomock.Any()).Return(nil), + ) + + mockQ := queuemock.NewMockQueue(ctrl) + mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() + + registry, err := consumer.NewTopicRegistry([]consumer.TopicConfig{ + {Key: topickey.TopicKeyLog, Name: "log", Queue: mockQ}, + {Key: topickey.TopicKeySpeculate, Name: "speculate", Queue: mockQ}, + }) + require.NoError(t, err) + + controller := NewController( + zaptest.NewLogger(t).Sugar(), + tally.NoopScope, + store, + scorermock.NewMockFactory(ctrl), + registry, + topickey.TopicKeyScore, + "orchestrator-score", + ) + + msg := entityqueue.NewMessage(batch.ID, batchIDPayload(t, batch.ID), batch.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)) +}