cmd/roachtest: add changefeed roachtests exercising execution_locality

in multi-region clusters

Previously, we were not testing to verify that a changefeed with
execution_locality enabled creates an acceptable plan and is able to
progress. Prior to 24.1, the changefeed would only put aggregators on
nodes that were leaseholders for the data. If execution locality is
enabled to restrict the changefeed to nodes without leaseholders for
that data, then all the ranges would be planned on the gateway node.
This roachtest verifies that aggregators are spread across nodes within
a region (no one node has a majority of changefeed aggregators).

Fixes: #153736
Epic: CRDB-38755

Release note: None

Author: log-head

pkg/cmd/roachtest/tests/BUILD.bazel

@@ -276,6 +276,7 @@ go_library(
         "//pkg/settings/cluster",
         "//pkg/sql",
         "//pkg/sql/catalog/catpb",
+        "//pkg/sql/execinfrapb",
         "//pkg/sql/pgwire/pgcode",
         "//pkg/sql/pgwire/pgerror",
         "//pkg/sql/sem/tree",

pkg/cmd/roachtest/tests/cdc.go

@@ -59,6 +59,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachprod/prometheus"
 	"github.com/cockroachdb/cockroach/pkg/roachprod/vm"
 	roachprodaws "github.com/cockroachdb/cockroach/pkg/roachprod/vm/aws"
+	"github.com/cockroachdb/cockroach/pkg/sql/execinfrapb"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
@@ -1899,7 +1900,182 @@ func configureDBForMultiTablePTSBenchmark(db *gosql.DB) error {
 	return nil
 }
 
+func getDiagramProcessors(ctx context.Context, db *gosql.DB) ([]any, error) {
+	var diagramURL string
+	diagramQuery := `SELECT value 
+	FROM system.job_info ji
+	INNER JOIN system.jobs j ON ji.job_id = j.id
+	WHERE j.job_type = 'CHANGEFEED' AND ji.info_key LIKE '~dsp-diag-url-%'`
+	if err := db.QueryRowContext(ctx, diagramQuery).Scan(&diagramURL); err != nil {
+		return nil, err
+	}
+	diagram, err := execinfrapb.FromURL(diagramURL)
+	if err != nil {
+		return nil, err
+	}
+	diagramJSON, err := json.Marshal(diagram)
+	if err != nil {
+		return nil, err
+	}
+	var flow map[string]any
+	if err := json.Unmarshal(diagramJSON, &flow); err != nil {
+		return nil, err
+	}
+	processors, ok := flow["processors"].([]any)
+	if !ok {
+		return nil, fmt.Errorf("processors not found in flow")
+	}
+	return processors, nil
+}
+
+type ChangefeedDistribution struct {
+	NodeToSpansWatched map[int]int
+	ZoneToSpansWatched map[string]int
+	TotalSpansWatched  int
+	TotalAggregators   int
+	TotalLeaseHolders  int
+	TotalRanges        int
+	NodeToZone         map[int]string
+}
+
+func getChangefeedDistribution(
+	processors []any, nodeToZone map[int]string, t test.Test,
+) ChangefeedDistribution {
+	changefeedDistribution := ChangefeedDistribution{
+		NodeToSpansWatched: make(map[int]int),
+		ZoneToSpansWatched: make(map[string]int),
+		TotalSpansWatched:  0,
+		TotalAggregators:   0,
+		TotalLeaseHolders:  0,
+		TotalRanges:        0,
+		NodeToZone:         nodeToZone,
+	}
+	for _, p := range processors {
+		procMap, ok := p.(map[string]any)
+		if !ok {
+			t.Fatalf("processor not a map")
+		}
+		nodeIdx, ok := procMap["nodeIdx"].(float64)
+		require.True(t, ok, "node idx not found in processor")
+		core, ok := procMap["core"].(map[string]any)
+		require.True(t, ok, "core not found in processor")
+		title, ok := core["title"].(string)
+		require.True(t, ok, "title not found in core")
+		if strings.HasPrefix(title, "ChangeAggregator") {
+			changefeedDistribution.TotalAggregators++
+			details := core["details"].([]any)
+			for _, detail := range details {
+				if strings.HasPrefix(detail.(string), "Watches") {
+					re := regexp.MustCompile(`Watches \[(\d+)\]:`)
+					matches := re.FindStringSubmatch(detail.(string))
+					if len(matches) > 1 {
+						numWatches, err := strconv.Atoi(matches[1])
+						require.NoError(t, err)
+						changefeedDistribution.NodeToSpansWatched[int(nodeIdx)] += numWatches
+						changefeedDistribution.TotalSpansWatched += numWatches
+						changefeedDistribution.ZoneToSpansWatched[changefeedDistribution.NodeToZone[int(nodeIdx)]] += numWatches
+
+					}
+				}
+			}
+		}
+	}
+	return changefeedDistribution
+}
+
+func veryifyLeaseHolderDistribution(
+	db *gosql.DB, t test.Test, nodeToZone map[int]string,
+) map[string]int {
+	var rows *gosql.Rows
+	// Get lease holders for all ranges in tpcc database.
+	leaseHolderQuery := `SELECT r.start_pretty, r.replicas, r.replica_localities, r.lease_holder 
+	FROM crdb_internal.ranges r 
+	JOIN crdb_internal.tables t ON r.start_pretty like concat('/Table/', t.table_id::STRING,'%')
+	WHERE t.database_name = 'tpcc'`
+	rows, err := db.Query(leaseHolderQuery)
+	zoneToLeaseHolderCount := make(map[string]int)
+	require.NoError(t, err)
+	defer rows.Close()
+	for rows.Next() {
+		var startKeyPretty string
+		var replicas []uint8
+		var replicaLocalities []uint8
+		var leaseHolder int
+		require.NoError(t, rows.Scan(&startKeyPretty, &replicas, &replicaLocalities, &leaseHolder))
+		for indx := range replicas {
+			require.NotEqual(t, replicas[indx], 0)
+			replicas[indx]--
+		}
+		leaseHolder--
+		zoneToLeaseHolderCount[nodeToZone[leaseHolder]]++
+	}
+	return zoneToLeaseHolderCount
+}
+
 func registerCDC(r registry.Registry) {
+	r.Add(registry.TestSpec{
+		// This test
+		// 1. Creates a cluster with 3 nodes each in us-east and us-west
+		// 2. Runs a tpcc workload, then sets tpcc database to primary region us-west
+		// 3. Creates a changefeed with execution locality set to us-east
+		// 4. Gets the changefeed diagram and creates mappings
+
+		// This test is used to verify that ranges are evenly distributed across
+		// change aggregators in the execution_locality region while targeting tables
+		// whose primary region is different. In issue #2955, in that scenario,
+		// a single change aggregator (on the gateway node) would watch all the ranges.
+		// The above scenario occured with the older bin-packing oracle rather than
+		// the bulk oracle.
+		Name:             "cdc/multi-region-execution-locality-tpcc",
+		Owner:            registry.OwnerCDC,
+		Cluster:          r.MakeClusterSpec(7, spec.Geo(), spec.GatherCores(), spec.GCEZones("us-east1-b,us-west1-b")),
+		CompatibleClouds: registry.OnlyGCE,
+		Suites:           registry.Suites(),
+		Run: func(ctx context.Context, t test.Test, c cluster.Cluster) {
+			nodeToZone := map[int]string{
+				0: "us-east1-b",
+				1: "us-east1-b",
+				2: "us-east1-b",
+				3: "us-west1-b",
+				4: "us-west1-b",
+				5: "us-west1-b",
+			}
+			ct := newCDCTester(ctx, t, c)
+			defer ct.Close()
+
+			ct.runTPCCWorkload(tpccArgs{warehouses: 100})
+
+			var err error
+			_, err = ct.DB().Exec("ALTER DATABASE tpcc SET PRIMARY REGION 'us-west1'")
+			require.NoError(t, err)
+
+			feed := ct.newChangefeed(feedArgs{
+				sinkType: cloudStorageSink,
+				targets:  allTpccTargets,
+				opts: map[string]string{
+					"execution_locality": "'region=us-east1'",
+				},
+			})
+			ct.waitForWorkload()
+			feed.waitForCompletion()
+
+			processors, err := getDiagramProcessors(ctx, ct.DB())
+			require.NoError(t, err)
+
+			changefeedDistribution := getChangefeedDistribution(processors, nodeToZone, t)
+			require.Greater(t, changefeedDistribution.TotalAggregators, 1)
+			for nodeIdx, spansWatched := range changefeedDistribution.NodeToSpansWatched {
+				require.LessOrEqual(t, spansWatched, changefeedDistribution.TotalSpansWatched/2, "nodeIdx %d watched %d spans, total spans watched %d", nodeIdx, spansWatched, changefeedDistribution.TotalSpansWatched)
+			}
+			require.Equal(t, 1, len(changefeedDistribution.ZoneToSpansWatched))
+			require.Equal(t, changefeedDistribution.ZoneToSpansWatched["us-east1-b"], changefeedDistribution.TotalSpansWatched)
+			zoneToLeaseHolderCount := veryifyLeaseHolderDistribution(ct.DB(), t, nodeToZone)
+			// Majority of lease holders should be in us-west1-b. Some may not, but most should.
+			if zoneToLeaseHolderCount["us-east1-b"] != 0 {
+				require.Greater(t, zoneToLeaseHolderCount["us-west1-b"]/zoneToLeaseHolderCount["us-east1-b"], 10)
+			}
+		},
+	})
 	r.Add(registry.TestSpec{
 		Name:      "cdc/initial-scan-only",
 		Owner:     registry.OwnerCDC,