Skip to content

Commit 6ef5943

Browse files
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 <cursoragent@cursor.com>
1 parent 5607a83 commit 6ef5943

5 files changed

Lines changed: 71 additions & 15 deletions

File tree

submitqueue/entity/request_log.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,9 @@ const (
4949
// RequestStatusBatched indicates that the request has been included in a new batch and will be sent to speculation.
5050
RequestStatusBatched RequestStatus = "batched"
5151

52+
// RequestStatusScoring indicates that the batch containing the request is being scored for build success probability.
53+
RequestStatusScoring RequestStatus = "scoring"
54+
5255
// RequestStatusScored indicates that the batch containing the request has been scored for build success probability.
5356
RequestStatusScored RequestStatus = "scored"
5457

submitqueue/orchestrator/controller/batch/batch.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -118,6 +118,12 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r
118118
return nil
119119
}
120120

121+
batchingLog := entity.NewRequestLog(request.ID, entity.RequestStatusBatching, 0, "", nil)
122+
if err := corerequest.PublishLog(ctx, c.registry, batchingLog, request.ID); err != nil {
123+
metrics.NamedCounter(c.metricsScope, opName, "request_log_errors", 1)
124+
return fmt.Errorf("failed to publish batching request log for request %s: %w", request.ID, err)
125+
}
126+
121127
// 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.
122128

123129
// Generate a globally unique batch ID.

submitqueue/orchestrator/controller/batch/batch_test.go

Lines changed: 50 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,18 @@ func requestIDPayload(t *testing.T, id string) []byte {
4848
return payload
4949
}
5050

51+
func requestLogsFromMessages(t *testing.T, msgs []entityqueue.Message) []entity.RequestLog {
52+
t.Helper()
53+
54+
logs := make([]entity.RequestLog, 0, len(msgs))
55+
for _, msg := range msgs {
56+
logEntry, err := entity.RequestLogFromBytes(msg.Payload)
57+
require.NoError(t, err)
58+
logs = append(logs, logEntry)
59+
}
60+
return logs
61+
}
62+
5163
// newSequentialCounter returns a mock counter that returns incrementing values starting at 1.
5264
func newSequentialCounter(ctrl *gomock.Controller) *countermock.MockCounter {
5365
var seq int64
@@ -157,10 +169,10 @@ func TestController_Process_Success(t *testing.T) {
157169
require.NoError(t, err)
158170
}
159171

160-
// TestController_Process_PublishesBatchedLog asserts the controller emits a
161-
// "batched" request log carrying the request ID, the post-CAS request version,
162-
// and the batch ID it was placed into.
163-
func TestController_Process_PublishesBatchedLog(t *testing.T) {
172+
// TestController_Process_PublishesBatchingAndBatchedLogs asserts the controller
173+
// emits a progress log when active batching starts, then the versioned "batched"
174+
// log once the request has been claimed into a persisted batch.
175+
func TestController_Process_PublishesBatchingAndBatchedLogs(t *testing.T) {
164176
ctrl := gomock.NewController(t)
165177

166178
request := testRequest()
@@ -217,13 +229,24 @@ func TestController_Process_PublishesBatchedLog(t *testing.T) {
217229

218230
require.NoError(t, controller.Process(context.Background(), delivery))
219231

220-
require.Len(t, logMsgs, 1)
221-
logEntry, err := entity.RequestLogFromBytes(logMsgs[0].Payload)
222-
require.NoError(t, err)
223-
assert.Equal(t, request.ID, logEntry.RequestID)
224-
assert.Equal(t, entity.RequestStatusBatched, logEntry.Status)
225-
assert.Equal(t, request.Version+1, logEntry.RequestVersion)
226-
assert.Equal(t, "test-queue/batch/1", logEntry.Metadata["batch_id"])
232+
require.Len(t, logMsgs, 2)
233+
logs := requestLogsFromMessages(t, logMsgs)
234+
for i, want := range []struct {
235+
status entity.RequestStatus
236+
requestVersion int32
237+
batchID string
238+
}{
239+
{status: entity.RequestStatusBatching},
240+
{status: entity.RequestStatusBatched, requestVersion: request.Version + 1, batchID: "test-queue/batch/1"},
241+
} {
242+
logEntry := logs[i]
243+
assert.Equal(t, request.ID, logEntry.RequestID)
244+
assert.Equal(t, want.status, logEntry.Status)
245+
assert.Equal(t, want.requestVersion, logEntry.RequestVersion)
246+
if want.batchID != "" {
247+
assert.Equal(t, want.batchID, logEntry.Metadata["batch_id"])
248+
}
249+
}
227250
}
228251

229252
func TestController_Process_StorageFailure(t *testing.T) {
@@ -510,13 +533,23 @@ func TestController_Process_CASLostToCancel(t *testing.T) {
510533
mockStorage.EXPECT().GetBatchDependentStore().Return(mockBatchDependentStore).AnyTimes()
511534
mockStorage.EXPECT().GetRequestStore().Return(mockReqStore).AnyTimes()
512535

513-
// Publisher with no EXPECTs — must not be called.
536+
// Allow only the early batching log publish; score must not be called.
537+
var logMsg entityqueue.Message
514538
mockPub := queuemock.NewMockPublisher(ctrl)
539+
mockPub.EXPECT().Publish(gomock.Any(), "log", gomock.Any()).DoAndReturn(
540+
func(_ context.Context, _ string, msg entityqueue.Message) error {
541+
logMsg = msg
542+
return nil
543+
},
544+
)
515545
mockQ := queuemock.NewMockQueue(ctrl)
516546
mockQ.EXPECT().Publisher().Return(mockPub).AnyTimes()
517547

518548
registry, err := consumer.NewTopicRegistry(
519-
[]consumer.TopicConfig{{Key: topickey.TopicKeyScore, Name: "score", Queue: mockQ}},
549+
[]consumer.TopicConfig{
550+
{Key: topickey.TopicKeyScore, Name: "score", Queue: mockQ},
551+
{Key: topickey.TopicKeyLog, Name: "log", Queue: mockQ},
552+
},
520553
)
521554
require.NoError(t, err)
522555

@@ -533,6 +566,10 @@ func TestController_Process_CASLostToCancel(t *testing.T) {
533566
delivery.EXPECT().Attempt().Return(1).AnyTimes()
534567

535568
require.NoError(t, controller.Process(context.Background(), delivery))
569+
logEntry, err := entity.RequestLogFromBytes(logMsg.Payload)
570+
require.NoError(t, err)
571+
assert.Equal(t, request.ID, logEntry.RequestID)
572+
assert.Equal(t, entity.RequestStatusBatching, logEntry.Status)
536573
}
537574

538575
// Race-unexpected-error: any CAS failure other than ErrVersionMismatch (e.g.

submitqueue/orchestrator/controller/score/score.go

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -131,6 +131,13 @@ func (c *Controller) Process(ctx context.Context, delivery consumer.Delivery) (r
131131
return nil
132132
}
133133

134+
if err := corerequest.PublishBatchLogs(ctx, c.registry, batch.Contains, entity.RequestStatusScoring, map[string]string{
135+
"batch_id": batch.ID,
136+
}); err != nil {
137+
metrics.NamedCounter(c.metricsScope, opName, "request_log_errors", 1)
138+
return fmt.Errorf("failed to publish scoring request logs for batch %s: %w", batch.ID, err)
139+
}
140+
134141
// Score the batch. The scorer resolves the batch's changes itself.
135142
batchScore, err := c.scoreBatch(ctx, batch)
136143
if err != nil {

submitqueue/orchestrator/controller/score/score_test.go

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -104,14 +104,17 @@ func newMockStorage(ctrl *gomock.Controller, batch entity.Batch, request entity.
104104
}
105105

106106
// newTestController creates a controller with test dependencies.
107-
func newTestController(t *testing.T, ctrl *gomock.Controller, store *storagemock.MockStorage, scorer *scorermock.MockScorer, publishErr error) *Controller {
107+
func newTestController(t *testing.T, ctrl *gomock.Controller, store *storagemock.MockStorage, scorer *scorermock.MockScorer, speculatePublishErr error) *Controller {
108108
logger := zaptest.NewLogger(t).Sugar()
109109
scope := tally.NoopScope
110110

111111
mockPub := queuemock.NewMockPublisher(ctrl)
112112
mockPub.EXPECT().Publish(gomock.Any(), gomock.Any(), gomock.Any()).DoAndReturn(
113113
func(ctx context.Context, topic string, msg entityqueue.Message) error {
114-
return publishErr
114+
if topic == "speculate" {
115+
return speculatePublishErr
116+
}
117+
return nil
115118
},
116119
).AnyTimes()
117120

0 commit comments

Comments
 (0)