Merge #141736 #143750

141736: server,sqlstats: add ssRemote.SQLStats r=kyle-a-wong a=kyle-a-wong

ssRemote.SQLStats implements SSDrainer and does a cluster wide DrainStats. This is done using a new gRPC via /_status/drainstats.

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.

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

143750: schemachanger/scplan: add deprules to eval zone cfgs in seqNum order r=annrpom a=annrpom

This patch adds deprules to ensure that zone configs are depended by
each other in increasing`seqNum` order.

Epic: none
Fixes: #129482

Release note: None

Co-authored-by: Kyle Wong <[email protected]>
Co-authored-by: Annie Pompa <[email protected]>

Author: 3 people

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) = {