server,sqlstats: add ssRemote.SQLStats

ssRemote.SQLStats implements SSDrainer and does a
cluster wide DrainStats. This is done using a new
gRPC DrainSQLStats API.

This implementation is intended to be used in the
new sql activity flush job to do the flush operation
on a single node. By collecting all the stats and
merging them on a single node, it reduces the amount
of sql upserts done compared to individual node
flushing.

This commit also adds a new iterateNodesExt to allow
more configurability for the iterateNodes operation.
For now, it adds support for configuring the
concurrency level of cluster wide fanouts. The new
DrainSQLStats sets this concurrency to 4 to help
reduce memory usage for cluster wide fanouts of large
clusters.

Resolves: #142364
Epic: CRDB-45771
Release note: None

Author: kyle-a-wong

docs/generated/http/full.md

@@ -4854,6 +4854,52 @@ Response object returned by ResetSQLStats.
 
 
 
+## DrainSqlStats
+
+
+
+
+
+Support status: [reserved](#support-status)
+
+#### Request Parameters
+
+
+
+
+
+
+
+| Field | Type | Label | Description | Support status |
+| ----- | ---- | ----- | ----------- | -------------- |
+| node_id | [string](#cockroach.server.serverpb.DrainSqlStatsRequest-string) |  |  | [reserved](#support-status) |
+
+
+
+
+
+
+
+#### Response Parameters
+
+
+
+
+
+
+
+| Field | Type | Label | Description | Support status |
+| ----- | ---- | ----- | ----------- | -------------- |
+| statements | [cockroach.sql.CollectedStatementStatistics](#cockroach.server.serverpb.DrainStatsResponse-cockroach.sql.CollectedStatementStatistics) | repeated |  | [reserved](#support-status) |
+| transactions | [cockroach.sql.CollectedTransactionStatistics](#cockroach.server.serverpb.DrainStatsResponse-cockroach.sql.CollectedTransactionStatistics) | repeated |  | [reserved](#support-status) |
+| fingerprint_count | [int64](#cockroach.server.serverpb.DrainStatsResponse-int64) |  |  | [reserved](#support-status) |
+
+
+
+
+
+
+
 ## IndexUsageStatistics
 
 `GET /_status/indexusagestatistics`

pkg/BUILD.bazel

@@ -639,6 +639,7 @@ ALL_TESTS = [
     "//pkg/sql/sqlstats/persistedsqlstats:persistedsqlstats_test",
     "//pkg/sql/sqlstats/sslocal:sslocal_test",
     "//pkg/sql/sqlstats/ssmemstorage:ssmemstorage_test",
+    "//pkg/sql/sqlstats/ssremote:ssremote_test",
     "//pkg/sql/sqltestutils:sqltestutils_test",
     "//pkg/sql/stats:stats_test",
     "//pkg/sql/stmtdiagnostics:stmtdiagnostics_test",
@@ -2325,6 +2326,8 @@ GO_TARGETS = [
     "//pkg/sql/sqlstats/sslocal:sslocal_test",
     "//pkg/sql/sqlstats/ssmemstorage:ssmemstorage",
     "//pkg/sql/sqlstats/ssmemstorage:ssmemstorage_test",
+    "//pkg/sql/sqlstats/ssremote:ssremote",
+    "//pkg/sql/sqlstats/ssremote:ssremote_test",
     "//pkg/sql/sqlstats:sqlstats",
     "//pkg/sql/sqltelemetry:sqltelemetry",
     "//pkg/sql/sqltestutils:sqltestutils",

pkg/server/BUILD.bazel

@@ -380,6 +380,7 @@ go_library(
         "@org_golang_google_grpc//codes",
         "@org_golang_google_grpc//metadata",
         "@org_golang_google_grpc//status",
+        "@org_golang_x_exp//maps",
     ] + select({
         "@io_bazel_rules_go//go/platform:aix": [
             "@org_golang_x_sys//unix",

pkg/server/application_api/sql_stats_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/security/username"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/server/apiconstants"
+	"github.com/cockroachdb/cockroach/pkg/server/authserver"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/server/srvtestutils"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -1540,6 +1541,93 @@ func TestCombinedStatementUsesCorrectSourceTable(t *testing.T) {
 	}
 }
 
+func TestDrainSqlStats(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	appName := "drain_stats_app"
+	testCluster := serverutils.StartCluster(t, 3, base.TestClusterArgs{})
+	ctx := context.Background()
+	defer testCluster.Stopper().Stop(ctx)
+
+	conn1 := sqlutils.MakeSQLRunner(testCluster.ServerConn(0))
+	conn2 := sqlutils.MakeSQLRunner(testCluster.ServerConn(1))
+	conn3 := sqlutils.MakeSQLRunner(testCluster.ServerConn(2))
+
+	for _, conn := range []*sqlutils.SQLRunner{conn1, conn2, conn3} {
+		conn.Exec(t, fmt.Sprintf("SET application_name = '%s'", appName))
+		conn.Exec(t, "SELECT 1")
+	}
+
+	statusServer := testCluster.Server(0).GetStatusClient(t)
+	resp, err := statusServer.DrainSqlStats(ctx, &serverpb.DrainSqlStatsRequest{})
+	require.NoError(t, err)
+	checkFingerprintCount(t, resp)
+	stmts, txns := filterStatementStatsByAppName(resp.Statements, resp.Transactions, appName)
+	require.Len(t, stmts, 1)
+	require.Equal(t, int64(3), stmts[0].Stats.Count)
+	require.Equal(t, "SELECT _", stmts[0].Key.Query)
+	require.Len(t, txns, 1)
+	require.Equal(t, int64(3), txns[0].Stats.Count)
+	require.Len(t, txns[0].StatementFingerprintIDs, 1)
+	require.Equal(t, stmts[0].ID, txns[0].StatementFingerprintIDs[0])
+
+	// Check that the stats are cleared.
+	resp, err = statusServer.DrainSqlStats(ctx, &serverpb.DrainSqlStatsRequest{})
+	require.NoError(t, err)
+	stmts, txns = filterStatementStatsByAppName(resp.Statements, resp.Transactions, appName)
+	require.Empty(t, stmts)
+	require.Empty(t, txns)
+}
+
+func TestDrainSqlStats_partialOutage(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	appName := "drain_stats_app"
+	testCluster := serverutils.StartCluster(t, 3, base.TestClusterArgs{})
+	ctx := context.Background()
+	defer testCluster.Stopper().Stop(ctx)
+
+	conn1 := sqlutils.MakeSQLRunner(testCluster.ServerConn(0))
+	conn2 := sqlutils.MakeSQLRunner(testCluster.ServerConn(1))
+	conn3 := sqlutils.MakeSQLRunner(testCluster.ServerConn(2))
+
+	for _, conn := range []*sqlutils.SQLRunner{conn1, conn2, conn3} {
+		conn.Exec(t, fmt.Sprintf("SET application_name = '%s'", appName))
+		conn.Exec(t, "SELECT 1")
+	}
+
+	// Stop server 2 to simulate a partial outage
+	testCluster.StopServer(2)
+	statusServer := testCluster.Server(0).GetStatusClient(t)
+	resp, err := statusServer.DrainSqlStats(ctx, &serverpb.DrainSqlStatsRequest{})
+	require.NoError(t, err)
+	checkFingerprintCount(t, resp)
+	stmts, txns := filterStatementStatsByAppName(resp.Statements, resp.Transactions, appName)
+	require.Len(t, stmts, 1)
+	require.Equal(t, int64(2), stmts[0].Stats.Count)
+	require.Equal(t, "SELECT _", stmts[0].Key.Query)
+	require.Len(t, txns, 1)
+	require.Equal(t, int64(2), txns[0].Stats.Count)
+}
+
+func TestDrainSqlStatsPermissionDenied(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+	ts := serverutils.StartServerOnly(t, base.TestServerArgs{})
+	ctx := context.Background()
+	nonRootUser := apiconstants.TestingUserNameNoAdmin()
+	sqlutils.MakeSQLRunner(ts.SQLConn(t)).Exec(t, fmt.Sprintf("CREATE USER IF NOT EXISTS %s", nonRootUser))
+	ctx = authserver.ContextWithHTTPAuthInfo(ctx, nonRootUser.Normalized(), 1)
+	ctx = authserver.ForwardHTTPAuthInfoToRPCCalls(ctx, nil)
+
+	statusClient := ts.GetStatusClient(t)
+	defer ts.Stopper().Stop(ctx)
+	_, err := statusClient.DrainSqlStats(ctx, &serverpb.DrainSqlStatsRequest{})
+
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "user does not have admin role")
+}
+
 func createStmtFetchMode(
 	sort serverpb.StatsSortOptions,
 ) *serverpb.CombinedStatementsStatsRequest_FetchMode {
@@ -1860,3 +1948,44 @@ VALUES (
 	}
 	db.Exec(t, query, args...)
 }
+
+func filterStatementStatsByAppName(
+	statements []*appstatspb.CollectedStatementStatistics,
+	transactions []*appstatspb.CollectedTransactionStatistics,
+	appName string,
+) ([]*appstatspb.CollectedStatementStatistics, []*appstatspb.CollectedTransactionStatistics) {
+	var filteredStatements []*appstatspb.CollectedStatementStatistics
+	var filteredTransactions []*appstatspb.CollectedTransactionStatistics
+
+	for _, stmt := range statements {
+		if stmt.Key.App == appName {
+			filteredStatements = append(filteredStatements, stmt)
+		}
+	}
+
+	for _, txn := range transactions {
+		if txn.App == appName {
+			filteredTransactions = append(filteredTransactions, txn)
+		}
+	}
+
+	return filteredStatements, filteredTransactions
+}
+
+func checkFingerprintCount(t *testing.T, resp *serverpb.DrainStatsResponse) {
+	stmtFingerprints := make(map[appstatspb.StmtFingerprintID]struct{})
+	txnFingerprints := make(map[appstatspb.TransactionFingerprintID]struct{})
+	for _, stmt := range resp.Statements {
+		if _, ok := stmtFingerprints[stmt.ID]; !ok {
+			stmtFingerprints[stmt.ID] = struct{}{}
+		}
+	}
+
+	for _, txn := range resp.Transactions {
+		if _, ok := txnFingerprints[txn.TransactionFingerprintID]; !ok {
+			txnFingerprints[txn.TransactionFingerprintID] = struct{}{}
+		}
+	}
+	actualFpCount := len(stmtFingerprints) + len(txnFingerprints)
+	require.Equal(t, resp.FingerprintCount, int64(actualFpCount))
+}

pkg/server/serverpb/status.go

@@ -46,6 +46,7 @@ type SQLStatusServer interface {
 	TenantServiceStatus(context.Context, *TenantServiceStatusRequest) (*TenantServiceStatusResponse, error)
 	UpdateTableMetadataCache(context.Context, *UpdateTableMetadataCacheRequest) (*UpdateTableMetadataCacheResponse, error)
 	GetUpdateTableMetadataCacheSignal() chan struct{}
+	DrainSqlStats(context.Context, *DrainSqlStatsRequest) (*DrainStatsResponse, error)
 }
 
 // OptionalNodesStatusServer is a StatusServer that is only optionally present

pkg/server/serverpb/status.proto

@@ -1599,6 +1599,16 @@ message StatementsRequest {
   FetchMode fetch_mode = 5;
 }
 
+message DrainSqlStatsRequest {
+  string node_id = 1 [(gogoproto.customname) = "NodeID"];
+}
+
+message DrainStatsResponse {
+  repeated cockroach.sql.CollectedStatementStatistics statements = 1;
+  repeated cockroach.sql.CollectedTransactionStatistics transactions = 2;
+  int64 fingerprint_count = 3;
+}
+
 message StatementsResponse {
   message ExtendedStatementStatisticsKey {
     cockroach.sql.StatementStatisticsKey key_data = 1 [(gogoproto.nullable) = false];
@@ -2567,6 +2577,8 @@ service Status {
       body: "*"
     };
   }
+  rpc DrainSqlStats(DrainSqlStatsRequest) returns (DrainStatsResponse) {
+  }
 
   rpc IndexUsageStatistics(IndexUsageStatisticsRequest) returns (IndexUsageStatisticsResponse) {
     option (google.api.http) = {