From 5a36a0329e87bd904d2aa667340431087a124efa Mon Sep 17 00:00:00 2001 From: Albert Wu Date: Mon, 6 Jul 2026 14:07:03 -0700 Subject: [PATCH] feat: publish batching and scoring request logs Emit gateway-visible progress logs when requests enter active batching and scoring, while keeping tests focused on existing controller behavior. Co-authored-by: Cursor --- submitqueue/entity/request_log.go | 3 + .../orchestrator/controller/batch/batch.go | 6 ++ .../controller/batch/batch_test.go | 63 +++++++++++++---- .../orchestrator/controller/score/score.go | 7 ++ .../controller/score/score_test.go | 68 ++++++++++++++++++- 5 files changed, 132 insertions(+), 15 deletions(-) diff --git a/submitqueue/entity/request_log.go b/submitqueue/entity/request_log.go index 6dee8839..3b3ce712 100644 --- a/submitqueue/entity/request_log.go +++ b/submitqueue/entity/request_log.go @@ -49,6 +49,9 @@ const ( // RequestStatusBatched indicates that the request has been included in a new batch and will be sent to speculation. RequestStatusBatched RequestStatus = "batched" + // RequestStatusScoring indicates that the batch containing the request is being scored for build success probability. + RequestStatusScoring RequestStatus = "scoring" + // RequestStatusScored indicates that the batch containing the request has been scored for build success probability. RequestStatusScored RequestStatus = "scored" diff --git a/submitqueue/orchestrator/controller/batch/batch.go b/submitqueue/orchestrator/controller/batch/batch.go index 9bb13e85..23a47e8c 100644 --- a/submitqueue/orchestrator/controller/batch/batch.go +++ b/submitqueue/orchestrator/controller/batch/batch.go @@ -118,6 +118,12 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r return nil } + batchingLog := entity.NewRequestLog(request.ID, entity.RequestStatusBatching, 0, "", nil) + if err := corerequest.PublishLog(ctx, c.registry, batchingLog, request.ID); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "request_log_errors", 1) + return fmt.Errorf("failed to publish batching request log for request %s: %w", request.ID, err) + } + // TODO: if capacity is full, wait here for other requests to accumulate to batch them together, or include a request into an existing batch if it's not too late. // Generate a globally unique batch ID. diff --git a/submitqueue/orchestrator/controller/batch/batch_test.go b/submitqueue/orchestrator/controller/batch/batch_test.go index f1651c5d..82aec161 100644 --- a/submitqueue/orchestrator/controller/batch/batch_test.go +++ b/submitqueue/orchestrator/controller/batch/batch_test.go @@ -48,6 +48,18 @@ func requestIDPayload(t *testing.T, id string) []byte { return payload } +func requestLogsFromMessages(t *testing.T, msgs []entityqueue.Message) []entity.RequestLog { + t.Helper() + + logs := make([]entity.RequestLog, 0, len(msgs)) + for _, msg := range msgs { + logEntry, err := entity.RequestLogFromBytes(msg.Payload) + require.NoError(t, err) + logs = append(logs, logEntry) + } + return logs +} + // newSequentialCounter returns a mock counter that returns incrementing values starting at 1. func newSequentialCounter(ctrl *gomock.Controller) *countermock.MockCounter { var seq int64 @@ -157,10 +169,10 @@ func TestController_Process_Success(t *testing.T) { require.NoError(t, err) } -// TestController_Process_PublishesBatchedLog asserts the controller emits a -// "batched" request log carrying the request ID, the post-CAS request version, -// and the batch ID it was placed into. -func TestController_Process_PublishesBatchedLog(t *testing.T) { +// TestController_Process_PublishesBatchingAndBatchedLogs asserts the controller +// emits a progress log when active batching starts, then the versioned "batched" +// log once the request has been claimed into a persisted batch. +func TestController_Process_PublishesBatchingAndBatchedLogs(t *testing.T) { ctrl := gomock.NewController(t) request := testRequest() @@ -217,13 +229,24 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) { require.NoError(t, controller.Process(context.Background(), delivery)) - require.Len(t, logMsgs, 1) - logEntry, err := entity.RequestLogFromBytes(logMsgs[0].Payload) - require.NoError(t, err) - assert.Equal(t, request.ID, logEntry.RequestID) - assert.Equal(t, entity.RequestStatusBatched, logEntry.Status) - assert.Equal(t, request.Version+1, logEntry.RequestVersion) - assert.Equal(t, "test-queue/batch/1", logEntry.Metadata["batch_id"]) + require.Len(t, logMsgs, 2) + logs := requestLogsFromMessages(t, logMsgs) + for i, want := range []struct { + status entity.RequestStatus + requestVersion int32 + batchID string + }{ + {status: entity.RequestStatusBatching}, + {status: entity.RequestStatusBatched, requestVersion: request.Version + 1, batchID: "test-queue/batch/1"}, + } { + logEntry := logs[i] + assert.Equal(t, request.ID, logEntry.RequestID) + assert.Equal(t, want.status, logEntry.Status) + assert.Equal(t, want.requestVersion, logEntry.RequestVersion) + if want.batchID != "" { + assert.Equal(t, want.batchID, logEntry.Metadata["batch_id"]) + } + } } func TestController_Process_StorageFailure(t *testing.T) { @@ -510,13 +533,23 @@ func TestController_Process_CASLostToCancel(t *testing.T) { mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes() mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes() - // Publisher with no EXPECTs — must not be called. + // Allow only the early batching log publish; score must not be called. + var logMsg entityqueue.Message mockPub := queuemock.NewMockPublisher(ctrl) + mockPub.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).DoAndReturn( + func(_ context.Context, _ string, msg entityqueue.Message) error { + logMsg = msg + return nil + }, + ) mockQ := queuemock.NewMockQueue(ctrl) mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() registry, err := consumer.NewTopicRegistry( - []consumer.TopicConfig{{Key: topickey.TopicKeyScore, Name: "score", Queue: mockQ}}, + []consumer.TopicConfig{ + {Key: topickey.TopicKeyScore, Name: "score", Queue: mockQ}, + {Key: topickey.TopicKeyLog, Name: "log", Queue: mockQ}, + }, ) require.NoError(t, err) @@ -533,6 +566,10 @@ func TestController_Process_CASLostToCancel(t *testing.T) { delivery.EXPECT().Attempt().Return(1).AnyTimes() require.NoError(t, controller.Process(context.Background(), delivery)) + logEntry, err := entity.RequestLogFromBytes(logMsg.Payload) + require.NoError(t, err) + assert.Equal(t, request.ID, logEntry.RequestID) + assert.Equal(t, entity.RequestStatusBatching, logEntry.Status) } // Race-unexpected-error: any CAS failure other than ErrVersionMismatch (e.g. diff --git a/submitqueue/orchestrator/controller/score/score.go b/submitqueue/orchestrator/controller/score/score.go index 257ae86d..080603de 100644 --- a/submitqueue/orchestrator/controller/score/score.go +++ b/submitqueue/orchestrator/controller/score/score.go @@ -131,6 +131,13 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r return nil } + if err := corerequest.PublishBatchLogs(ctx, c.registry, batch.Contains, entity.RequestStatusScoring, map[string]string{ + "batch_id": batch.ID, + }); err != nil { + metrics.NamedCounter(c.metricsScope, opName, "request_log_errors", 1) + return fmt.Errorf("failed to publish scoring request logs for batch %s: %w", batch.ID, err) + } + // Score the batch. The scorer resolves the batch's changes itself. batchScore, err := c.scoreBatch(ctx, batch) if err != nil { diff --git a/submitqueue/orchestrator/controller/score/score_test.go b/submitqueue/orchestrator/controller/score/score_test.go index 63671aff..03cc94cb 100644 --- a/submitqueue/orchestrator/controller/score/score_test.go +++ b/submitqueue/orchestrator/controller/score/score_test.go @@ -104,14 +104,17 @@ func newMockStorage(ctrl *gomock.Controller, batch entity.Batch, request entity. } // newTestController creates a controller with test dependencies. -func newTestController(t *testing.T, ctrl *gomock.Controller, store *storagemock.MockStorage, scorer *scorermock.MockScorer, publishErr error) *Controller { +func newTestController(t *testing.T, ctrl *gomock.Controller, store *storagemock.MockStorage, scorer *scorermock.MockScorer, speculatePublishErr error) *Controller { logger := zaptest.NewLogger(t).Sugar() scope := tally.NoopScope mockPub := queuemock.NewMockPublisher(ctrl) mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( func(ctx context.Context, topic string, msg entityqueue.Message) error { - return publishErr + if topic == "speculate" { + return speculatePublishErr + } + return nil }, ).AnyTimes() @@ -211,6 +214,67 @@ func TestController_Process_BatchLevelScore(t *testing.T) { require.NoError(t, err) } +func TestController_Process_PublishesScoringAndScoredLogs(t *testing.T) { + ctrl := gomock.NewController(t) + + batch := testBatch() + mockBatchStore := storagemock.NewMockBatchStore(ctrl) + mockBatchStore.EXPECT().Get(gomock.Any(), batch.ID).Return(batch, nil) + mockBatchStore.EXPECT().UpdateScoreAndState(gomock.Any(), batch.ID, batch.Version, batch.Version+1, 0.7, entity.BatchStateScored).Return(nil) + + store := storagemock.NewMockStorage(ctrl) + store.EXPECT().GetBatchStore().Return(mockBatchStore).AnyTimes() + + mockScorer := scorermock.NewMockScorer(ctrl) + mockScorer.EXPECT().Score(gomock.Any(), gomock.Any()).Return(0.7, nil) + + var logMsgs []entityqueue.Message + mockPub := queuemock.NewMockPublisher(ctrl) + mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn( + func(_ context.Context, topic string, msg entityqueue.Message) error { + if topic == "log" { + logMsgs = append(logMsgs, msg) + } + return nil + }, + ).AnyTimes() + mockQ := queuemock.NewMockQueue(ctrl) + mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes() + + registry, err := consumer.NewTopicRegistry( + []consumer.TopicConfig{ + {Key: topickey.TopicKeySpeculate, Name: "speculate", Queue: mockQ}, + {Key: topickey.TopicKeyLog, Name: "log", Queue: mockQ}, + }, + ) + require.NoError(t, err) + + scorerFactory := scorermock.NewMockFactory(ctrl) + scorerFactory.EXPECT().For(gomock.Any()).Return(mockScorer, nil).AnyTimes() + controller := NewController(zaptest.NewLogger(t).Sugar(), tally.NoopScope, store, scorerFactory, 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(1).AnyTimes() + + require.NoError(t, controller.Process(context.Background(), delivery)) + require.Len(t, logMsgs, 2) + + scoringLog, err := entity.RequestLogFromBytes(logMsgs[0].Payload) + require.NoError(t, err) + assert.Equal(t, batch.Contains[0], scoringLog.RequestID) + assert.Equal(t, entity.RequestStatusScoring, scoringLog.Status) + assert.Equal(t, batch.ID, scoringLog.Metadata["batch_id"]) + + scoredLog, err := entity.RequestLogFromBytes(logMsgs[1].Payload) + require.NoError(t, err) + assert.Equal(t, batch.Contains[0], scoredLog.RequestID) + assert.Equal(t, entity.RequestStatusScored, scoredLog.Status) + assert.Equal(t, batch.ID, scoredLog.Metadata["batch_id"]) + assert.Equal(t, "0.7000", scoredLog.Metadata["score"]) +} + func TestController_Process_StorageFailure(t *testing.T) { ctrl := gomock.NewController(t)