From 910c843eb9bcb0619aa6eb96e529e60ea99b0def Mon Sep 17 00:00:00 2001
From: Yahor Yuzefovich <yahor@cockroachlabs.com>
Date: Thu, 22 Sep 2022 12:12:38 -0700
Subject: [PATCH] persistedsqlstats: speed up a test

Previously, a single unit test could take on the order of 4 minutes (or
even exceed 5 minute timeout, rarely) because the job monitor checks
whether a cluster setting has been updated only every minute, and we
update the cluster setting twice in a unit test. This commit makes it so
that in a testing setup the check happens every second.

Release note: None
---
 pkg/sql/sqlstats/persistedsqlstats/provider.go         |  1 +
 .../persistedsqlstats/scheduled_job_monitor.go         | 10 +++++++++-
 .../scheduled_sql_stats_compaction_test.go             |  9 +++------
 pkg/sql/sqlstats/test_utils.go                         |  5 +++++
 4 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/pkg/sql/sqlstats/persistedsqlstats/provider.go b/pkg/sql/sqlstats/persistedsqlstats/provider.go
index eaaa3a665ceb..9750d5373f33 100644
--- a/pkg/sql/sqlstats/persistedsqlstats/provider.go
+++ b/pkg/sql/sqlstats/persistedsqlstats/provider.go
@@ -93,6 +93,7 @@ func New(cfg *Config, memSQLStats *sslocal.SQLStats) *PersistedSQLStats {
 		scanInterval: defaultScanInterval,
 		jitterFn:     p.jitterInterval,
 	}
+	p.jobMonitor.testingKnobs.updateCheckInterval = cfg.Knobs.JobMonitorUpdateCheckInterval
 
 	return p
 }
diff --git a/pkg/sql/sqlstats/persistedsqlstats/scheduled_job_monitor.go b/pkg/sql/sqlstats/persistedsqlstats/scheduled_job_monitor.go
index d61779d3bff6..0c0b2694ed11 100644
--- a/pkg/sql/sqlstats/persistedsqlstats/scheduled_job_monitor.go
+++ b/pkg/sql/sqlstats/persistedsqlstats/scheduled_job_monitor.go
@@ -65,6 +65,9 @@ type jobMonitor struct {
 	db           *kv.DB
 	scanInterval time.Duration
 	jitterFn     func(time.Duration) time.Duration
+	testingKnobs struct {
+		updateCheckInterval time.Duration
+	}
 }
 
 func (j *jobMonitor) start(ctx context.Context, stopper *stop.Stopper) {
@@ -80,6 +83,11 @@ func (j *jobMonitor) start(ctx context.Context, stopper *stop.Stopper) {
 		timer.Reset(0)
 		defer timer.Stop()
 
+		updateCheckInterval := time.Minute
+		if j.testingKnobs.updateCheckInterval != 0 {
+			updateCheckInterval = j.testingKnobs.updateCheckInterval
+		}
+
 		// This loop runs every minute to check if we need to update the job schedule.
 		// We only hit the jobs table if the schedule needs to be updated due to a
 		// change in the recurrence cluster setting or as a scheduled check to
@@ -97,7 +105,7 @@ func (j *jobMonitor) start(ctx context.Context, stopper *stop.Stopper) {
 				currentRecurrence = SQLStatsCleanupRecurrence.Get(&j.st.SV)
 			}
 
-			timer.Reset(time.Minute)
+			timer.Reset(updateCheckInterval)
 		}
 	})
 }
diff --git a/pkg/sql/sqlstats/persistedsqlstats/scheduled_sql_stats_compaction_test.go b/pkg/sql/sqlstats/persistedsqlstats/scheduled_sql_stats_compaction_test.go
index 47fdfb692b28..36725ec7aba4 100644
--- a/pkg/sql/sqlstats/persistedsqlstats/scheduled_sql_stats_compaction_test.go
+++ b/pkg/sql/sqlstats/persistedsqlstats/scheduled_sql_stats_compaction_test.go
@@ -29,7 +29,6 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/tests"
 	"github.com/cockroachdb/cockroach/pkg/testutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
-	"github.com/cockroachdb/cockroach/pkg/testutils/skip"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
@@ -178,11 +177,9 @@ func TestScheduledSQLStatsCompaction(t *testing.T) {
 func TestSQLStatsScheduleOperations(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	defer log.Scope(t).Close(t)
-	skip.UnderStressRace(t, "test is too slow to run under race")
 
 	ctx := context.Background()
-	helper, helperCleanup := newTestHelper(t, nil /* sqlStatsKnobs */)
-	helper.sqlDB.SucceedsSoonDuration = 2 * time.Minute
+	helper, helperCleanup := newTestHelper(t, &sqlstats.TestingKnobs{JobMonitorUpdateCheckInterval: time.Second})
 	defer helperCleanup()
 
 	schedID := getSQLStatsCompactionSchedule(t, helper).ScheduleID()
@@ -218,7 +215,7 @@ func TestSQLStatsScheduleOperations(t *testing.T) {
 			helper.sqlDB.Exec(t, "SET CLUSTER SETTING sql.stats.cleanup.recurrence = $1", expr)
 
 			var err error
-			testutils.SucceedsWithin(t, func() error {
+			testutils.SucceedsSoon(t, func() error {
 				// Reload schedule from DB.
 				sj := getSQLStatsCompactionSchedule(t, helper)
 				err = persistedsqlstats.CheckScheduleAnomaly(sj)
@@ -227,7 +224,7 @@ func TestSQLStatsScheduleOperations(t *testing.T) {
 				}
 				require.Equal(t, expr, sj.ScheduleExpr())
 				return nil
-			}, time.Minute*2)
+			})
 
 			require.True(t, errors.Is(
 				errors.Unwrap(err), persistedsqlstats.ErrScheduleIntervalTooLong),
diff --git a/pkg/sql/sqlstats/test_utils.go b/pkg/sql/sqlstats/test_utils.go
index 47f12a714c20..b024e6f30886 100644
--- a/pkg/sql/sqlstats/test_utils.go
+++ b/pkg/sql/sqlstats/test_utils.go
@@ -33,6 +33,11 @@ type TestingKnobs struct {
 	// AOSTClause overrides the AS OF SYSTEM TIME clause in queries used in
 	// persistedsqlstats.
 	AOSTClause string
+
+	// JobMonitorUpdateCheckInterval if non-zero indicates the frequency at
+	// which the job monitor needs to check whether the schedule needs to be
+	// updated.
+	JobMonitorUpdateCheckInterval time.Duration
 }
 
 // ModuleTestingKnobs implements base.ModuleTestingKnobs interface.