CRE-362: dynamic batching based on observation sizes

pkg/capabilities/consensus/ocr3/factory.go

@@ -19,6 +19,7 @@ const (
 	defaultMaxPhaseOutputBytes     = 1000000 // 1 MB
 	defaultMaxReportCount          = 20
 	defaultBatchSize               = 20
+	defaultBatchSizeMiB            = 1024 * 1024 * 5
 	defaultOutcomePruningThreshold = 3600
 	defaultRequestExpiry           = 20 * time.Second
 )

pkg/capabilities/consensus/ocr3/reporting_plugin.go

@@ -57,33 +57,24 @@ func NewReportingPlugin(s *requests.Store, r CapabilityIface, batchSize int, con
 	}, nil
 }
 
-func (r *reportingPlugin) Query(ctx context.Context, outctx ocr3types.OutcomeContext) (types.Query, error) {
-	batch, err := r.s.FirstN(r.batchSize)
+func (r *reportingPlugin) Query(_ context.Context, _ ocr3types.OutcomeContext) (types.Query, error) {
+	all, err := r.s.All()
 	if err != nil {
-		r.lggr.Errorw("could not retrieve batch", "error", err)
 		return nil, err
 	}
 
-	ids := []*pbtypes.Id{}
-	allExecutionIDs := []string{}
-	for _, rq := range batch {
-		ids = append(ids, &pbtypes.Id{
-			WorkflowExecutionId:      rq.WorkflowExecutionID,
-			WorkflowId:               rq.WorkflowID,
-			WorkflowOwner:            rq.WorkflowOwner,
-			WorkflowName:             rq.WorkflowName,
-			WorkflowDonId:            rq.WorkflowDonID,
-			WorkflowDonConfigVersion: rq.WorkflowDonConfigVersion,
-			ReportId:                 rq.ReportID,
-			KeyId:                    rq.KeyID,
-		})
-		allExecutionIDs = append(allExecutionIDs, rq.WorkflowExecutionID)
+	allExecutionIDs, serializedBatch, err := packToSizeLimit(
+		r.lggr,
+		QueriesSerializable{all},
+	)
+	if err != nil {
+		r.lggr.Errorw("could not serialize query batch", "error", err)
+		return nil, err
 	}
 
-	r.lggr.Debugw("Query complete", "len", len(ids), "allExecutionIDs", allExecutionIDs)
-	return proto.MarshalOptions{Deterministic: true}.Marshal(&pbtypes.Query{
-		Ids: ids,
-	})
+	r.lggr.Debugw("Query complete", "len", len(serializedBatch), "allExecutionIDs", allExecutionIDs)
+
+	return serializedBatch, nil
 }
 
 func (r *reportingPlugin) Observation(ctx context.Context, outctx ocr3types.OutcomeContext, query types.Query) (types.Observation, error) {
@@ -93,7 +84,7 @@ func (r *reportingPlugin) Observation(ctx context.Context, outctx ocr3types.Outc
 		return nil, err
 	}
 
-	weids := []string{}
+	weids := make([]string, 0)
 	for _, q := range queryReq.Ids {
 		if q == nil {
 			r.lggr.Debugw("skipping nil id for query", "query", queryReq)
@@ -108,57 +99,28 @@ func (r *reportingPlugin) Observation(ctx context.Context, outctx ocr3types.Outc
 		reqMap[req.WorkflowExecutionID] = req
 	}
 
-	obs := &pbtypes.Observations{}
-	allExecutionIDs := []string{}
-	for _, weid := range weids {
-		rq, ok := reqMap[weid]
-		if !ok {
-			r.lggr.Debugw("could not find local observations for weid requested in the query", "executionID", weid)
-			continue
-		}
+	registeredWorkflowIDs := r.r.GetRegisteredWorkflowsIDs()
 
-		lggr := logger.With(
-			r.lggr,
-			"executionID", rq.WorkflowExecutionID,
-			"workflowID", rq.WorkflowID,
-		)
-
-		listProto := values.Proto(rq.Observations).GetListValue()
-		if listProto == nil {
-			lggr.Errorw("observations are not a list")
-			continue
-		}
-
-		var cfgProto *pb.Map
-		if rq.OverriddenEncoderConfig != nil {
-			cp := values.Proto(rq.OverriddenEncoderConfig).GetMapValue()
-			cfgProto = cp
-		}
-
-		newOb := &pbtypes.Observation{
-			Observations: listProto,
-			Id: &pbtypes.Id{
-				WorkflowExecutionId:      rq.WorkflowExecutionID,
-				WorkflowId:               rq.WorkflowID,
-				WorkflowOwner:            rq.WorkflowOwner,
-				WorkflowName:             rq.WorkflowName,
-				WorkflowDonId:            rq.WorkflowDonID,
-				WorkflowDonConfigVersion: rq.WorkflowDonConfigVersion,
-				ReportId:                 rq.ReportID,
-				KeyId:                    rq.KeyID,
-			},
-			OverriddenEncoderName:   rq.OverriddenEncoderName,
-			OverriddenEncoderConfig: cfgProto,
-		}
-
-		obs.Observations = append(obs.Observations, newOb)
-		allExecutionIDs = append(allExecutionIDs, rq.WorkflowExecutionID)
+	allExecutionIDs, obs, err := packToSizeLimit(
+		r.lggr,
+		ObservationSerializable{reqMap, registeredWorkflowIDs},
+	)
+	if err != nil {
+		r.lggr.Errorw("could not serialize observations batch", "error", err)
+		return nil, err
 	}
-	obs.RegisteredWorkflowIds = r.r.GetRegisteredWorkflowsIDs()
-	obs.Timestamp = timestamppb.New(time.Now())
 
-	r.lggr.Debugw("Observation complete", "len", len(obs.Observations), "queryLen", len(queryReq.Ids), "allExecutionIDs", allExecutionIDs)
-	return proto.MarshalOptions{Deterministic: true}.Marshal(obs)
+	r.lggr.Debugw(
+		"Observation complete",
+		"len",
+		len(obs),
+		"queryLen",
+		len(queryReq.Ids),
+		"allExecutionIDs",
+		allExecutionIDs,
+	)
+
+	return obs, nil
 }
 
 func (r *reportingPlugin) ValidateObservation(ctx context.Context, outctx ocr3types.OutcomeContext, query types.Query, ao types.AttributedObservation) error {
@@ -306,14 +268,76 @@ func (r *reportingPlugin) Outcome(ctx context.Context, outctx ocr3types.OutcomeC
 	// every time since we only want to transmit reports that
 	// are part of the current Query.
 	previousOutcome.CurrentReports = []*pbtypes.Report{}
-	var allExecutionIDs []string
 
+	prepWeids, prepOutcomes := r.OutcomePrep(
+		execIDToOracleObservations,
+		execIDToEncoderShaToCount,
+		shaToEncoder,
+		previousOutcome,
+		q,
+		finalTimestamp,
+	)
+
+	previousOutcome.Outcomes = prepOutcomes
+
+	// We need to prune outcomes from previous workflows that are no longer relevant.
+	for workflowID, outcome := range previousOutcome.Outcomes {
+		// Update the last seen round for this outcome. But this should only happen if the workflow is seen by F+1 nodes.
+		if seenWorkflowIDs[workflowID] >= (r.config.F + 1) {
+			r.lggr.Debugw("updating last seen round of outcome for workflow", "workflowID", workflowID)
+			outcome.LastSeenAt = outctx.SeqNr
+		} else if outctx.SeqNr-outcome.LastSeenAt > r.outcomePruningThreshold {
+			r.lggr.Debugw("pruning outcome for workflow", "workflowID", workflowID, "SeqNr", outctx.SeqNr, "lastSeenAt", outcome.LastSeenAt)
+			delete(previousOutcome.Outcomes, workflowID)
+			r.r.UnregisterWorkflowID(workflowID)
+		}
+	}
+
+	outcomeSerializable := &OutcomeSerializable{
+		previousOutcome: previousOutcome,
+		weids:           prepWeids,
+		outcomes:        prepOutcomes,
+	}
+
+	allExecutionIDs, rawOutcome, err := packToSizeLimit(
+		r.lggr,
+		outcomeSerializable,
+	)
+
+	previousOutcome.Outcomes = outcomeSerializable.previousOutcome.Outcomes
+	previousOutcome.CurrentReports = outcomeSerializable.previousOutcome.CurrentReports
+
+	if err != nil {
+		r.lggr.Errorw("could not serialize outcomes batch", "error", err)
+		return nil, err
+	}
+
+	h := sha256.New()
+	h.Write(rawOutcome)
+	outcomeHash := h.Sum(nil)
+	r.lggr.Debugw("Outcome complete", "len", len(previousOutcome.Outcomes), "nAggregatedWorkflowExecutions", len(previousOutcome.CurrentReports), "allExecutionIDs", allExecutionIDs, "outcomeHash", hex.EncodeToString(outcomeHash), "err", err)
+	return rawOutcome, err
+}
+
+func (r *reportingPlugin) OutcomePrep(
+	execIDToOracleObservations map[string]map[ocrcommon.OracleID][]values.Value,
+	execIDToEncoderShaToCount map[string]map[string]int,
+	shaToEncoder map[string]encoderConfig,
+	previousOutcome *pbtypes.Outcome,
+	q *pbtypes.Query,
+	finalTimestamp *timestamppb.Timestamp,
+) ([]*pbtypes.Id, map[string]*pbtypes.AggregationOutcome) {
+	lggr := r.lggr
+	weids := make([]*pbtypes.Id, 0)
+	outcomes := make(map[string]*pbtypes.AggregationOutcome, 0)
+	rc := r.r
+	F := r.config.F
 	for _, weid := range q.Ids {
 		if weid == nil {
-			r.lggr.Debugw("skipping nil id in query", "query", q)
+			lggr.Debugw("skipping nil id in query", "query", q)
 			continue
 		}
-		lggr := logger.With(r.lggr, "executionID", weid.WorkflowExecutionId, "workflowID", weid.WorkflowId)
+		lggr := logger.With(lggr, "executionID", weid.WorkflowExecutionId, "workflowID", weid.WorkflowId)
 		obs, ok := execIDToOracleObservations[weid.WorkflowExecutionId]
 		if !ok {
 			lggr.Debugw("could not find any observations matching weid requested in the query")
@@ -325,18 +349,18 @@ func (r *reportingPlugin) Outcome(ctx context.Context, outctx ocr3types.OutcomeC
 			lggr.Debugw("could not find existing outcome for workflow, aggregator will create a new one")
 		}
 
-		if len(obs) < (2*r.config.F + 1) {
+		if len(obs) < (2*F + 1) {
 			lggr.Debugw("insufficient observations for workflow execution id")
 			continue
 		}
 
-		agg, err2 := r.r.GetAggregator(weid.WorkflowId)
+		agg, err2 := rc.GetAggregator(weid.WorkflowId)
 		if err2 != nil {
 			lggr.Errorw("could not retrieve aggregator for workflow", "error", err2)
 			continue
 		}
 
-		outcome, err2 := agg.Aggregate(lggr, workflowOutcome, obs, r.config.F)
+		outcome, err2 := agg.Aggregate(lggr, workflowOutcome, obs, F)
 		if err2 != nil {
 			lggr.Errorw("error aggregating outcome", "error", err2)
 			continue
@@ -361,7 +385,7 @@ func (r *reportingPlugin) Outcome(ctx context.Context, outctx ocr3types.OutcomeC
 		// we've checked this above when we checked the observations count.
 		var encCfg *encoderConfig
 		for sha, count := range shaToCount {
-			if count >= 2*r.config.F+1 {
+			if count >= 2*F+1 {
 				encoderCfg, ok := shaToEncoder[sha]
 				if !ok {
 					lggr.Debugw("could not find encoder matching sha")
@@ -380,35 +404,10 @@ func (r *reportingPlugin) Outcome(ctx context.Context, outctx ocr3types.OutcomeC
 			outcome.EncoderConfig = encCfg.config
 		}
 
-		report := &pbtypes.Report{
-			Outcome: outcome,
-			Id:      weid,
-		}
-		previousOutcome.CurrentReports = append(previousOutcome.CurrentReports, report)
-		allExecutionIDs = append(allExecutionIDs, weid.WorkflowExecutionId)
-
-		previousOutcome.Outcomes[weid.WorkflowId] = outcome
-	}
-
-	// We need to prune outcomes from previous workflows that are no longer relevant.
-	for workflowID, outcome := range previousOutcome.Outcomes {
-		// Update the last seen round for this outcome. But this should only happen if the workflow is seen by F+1 nodes.
-		if seenWorkflowIDs[workflowID] >= (r.config.F + 1) {
-			r.lggr.Debugw("updating last seen round of outcome for workflow", "workflowID", workflowID)
-			outcome.LastSeenAt = outctx.SeqNr
-		} else if outctx.SeqNr-outcome.LastSeenAt > r.outcomePruningThreshold {
-			r.lggr.Debugw("pruning outcome for workflow", "workflowID", workflowID, "SeqNr", outctx.SeqNr, "lastSeenAt", outcome.LastSeenAt)
-			delete(previousOutcome.Outcomes, workflowID)
-			r.r.UnregisterWorkflowID(workflowID)
-		}
+		weids = append(weids, weid)
+		outcomes[weid.WorkflowExecutionId] = outcome
 	}
-
-	rawOutcome, err := proto.MarshalOptions{Deterministic: true}.Marshal(previousOutcome)
-	h := sha256.New()
-	h.Write(rawOutcome)
-	outcomeHash := h.Sum(nil)
-	r.lggr.Debugw("Outcome complete", "len", len(previousOutcome.Outcomes), "nAggregatedWorkflowExecutions", len(previousOutcome.CurrentReports), "allExecutionIDs", allExecutionIDs, "outcomeHash", hex.EncodeToString(outcomeHash), "err", err)
-	return rawOutcome, err
+	return weids, outcomes
 }
 
 func marshalReportInfo(info *pbtypes.ReportInfo, keyID string) ([]byte, error) {