[CRE][DONTime] Respect batchSize config (#2184)

Author: bolekk

pkg/workflows/dontime/pb/dontime.pb.go

@@ -5,8 +5,8 @@ option go_package = "github.com/smartcontractkit/chainlink-common/pkg/workflows/
 message Observation {
   int64 timestamp = 1;
   map<string, int64> requests = 2;
-  // Flag to roll out execution pruning fix. Can be removed after rollout (once unused in the outcome phase).
-  bool prune_executions = 3;
+  reserved 3;
+  bool limit_by_batch_size_flag = 4;
 }
 
 message ObservedDonTimes {

pkg/workflows/dontime/pb/dontime.proto

@@ -22,9 +22,11 @@ import (
 )
 
 type pluginMetrics struct {
-	donTime        metric.Int64Gauge
-	donTimeEntries metric.Int64Gauge
-	outcomeSize    metric.Int64Gauge
+	donTime                  metric.Int64Gauge
+	donTimeEntries           metric.Int64Gauge
+	outcomeSize              metric.Int64Gauge
+	observationBatchOverflow metric.Int64Gauge
+	outcomeBatchOverflow     metric.Int64Gauge
 }
 
 func newPluginMetrics() (pluginMetrics, error) {
@@ -54,10 +56,28 @@ func newPluginMetrics() (pluginMetrics, error) {
 		return pluginMetrics{}, fmt.Errorf("failed to create outcome_size gauge: %w", err)
 	}
 
+	observationBatchOverflow, err := meter.Int64Gauge("platform_dontime_observation_batch_overflow",
+		metric.WithDescription("Number of pending requests excluded from the observation due to batch size limit"),
+		metric.WithUnit("{request}"),
+	)
+	if err != nil {
+		return pluginMetrics{}, fmt.Errorf("failed to create observation_batch_overflow gauge: %w", err)
+	}
+
+	outcomeBatchOverflow, err := meter.Int64Gauge("platform_dontime_outcome_batch_overflow",
+		metric.WithDescription("Number of workflow execution entries removed from the outcome due to batch size limit"),
+		metric.WithUnit("{entry}"),
+	)
+	if err != nil {
+		return pluginMetrics{}, fmt.Errorf("failed to create outcome_batch_overflow gauge: %w", err)
+	}
+
 	return pluginMetrics{
-		donTime:        donTime,
-		donTimeEntries: donTimeEntries,
-		outcomeSize:    outcomeSize,
+		donTime:                  donTime,
+		donTimeEntries:           donTimeEntries,
+		outcomeSize:              outcomeSize,
+		observationBatchOverflow: observationBatchOverflow,
+		outcomeBatchOverflow:     outcomeBatchOverflow,
 	}, nil
 }
 
@@ -106,14 +126,33 @@ func (p *Plugin) Query(_ context.Context, _ ocr3types.OutcomeContext) (types.Que
 	return nil, nil
 }
 
-func (p *Plugin) Observation(_ context.Context, outctx ocr3types.OutcomeContext, query types.Query) (types.Observation, error) {
+func sortedRequests(requests map[string]*Request) []*Request {
+	if len(requests) == 0 {
+		return nil
+	}
+
+	ids := make([]string, 0, len(requests))
+	for id := range requests {
+		ids = append(ids, id)
+	}
+	slices.Sort(ids)
+
+	sorted := make([]*Request, 0, len(ids))
+	for _, id := range ids {
+		sorted = append(sorted, requests[id])
+	}
+	return sorted
+}
+
+func (p *Plugin) Observation(ctx context.Context, outctx ocr3types.OutcomeContext, query types.Query) (types.Observation, error) {
 	previousOutcome := &pb.Outcome{}
 	if err := proto.Unmarshal(outctx.PreviousOutcome, previousOutcome); err != nil {
 		p.lggr.Errorf("failed to unmarshal previous outcome in Observation phase")
 	}
 
+	sortedRequests := sortedRequests(p.store.GetRequests())
 	requests := map[string]int64{} // Maps executionID --> seqNum
-	for _, req := range p.store.GetRequests() {
+	for _, req := range sortedRequests {
 		// Validate request sequence number
 		numObservedDonTimes := 0
 		times, ok := previousOutcome.ObservedDonTimes[req.WorkflowExecutionID]
@@ -135,12 +174,27 @@ func (p *Plugin) Observation(_ context.Context, outctx ocr3types.OutcomeContext,
 		}
 
 		requests[req.WorkflowExecutionID] = int64(req.SeqNum)
+		if len(requests) >= p.batchSize {
+			break
+		}
 	}
 
+	overflowCount := len(sortedRequests) - len(requests)
+	p.lggr.Debugw("Observation batch processed",
+		"inputRequests", len(sortedRequests),
+		"batchSize", p.batchSize,
+		"includedRequests", len(requests),
+		"overflowRequests", overflowCount,
+	)
+	if overflowCount > 0 {
+		p.lggr.Warnw("Observation batch overflow", "overflowRequests", overflowCount)
+	}
+	p.metrics.observationBatchOverflow.Record(ctx, int64(overflowCount))
+
 	observation := &pb.Observation{
-		Timestamp:       time.Now().UTC().UnixMilli(),
-		Requests:        requests,
-		PruneExecutions: true,
+		Timestamp:            time.Now().UTC().UnixMilli(),
+		Requests:             requests,
+		LimitByBatchSizeFlag: true,
 	}
 
 	return proto.MarshalOptions{Deterministic: true}.Marshal(observation)
@@ -162,6 +216,7 @@ func (p *Plugin) Outcome(ctx context.Context, outctx ocr3types.OutcomeContext, _
 		OffsetFromMedian int64
 	}
 	var timestampNodePairs []timestampNodePair
+	limitByBatchSizeFlagEnabled := true
 
 	prevOutcome := &pb.Outcome{}
 	if err := proto.Unmarshal(outctx.PreviousOutcome, prevOutcome); err != nil {
@@ -178,6 +233,10 @@ func (p *Plugin) Outcome(ctx context.Context, outctx ocr3types.OutcomeContext, _
 			continue
 		}
 
+		if !observation.GetLimitByBatchSizeFlag() {
+			limitByBatchSizeFlagEnabled = false
+		}
+
 		for id, requestSeqNum := range observation.Requests {
 			var currSeqNum int64
 			if times, ok := prevOutcome.ObservedDonTimes[id]; ok {
@@ -249,13 +308,31 @@ func (p *Plugin) Outcome(ctx context.Context, outctx ocr3types.OutcomeContext, _
 		}
 	}
 
+	var outcomeBatchOverflowCount int64
+	if len(outcome.ObservedDonTimes) > p.batchSize && limitByBatchSizeFlagEnabled {
+		ids := make([]string, 0, len(outcome.ObservedDonTimes))
+		for id := range outcome.ObservedDonTimes {
+			ids = append(ids, id)
+		}
+		slices.Sort(ids)
+		outcomeBatchOverflowCount = int64(len(ids) - p.batchSize)
+		for _, id := range ids[p.batchSize:] {
+			delete(outcome.ObservedDonTimes, id)
+		}
+		p.lggr.Warnw("Trimmed outcome observed don times to batch size",
+			"batchSize", p.batchSize,
+			"removedEntries", outcomeBatchOverflowCount,
+		)
+	}
+
 	outcomeBytes, err := proto.MarshalOptions{Deterministic: true}.Marshal(outcome)
 	p.lggr.Infow("Outcome computed",
 		"observedDonTimesEntries", len(outcome.ObservedDonTimes),
 		"outcomeSizeBytes", len(outcomeBytes),
 	)
 	p.metrics.donTime.Record(ctx, outcome.Timestamp)
 	p.metrics.donTimeEntries.Record(ctx, int64(len(outcome.ObservedDonTimes)))
+	p.metrics.outcomeBatchOverflow.Record(ctx, outcomeBatchOverflowCount)
 	p.metrics.outcomeSize.Record(ctx, int64(len(outcomeBytes)))
 	return outcomeBytes, err
 }

pkg/workflows/dontime/plugin.go

@@ -82,6 +82,35 @@ func TestPlugin_Observation(t *testing.T) {
 		require.Equal(t, expectedRequests, obsProto.Requests)
 		store.deleteExecutionID("workflow-123")
 	})
+
+	t.Run("batch limit excludes overflow requests", func(t *testing.T) {
+		offchainCfg := newTestPluginOffchainConfig(t)
+		offchainCfg.MaxBatchSize = 2
+		plugin, err := NewPlugin(store, config, offchainCfg, lggr)
+		require.NoError(t, err)
+
+		for _, id := range []string{"workflow-c", "workflow-a", "workflow-b"} {
+			_ = store.RequestDonTime(id, 0)
+		}
+
+		observation, err := plugin.Observation(ctx, outcomeCtx, query)
+		require.NoError(t, err)
+
+		obsProto := &pb.Observation{}
+		err = proto.Unmarshal(observation, obsProto)
+		require.NoError(t, err)
+
+		expectedRequests := map[string]int64{
+			"workflow-a": 0,
+			"workflow-b": 0,
+		}
+		require.Equal(t, expectedRequests, obsProto.Requests)
+		require.NotNil(t, store.GetRequest("workflow-c"))
+
+		for _, id := range []string{"workflow-a", "workflow-b", "workflow-c"} {
+			store.deleteExecutionID(id)
+		}
+	})
 }
 
 func TestPlugin_ValidateObservation(t *testing.T) {
@@ -518,6 +547,73 @@ func TestPlugin_FinishedExecutions(t *testing.T) {
 	})
 }
 
+func TestPlugin_Outcome_TrimByBatchSize(t *testing.T) {
+	lggr := logger.Test(t)
+	store := NewStore(DefaultRequestTimeout)
+	config, offchainCfg := newTestPluginConfig(t), newTestPluginOffchainConfig(t)
+	offchainCfg.MaxBatchSize = 2
+	ctx := t.Context()
+
+	plugin, err := NewPlugin(store, config, offchainCfg, lggr)
+	require.NoError(t, err)
+
+	query, err := plugin.Query(ctx, ocr3types.OutcomeContext{PreviousOutcome: []byte("")})
+	require.NoError(t, err)
+
+	timestamp := time.Now().UnixMilli()
+	makeObservations := func(limitByBatchSize bool) []types.AttributedObservation {
+		aos := make([]types.AttributedObservation, 4)
+		for i := 0; i < 4; i++ {
+			obs := &pb.Observation{
+				Timestamp:            timestamp + int64(i),
+				Requests:             map[string]int64{},
+				LimitByBatchSizeFlag: limitByBatchSize,
+			}
+			rawObs, err := proto.Marshal(obs)
+			require.NoError(t, err)
+			aos[i] = types.AttributedObservation{
+				Observation: rawObs,
+				Observer:    commontypes.OracleID(i),
+			}
+		}
+		return aos
+	}
+
+	prevOutcome := &pb.Outcome{
+		Timestamp: timestamp - 1000,
+		ObservedDonTimes: map[string]*pb.ObservedDonTimes{
+			"workflow-a": {Timestamps: []int64{timestamp - 1000}},
+			"workflow-b": {Timestamps: []int64{timestamp - 1000}},
+			"workflow-c": {Timestamps: []int64{timestamp - 1000}},
+		},
+	}
+	prevOutcomeBytes, err := proto.Marshal(prevOutcome)
+	require.NoError(t, err)
+
+	t.Run("trims when all observations set batch size flag", func(t *testing.T) {
+		outcome, err := plugin.Outcome(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes}, query, makeObservations(true))
+		require.NoError(t, err)
+
+		outcomeProto := &pb.Outcome{}
+		err = proto.Unmarshal(outcome, outcomeProto)
+		require.NoError(t, err)
+		require.Len(t, outcomeProto.ObservedDonTimes, 2)
+		require.Contains(t, outcomeProto.ObservedDonTimes, "workflow-a")
+		require.Contains(t, outcomeProto.ObservedDonTimes, "workflow-b")
+		require.NotContains(t, outcomeProto.ObservedDonTimes, "workflow-c")
+	})
+
+	t.Run("does not trim when batch size flag is missing", func(t *testing.T) {
+		outcome, err := plugin.Outcome(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes}, query, makeObservations(false))
+		require.NoError(t, err)
+
+		outcomeProto := &pb.Outcome{}
+		err = proto.Unmarshal(outcome, outcomeProto)
+		require.NoError(t, err)
+		require.Len(t, outcomeProto.ObservedDonTimes, 3)
+	})
+}
+
 func TestPlugin_ExpiredRequest(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore(0)