sqlstats: add retries to stats test on locked table

pkg/sql/sqlstats/persistedsqlstats/cluster_settings.go

@@ -137,10 +137,10 @@ var sqlStatsLimitTableSizeEnabled = settings.RegisterBoolSetting(
 	true,
 )
 
-// sqlStatsLimitTableCheckInterval is the cluster setting the controls what
+// SQLStatsLimitTableCheckInterval is the cluster setting the controls what
 // interval the check is done if the statement and transaction statistics
 // tables have grown past the sql.stats.persisted_rows.max.
-var sqlStatsLimitTableCheckInterval = settings.RegisterDurationSetting(
+var SQLStatsLimitTableCheckInterval = settings.RegisterDurationSetting(
 	settings.ApplicationLevel,
 	"sql.stats.limit_table_size_check.interval",
 	"controls what interval the check is done if the statement and "+

pkg/sql/sqlstats/persistedsqlstats/flush.go

@@ -111,7 +111,7 @@ func (s *PersistedSQLStats) Flush(ctx context.Context, stopper *stop.Stopper) {
 func (s *PersistedSQLStats) StmtsLimitSizeReached(ctx context.Context) (bool, error) {
 	// Doing a count check on every flush for every node adds a lot of overhead.
 	// To reduce the overhead only do the check once an hour by default.
-	intervalToCheck := sqlStatsLimitTableCheckInterval.Get(&s.cfg.Settings.SV)
+	intervalToCheck := SQLStatsLimitTableCheckInterval.Get(&s.cfg.Settings.SV)
 	if !s.lastSizeCheck.IsZero() && s.lastSizeCheck.Add(intervalToCheck).After(timeutil.Now()) {
 		log.Infof(ctx, "PersistedSQLStats.StmtsLimitSizeReached skipped with last check at: %s and check interval: %s", s.lastSizeCheck, intervalToCheck)
 		return false, nil

pkg/sql/sqlstats/persistedsqlstats/flush_test.go

@@ -682,8 +682,7 @@ func TestSQLStatsReadLimitSizeOnLockedTable(t *testing.T) {
 	// Ensure we have some rows in system.statement_statistics
 	require.GreaterOrEqual(t, stmtStatsCountFlush, minNumExpectedStmts)
 
-	// Set sql.stats.persisted_rows.max
-	sqlConn.Exec(t, fmt.Sprintf("SET CLUSTER SETTING sql.stats.persisted_rows.max=%d", maxNumPersistedRows))
+	persistedsqlstats.SQLStatsMaxPersistedRows.Override(ctx, &s.ClusterSettings().SV, maxNumPersistedRows)
 
 	// We need SucceedsSoon here for the follower read timestamp to catch up
 	// enough for this state to be reached.
@@ -705,7 +704,7 @@ func TestSQLStatsReadLimitSizeOnLockedTable(t *testing.T) {
 
 	// Set table size check interval to .0000001 second. So the next check doesn't
 	// use the cached value.
-	sqlConn.Exec(t, "SET CLUSTER SETTING sql.stats.limit_table_size_check.interval='.0000001s'")
+	persistedsqlstats.SQLStatsLimitTableCheckInterval.Override(ctx, &s.ClusterSettings().SV, time.Nanosecond)
 
 	// Begin a transaction.
 	sqlConn.Exec(t, "BEGIN")
@@ -714,28 +713,35 @@ func TestSQLStatsReadLimitSizeOnLockedTable(t *testing.T) {
 
 	// Ensure that we can read from the table despite it being locked, due to the follower read (AOST).
 	// Expect that the number of statements in the table exceeds sql.stats.persisted_rows.max * 1.5
-	// (meaning that the limit will be reached) and no error. Loop to make sure that
-	// checking it multiple times still returns the correct value.
+	// (meaning that the limit will be reached) and no error. Every iteration picks a random shard, and we
+	// loop 3 times to make sure that we find at least one shard with a count over the limit. In the wild,
+	// we've observed individual shards only having a single statement recorded which makes this check fail
+	// otherwise.
+	foundLimit := false
 	for i := 0; i < 3; i++ {
 		limitReached, err = pss.StmtsLimitSizeReached(ctx)
 		require.NoError(t, err)
-		if !limitReached {
-			readStmt := `SELECT crdb_internal_aggregated_ts_app_name_fingerprint_id_node_id_plan_hash_transaction_fingerprint_id_shard_8, count(*)
+		if limitReached {
+			foundLimit = true
+		}
+	}
+
+	if !foundLimit {
+		readStmt := `SELECT crdb_internal_aggregated_ts_app_name_fingerprint_id_node_id_plan_hash_transaction_fingerprint_id_shard_8, count(*)
       FROM system.statement_statistics
       AS OF SYSTEM TIME follower_read_timestamp()
       GROUP BY crdb_internal_aggregated_ts_app_name_fingerprint_id_node_id_plan_hash_transaction_fingerprint_id_shard_8`
 
-			sqlConn2 := sqlutils.MakeSQLRunner(s.SQLConn(t))
-			rows := sqlConn2.Query(t, readStmt)
-			shard := make([]int64, 8)
-			count := make([]int64, 8)
-			for j := 0; rows.Next(); {
-				err := rows.Scan(&shard[j], &count[j])
-				require.NoError(t, err)
-				j += 1
-			}
-			t.Fatalf("limitReached should be true. loop: %d; shards: %d counts: %d", i, shard, count)
+		sqlConn2 := sqlutils.MakeSQLRunner(s.SQLConn(t))
+		rows := sqlConn2.Query(t, readStmt)
+		shard := make([]int64, 8)
+		count := make([]int64, 8)
+		for j := 0; rows.Next(); {
+			err := rows.Scan(&shard[j], &count[j])
+			require.NoError(t, err)
+			j += 1
 		}
+		t.Fatalf("limitReached should be true. shards: %d counts: %d", shard, count)
 	}
 
 	// Close the transaction.