fix(orchestrator): fix score reconciliation bugs

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",

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{

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))
+}