roachtest: add operation to query crdb_internal.probe_ranges

Since SRE uses crdb_internal.probe_ranges to test for prod
cluster health, we would like to add this as a roach operation
to make the DRT cluster as realistic as possible, and test for
potential issues with crdb_internal.probe_ranges.

Fixes: cockroachdb#102034
Release note: None
Epic: None

Author: noahstho

pkg/cmd/roachtest/operations/BUILD.bazel

@@ -17,6 +17,7 @@ go_library(
         "network_partition.go",
         "node_kill.go",
         "pause_job.go",
+        "probe_ranges.go",
         "register.go",
         "resize.go",
         "session_vars.go",

pkg/cmd/roachtest/operations/probe_ranges.go

@@ -0,0 +1,89 @@
+// Copyright 2024 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package operations
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/cluster"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operation"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/operations/helpers"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/option"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/registry"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachtest/roachtestflags"
+	"github.com/cockroachdb/cockroach/pkg/util/randutil"
+)
+
+func runProbeRanges(
+	ctx context.Context, o operation.Operation, c cluster.Cluster,
+) registry.OperationCleanup {
+	rng, _ := randutil.NewPseudoRand()
+
+	nodes := c.All()
+	nid := nodes[rng.Intn(len(nodes))]
+
+	conn := c.Conn(ctx, o.L(), nid, option.VirtualClusterName(roachtestflags.VirtualCluster))
+	defer conn.Close()
+
+	dbName := helpers.PickRandomDB(ctx, o, conn, helpers.SystemDBs)
+	tableName := helpers.PickRandomTable(ctx, o, conn, dbName)
+	sid := helpers.PickRandomStore(ctx, o, conn, nid)
+
+	compactionStmt := fmt.Sprintf(`SELECT crdb_internal.compact_engine_span(
+				%d, %d,
+				(SELECT raw_start_key FROM [SHOW RANGES FROM TABLE %s.%s WITH KEYS] LIMIT 1),
+				(SELECT raw_end_key FROM [SHOW RANGES FROM TABLE %s.%s WITH KEYS] LIMIT 1))`,
+		nid, sid, dbName, tableName, dbName, tableName)
+	o.Status(fmt.Sprintf("compacting a range for table %s.%s in n%d, s%d",
+		dbName, tableName, nid, sid))
+	_, err := conn.ExecContext(ctx, compactionStmt)
+	if err != nil {
+		o.Fatal(err)
+	}
+	return nil
+}
+
+// runKVProber tests that accessing crdb_internal_probe_ranges works across mixed version clusters.
+func runKVProberMixedVersion(ctx context.Context, o operation.Operation, c cluster.Cluster) registry.OperationCleanup {
+	rng, _ := randutil.NewPseudoRand()
+
+	nodes := c.All()
+	nid := nodes[rng.Intn(len(nodes))]
+
+	conn := c.Conn(ctx, o.L(), nid, option.VirtualClusterName(roachtestflags.VirtualCluster))
+	defer conn.Close()
+
+	kvProbeRead := fmt.Sprintf(`select count(error) from crdb_internal.probe_ranges(interval '1s', 'read') where node_id=%d"`, nid)
+	kvProbeWrite := fmt.Sprintf(`select count(error) from crdb_internal.probe_ranges(interval '1s', 'write') where node_id=%d"`, nid)
+
+	o.Status("probing reads to crdb_internal.probe_ranges")
+	_, read_err := conn.ExecContext(ctx, kvProbeRead)
+	if read_err != nil {
+		o.Fatal(read_err)
+	}
+
+	o.Status("probing writes to crdb_internal.probe_ranges")
+	_, write_err := conn.ExecContext(ctx, kvProbeWrite)
+	if write_err != nil {
+		o.Fatal(write_err)
+	}
+
+	return nil
+}
+
+func registerProbeRanges(r registry.Registry) {
+	r.AddOperation(registry.OperationSpec{
+		Name:               "manual-compaction",
+		Owner:              registry.OwnerStorage,
+		Timeout:            24 * time.Hour,
+		CompatibleClouds:   registry.OnlyGCE,
+		CanRunConcurrently: registry.OperationCanRunConcurrently,
+		Dependencies:       []registry.OperationDependency{registry.OperationRequiresZeroUnavailableRanges},
+		Run:                runProbeRanges,
+	})
+}