From 9e0dc9359db8cc9d5856e8a087456a703fc49701 Mon Sep 17 00:00:00 2001
From: Cameron Nunez <cameron@cockroachlabs.com>
Date: Thu, 10 Feb 2022 16:10:06 -0500
Subject: [PATCH] server: flush SQL stats during drain

Previously, SQL stats would be lost when a node drains. Now a drain
triggers a flush of the SQL stats into the statement statistics
system table while the SQL layer is being drained.

Release note (cli change): a drain of node now ensures that
SQL statistics are not lost during the process; they are now
preserved in the statement statistics system table.
---
 pkg/server/BUILD.bazel   |  1 +
 pkg/server/drain.go      |  4 +++
 pkg/server/drain_test.go | 64 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 69 insertions(+)

diff --git a/pkg/server/BUILD.bazel b/pkg/server/BUILD.bazel
index e2f70170086b..475d2d535a99 100644
--- a/pkg/server/BUILD.bazel
+++ b/pkg/server/BUILD.bazel
@@ -177,6 +177,7 @@ go_library(
         "//pkg/sql/sqlliveness",
         "//pkg/sql/sqlliveness/slprovider",
         "//pkg/sql/sqlstats",
+        "//pkg/sql/sqlstats/persistedsqlstats",
         "//pkg/sql/sqlstats/persistedsqlstats/sqlstatsutil",
         "//pkg/sql/sqlutil",
         "//pkg/sql/stats",
diff --git a/pkg/server/drain.go b/pkg/server/drain.go
index ab258f6c1106..87376ece178d 100644
--- a/pkg/server/drain.go
+++ b/pkg/server/drain.go
@@ -19,6 +19,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/liveness"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
+	"github.com/cockroachdb/cockroach/pkg/sql/sqlstats/persistedsqlstats"
 	"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/stop"
@@ -318,6 +319,9 @@ func (s *drainServer) drainClients(
 	// Stop ongoing SQL execution up to the queryWait timeout.
 	s.sqlServer.distSQLServer.Drain(ctx, queryMaxWait, reporter)
 
+	// Flush in-memory SQL stats into the statement stats system table.
+	s.sqlServer.pgServer.SQLServer.GetSQLStatsProvider().(*persistedsqlstats.PersistedSQLStats).Flush(ctx)
+
 	// Drain the SQL leases. This must be done after the pgServer has
 	// given sessions a chance to finish ongoing work.
 	s.sqlServer.leaseMgr.SetDraining(ctx, true /* drain */, reporter)
diff --git a/pkg/server/drain_test.go b/pkg/server/drain_test.go
index 91a9751ad0b8..12dd1e6d1bfc 100644
--- a/pkg/server/drain_test.go
+++ b/pkg/server/drain_test.go
@@ -17,16 +17,20 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/server/serverpb"
+	"github.com/cockroachdb/cockroach/pkg/sql"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
+	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
 	"github.com/cockroachdb/cockroach/pkg/util/grpcutil"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/errors"
 	"github.com/kr/pretty"
+	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 )
 
@@ -93,6 +97,66 @@ func doTestDrain(tt *testing.T) {
 	})
 }
 
+func TestEnsureSQLStatsAreFlushedDuringDrain(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+
+	var drainSleepCallCount = 0
+	drainCtx := newTestDrainContext(t, &drainSleepCallCount)
+	defer drainCtx.Close()
+
+	var (
+		ts    = drainCtx.tc.Server(0).SQLServer().(*sql.Server)
+		sqlDB = sqlutils.MakeSQLRunner(drainCtx.tc.ServerConn(0))
+	)
+
+	// Issue queries to be registered in stats.
+	sqlDB.Exec(t, `
+CREATE DATABASE t;
+CREATE TABLE t.test (x INT PRIMARY KEY);
+INSERT INTO t.test VALUES (1);
+INSERT INTO t.test VALUES (2);
+INSERT INTO t.test VALUES (3);
+`)
+
+	// Find the in-memory stats for the queries.
+	stats, err := ts.GetScrubbedStmtStats(ctx)
+	require.NoError(t, err)
+	require.Truef(t,
+		func(stats []roachpb.CollectedStatementStatistics) bool {
+			for _, stat := range stats {
+				if stat.Key.Query == "INSERT INTO _ VALUES (_)" {
+					return true
+				}
+			}
+			return false
+		}(stats),
+		"expected to find in-memory stats",
+	)
+
+	// Issue a drain.
+	drainCtx.sendDrainNoShutdown()
+
+	// Open a new SQL connection.
+	sqlDB = sqlutils.MakeSQLRunner(drainCtx.tc.ServerConn(1))
+
+	// Check that the stats were flushed into the statement statistics system table.
+	// Verify that the number of statistics for node 1 are non-zero.
+	require.NotEqualf(t,
+		func() int {
+			q := sqlDB.Query(t, `SELECT count(*) FROM system.statement_statistics WHERE node_id = 1`)
+			q.Next()
+			var c int
+			require.Nil(t, q.Scan(&c))
+			require.Nil(t, q.Close())
+			return c
+		}(),
+		0,
+		"expected to find system table stats")
+}
+
 type testDrainContext struct {
 	*testing.T
 	tc         *testcluster.TestCluster