fix(orchestrator): fix score reconciliation bugs

submitqueue/orchestrator/controller/score/BUILD.bazel

@@ -32,6 +32,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,6 +16,7 @@ package score
 
 import (
 	"context"
+	"errors"
 	"fmt"
 
 	"github.com/uber-go/tally"
@@ -104,53 +105,90 @@ 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.
-	if batch.State == entity.BatchStateCancelling {
-		c.metricsScope.Counter("skipped_cancelling").Inc(1)
-		c.logger.Infow("skipping score for cancelling batch",
-			"batch_id", batch.ID,
-		)
-		return nil
+	var batchScore float64
+	for {
+		// 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",
+				"batch_id", batch.ID,
+			)
+			return nil
+		}
+
+		// 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",
+				"batch_id", batch.ID,
+				"state", string(batch.State),
+			)
+			return nil
+		}
+
+		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) {
+				c.metricsScope.Counter("version_conflicts").Inc(1)
+				batch, err = c.store.GetBatchStore().Get(ctx, batch.ID)
+				if err != nil {
+					metrics.NamedCounter(c.metricsScope, opName, "storage_errors", 1)
+					return fmt.Errorf("failed to reload batch %s after version conflict: %w", bid.ID, err)
+				}
+				continue
+			}
+			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,
+			)
+
+		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
+			c.metricsScope.Counter("replayed_scored").Inc(1)
+
+		case entity.BatchStateSpeculating, entity.BatchStateMerging:
+			// A downstream controller consumed the handoff, so this score
+			// delivery is stale and must not regress the batch to Scored.
+			c.metricsScope.Counter("skipped_downstream").Inc(1)
+			return nil
+
+		default:
+			c.metricsScope.Counter("unexpected_state").Inc(1)
+			return fmt.Errorf("unexpected batch state %q for batch %s", batch.State, batch.ID)
+		}
+
+		break
 	}
 
-	// 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.
-	if batch.State.IsTerminal() {
-		c.metricsScope.Counter("skipped_terminal").Inc(1)
-		c.logger.Infow("skipping score for terminal batch",
-			"batch_id", batch.ID,
-			"state", string(batch.State),
-		)
-		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)
-	}
-
-	// 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
-
-	c.logger.Infow("scored batch",
-		"batch_id", batch.ID,
-		"score", batchScore,
-	)
-
 	// Publish request log entries for all requests in the batch
 	if err := corerequest.PublishBatchLogs(ctx, c.registry, batch.Contains, entity.RequestStatusScored, map[string]string{
 		"batch_id": batch.ID,

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"
@@ -433,3 +434,117 @@ 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_CASConflictReloadsScoredAndReplaysFanout(t *testing.T) {
+	ctrl := gomock.NewController(t)
+
+	created := testBatch()
+	const losingScore = 0.25
+
+	scored := created
+	scored.State = entity.BatchStateScored
+	scored.Version = created.Version + 1
+	scored.Score = 0.9
+
+	batchStore := storagemock.NewMockBatchStore(ctrl)
+	gomock.InOrder(
+		batchStore.EXPECT().Get(gomock.Any(), created.ID).Return(created, nil),
+		batchStore.EXPECT().UpdateScoreAndState(
+			gomock.Any(),
+			created.ID,
+			created.Version,
+			created.Version+1,
+			losingScore,
+			entity.BatchStateScored,
+		).Return(storage.ErrVersionMismatch),
+		batchStore.EXPECT().Get(gomock.Any(), created.ID).Return(scored, nil),
+	)
+
+	store := storagemock.NewMockStorage(ctrl)
+	store.EXPECT().GetBatchStore().Return(batchStore).AnyTimes()
+
+	mockScorer := scorermock.NewMockScorer(ctrl)
+	mockScorer.EXPECT().Score(gomock.Any(), created).Return(losingScore, nil)
+
+	scorerFactory := scorermock.NewMockFactory(ctrl)
+	scorerFactory.EXPECT().For(gomock.Any()).Return(mockScorer, nil)
+
+	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,
+		scorerFactory,
+		registry,
+		topickey.TopicKeyScore,
+		"orchestrator-score",
+	)
+
+	msg := entityqueue.NewMessage(created.ID, batchIDPayload(t, created.ID), created.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))
+}