crosscluster: move aggregate range stats collector into utils

This patch moves the aggregate range stats collector for LDR into the
replication utils package so that it can be used in PCR.

Epic: None

Release note: None

Author: kev-cao

pkg/crosscluster/logical/BUILD.bazel

@@ -15,7 +15,6 @@ go_library(
         "metrics.go",
         "offline_initial_scan_processor.go",
         "purgatory.go",
-        "range_stats.go",
         "replication_statements.go",
         "savepoint.go",
         "sql_crud_writer.go",
@@ -128,7 +127,6 @@ go_test(
         "lww_row_processor_test.go",
         "main_test.go",
         "purgatory_test.go",
-        "range_stats_test.go",
         "replication_statements_test.go",
         "savepoint_test.go",
         "sql_row_reader_test.go",

pkg/crosscluster/logical/logical_replication_job.go

@@ -289,8 +289,10 @@ func (r *logicalReplicationResumer) ingest(
 			settings:              &execCfg.Settings.SV,
 			job:                   r.job,
 			frontierUpdates:       heartbeatSender.FrontierUpdates,
-			rangeStats:            newRangeStatsCollector(planInfo.writeProcessorCount),
-			r:                     r,
+			rangeStats: replicationutils.NewAggregateRangeStatsCollector(
+				planInfo.writeProcessorCount,
+			),
+			r: r,
 		}
 		rowResultWriter := sql.NewCallbackResultWriter(rh.handleRow)
 		distSQLReceiver := sql.MakeDistSQLReceiver(
@@ -785,7 +787,7 @@ type rowHandler struct {
 	job                   *jobs.Job
 	frontierUpdates       chan hlc.Timestamp
 
-	rangeStats rangeStatsByProcessorID
+	rangeStats replicationutils.AggregateRangeStatsCollector
 
 	lastPartitionUpdate time.Time
 

pkg/crosscluster/logical/logical_replication_writer_processor.go

@@ -18,6 +18,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/ccl/changefeedccl/cdcevent"
 	"github.com/cockroachdb/cockroach/pkg/clusterversion"
 	"github.com/cockroachdb/cockroach/pkg/crosscluster"
+	"github.com/cockroachdb/cockroach/pkg/crosscluster/replicationutils"
 	"github.com/cockroachdb/cockroach/pkg/crosscluster/streamclient"
 	"github.com/cockroachdb/cockroach/pkg/jobs/jobspb"
 	"github.com/cockroachdb/cockroach/pkg/keys"
@@ -50,7 +51,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/util/tracing"
 	"github.com/cockroachdb/errors"
 	"github.com/cockroachdb/logtags"
-	pbtypes "github.com/gogo/protobuf/types"
 )
 
 var logicalReplicationWriterResultType = []*types.T{
@@ -395,7 +395,9 @@ func (lrw *logicalReplicationWriterProcessor) Next() (
 			lrw.FlowCtx.NodeID.SQLInstanceID(), lrw.FlowCtx.ID, lrw.agg)
 
 	case stats := <-lrw.rangeStatsCh:
-		meta, err := lrw.newRangeStatsProgressMeta(stats)
+		meta, err := replicationutils.StreamRangeStatsToProgressMeta(
+			lrw.FlowCtx, lrw.ProcessorID, stats,
+		)
 		if err != nil {
 			lrw.MoveToDrainingAndLogError(err)
 			return nil, lrw.DrainHelper()
@@ -579,23 +581,6 @@ func (lrw *logicalReplicationWriterProcessor) rangeStats(
 	}
 }
 
-func (lrw *logicalReplicationWriterProcessor) newRangeStatsProgressMeta(
-	stats *streampb.StreamEvent_RangeStats,
-) (*execinfrapb.ProducerMetadata, error) {
-	asAny, err := pbtypes.MarshalAny(stats)
-	if err != nil {
-		return nil, errors.Wrap(err, "unable to convert stats into any proto")
-	}
-	return &execinfrapb.ProducerMetadata{
-		BulkProcessorProgress: &execinfrapb.RemoteProducerMetadata_BulkProcessorProgress{
-			NodeID:          lrw.FlowCtx.NodeID.SQLInstanceID(),
-			FlowID:          lrw.FlowCtx.ID,
-			ProcessorID:     lrw.ProcessorID,
-			ProgressDetails: *asAny,
-		},
-	}, nil
-}
-
 func (lrw *logicalReplicationWriterProcessor) checkpoint(
 	ctx context.Context, resolvedSpans []jobspb.ResolvedSpan,
 ) error {

pkg/crosscluster/logical/range_stats.go

@@ -2,7 +2,10 @@ load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
 go_library(
     name = "replicationutils",
-    srcs = ["utils.go"],
+    srcs = [
+        "stats.go",
+        "utils.go",
+    ],
     importpath = "github.com/cockroachdb/cockroach/pkg/crosscluster/replicationutils",
     visibility = ["//visibility:public"],
     deps = [
@@ -18,6 +21,8 @@ go_library(
         "//pkg/sql/catalog/descs",
         "//pkg/sql/catalog/resolver",
         "//pkg/sql/catalog/tabledesc",
+        "//pkg/sql/execinfra",
+        "//pkg/sql/execinfrapb",
         "//pkg/sql/isql",
         "//pkg/sql/parser",
         "//pkg/sql/privilege",
@@ -28,18 +33,24 @@ go_library(
         "//pkg/testutils/fingerprintutils",
         "//pkg/util/ctxgroup",
         "//pkg/util/hlc",
+        "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
+        "@com_github_gogo_protobuf//types",
     ],
 )
 
 go_test(
     name = "replicationutils_test",
-    srcs = ["utils_test.go"],
+    srcs = [
+        "stats_test.go",
+        "utils_test.go",
+    ],
     embed = [":replicationutils"],
     deps = [
         "//pkg/clusterversion",
         "//pkg/kv/kvpb",
+        "//pkg/repstream/streampb",
         "//pkg/roachpb",
         "//pkg/settings/cluster",
         "//pkg/storage",

pkg/crosscluster/replicationutils/BUILD.bazel

@@ -0,0 +1,101 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package replicationutils
+
+import (
+	"fmt"
+
+	"github.com/cockroachdb/cockroach/pkg/repstream/streampb"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfra"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
+	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
+	"github.com/cockroachdb/errors"
+	pbtypes "github.com/gogo/protobuf/types"
+)
+
+// StreamRangeStatsToProgressMeta converts a range statistics from a rangefeed
+// StreamEvent and converts it to a ProducerMetadata that can be passed through
+// the DistSQL pipeline.
+func StreamRangeStatsToProgressMeta(
+	flowCtx *execinfra.FlowCtx, procID int32, stats *streampb.StreamEvent_RangeStats,
+) (*execinfrapb.ProducerMetadata, error) {
+	asAny, err := pbtypes.MarshalAny(stats)
+	if err != nil {
+		return nil, errors.Wrap(err, "unable to convert range stats to any proto")
+	}
+	return &execinfrapb.ProducerMetadata{
+		BulkProcessorProgress: &execinfrapb.RemoteProducerMetadata_BulkProcessorProgress{
+			NodeID:          flowCtx.NodeID.SQLInstanceID(),
+			FlowID:          flowCtx.ID,
+			ProcessorID:     procID,
+			ProgressDetails: *asAny,
+		},
+	}, nil
+}
+
+// AggregateRangeStatsCollector collects rangefeed StreamEvent range stats from
+// multiple processors and aggregates them into single metrics.
+type AggregateRangeStatsCollector struct {
+	mu             syncutil.Mutex
+	stats          map[int32]*streampb.StreamEvent_RangeStats
+	processorCount int
+}
+
+func NewAggregateRangeStatsCollector(processorCount int) AggregateRangeStatsCollector {
+	return AggregateRangeStatsCollector{
+		stats:          make(map[int32]*streampb.StreamEvent_RangeStats),
+		processorCount: processorCount,
+	}
+}
+
+// Add adds range states from a processor to the collector.
+func (r *AggregateRangeStatsCollector) Add(
+	processorID int32, stats *streampb.StreamEvent_RangeStats,
+) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	r.stats[processorID] = stats
+}
+
+// RollupStats aggregates the collected stats and returns the total range stats,
+// the fraction of ranges that have reached the steady state, and a human
+// readable status message.
+func (r *AggregateRangeStatsCollector) RollupStats() (
+	streampb.StreamEvent_RangeStats,
+	float32,
+	string,
+) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	var total streampb.StreamEvent_RangeStats
+	for _, producerStats := range r.stats {
+		total.RangeCount += producerStats.RangeCount
+		total.ScanningRangeCount += producerStats.ScanningRangeCount
+		total.LaggingRangeCount += producerStats.LaggingRangeCount
+	}
+	initialScanComplete := total.ScanningRangeCount == 0
+	incompleteCount := total.ScanningRangeCount
+	if initialScanComplete {
+		incompleteCount = total.LaggingRangeCount
+	}
+
+	fractionCompleted := max(
+		// Use a tiny fraction completed to start with a nearly empty
+		// progress bar until we get the first batch of range stats.
+		float32(0.0001),
+		(float32(total.RangeCount-incompleteCount) / float32(total.RangeCount)))
+
+	if len(r.stats) != r.processorCount || total.RangeCount == 0 {
+		return streampb.StreamEvent_RangeStats{}, 0, fmt.Sprintf("starting streams (%d out of %d)", len(r.stats), r.processorCount)
+	}
+	if !initialScanComplete {
+		return total, fractionCompleted, fmt.Sprintf("initial scan on %d out of %d ranges", total.ScanningRangeCount, total.RangeCount)
+	}
+	if total.LaggingRangeCount != 0 {
+		return total, fractionCompleted, fmt.Sprintf("catching up on %d out of %d ranges", total.LaggingRangeCount, total.RangeCount)
+	}
+	return total, 1, ""
+}

pkg/crosscluster/replicationutils/stats.go

@@ -1,9 +1,9 @@
-// Copyright 2024 The Cockroach Authors.
+// Copyright 2025 The Cockroach Authors.
 //
 // Use of this software is governed by the CockroachDB Software License
 // included in the /LICENSE file.
 
-package logical
+package replicationutils
 
 import (
 	"testing"
@@ -72,7 +72,7 @@ func TestRangeStats(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			r := newRangeStatsCollector(3)
+			r := NewAggregateRangeStatsCollector(3)
 			for id, stats := range tc.inputStats {
 				r.Add(id, stats)
 			}
@@ -82,15 +82,14 @@ func TestRangeStats(t *testing.T) {
 			require.Equal(t, tc.fraction, fraction)
 			require.Equal(t, tc.expectedMsg, msg)
 
-			rInitializing := newRangeStatsCollector(4)
+			rInitializing := NewAggregateRangeStatsCollector(4)
 			for id, stats := range tc.inputStats {
 				r.Add(id, stats)
 			}
 			total, fraction, msg = rInitializing.RollupStats()
 			require.Equal(t, total, streampb.StreamEvent_RangeStats{})
 			require.Equal(t, float32(0), fraction)
 			require.Contains(t, msg, "starting streams")
-
 		})
 	}
 }