scheduledlogging: fix TestCaptureIndexUsageStats for secondary tenants

pkg/sql/scheduledlogging/BUILD.bazel

@@ -37,6 +37,7 @@ go_test(
         "//pkg/security/securitytest",
         "//pkg/server",
         "//pkg/settings/cluster",
+        "//pkg/testutils",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/sqlutils",
         "//pkg/testutils/testcluster",

pkg/sql/scheduledlogging/captured_index_usage_stats.go

@@ -71,9 +71,6 @@ type CaptureIndexUsageStatsTestingKnobs struct {
 	// scheduled interval in the case that the logging duration exceeds the
 	// default scheduled interval duration.
 	getOverlapDuration func() time.Duration
-	// onScheduleComplete allows tests to hook into when the current schedule
-	// is completed to check for the expected logs.
-	onScheduleComplete func()
 }
 
 // ModuleTestingKnobs implements base.ModuleTestingKnobs interface.
@@ -158,9 +155,6 @@ func (s *CaptureIndexUsageStatsLoggingScheduler) start(ctx context.Context, stop
 					dur = time.Second
 				}
 				timer.Reset(dur)
-				if s.knobs != nil && s.knobs.onScheduleComplete != nil {
-					s.knobs.onScheduleComplete()
-				}
 			}
 		}
 	})

pkg/sql/scheduledlogging/captured_index_usage_stats_test.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"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/util/leaktest"
@@ -98,20 +99,12 @@ func TestCaptureIndexUsageStats(t *testing.T) {
 	// Configure the delay between each emission of index usage stats logs.
 	telemetryCaptureIndexUsageStatsLoggingDelay.Override(context.Background(), &settings.SV, stubLoggingDelay)
 
-	scheduleCompleteChan := make(chan struct{})
-
 	srv, sqlDB, _ := serverutils.StartServer(t, base.TestServerArgs{
-		DefaultTestTenant: base.TestIsForStuffThatShouldWorkWithSecondaryTenantsButDoesntYet(109465),
-
 		Settings: settings,
 		Knobs: base.TestingKnobs{
 			CapturedIndexUsageStatsKnobs: &CaptureIndexUsageStatsTestingKnobs{
 				getLoggingDuration: sd.getLoggingDuration,
 				getOverlapDuration: sd.getOverlapDuration,
-				onScheduleComplete: func() {
-					scheduleCompleteChan <- struct{}{}
-					<-scheduleCompleteChan
-				},
 			},
 		},
 	})
@@ -161,18 +154,14 @@ func TestCaptureIndexUsageStats(t *testing.T) {
 	expectedTotalNumEntriesInSingleInterval := 8
 	expectedNumberOfIndividualIndexEntriesInSingleInterval := 1
 
-	// Wait for channel value from end of 1st schedule.
-	<-scheduleCompleteChan
-
 	// Check the expected number of entries from the 1st schedule.
-	require.NoError(t, checkNumTotalEntriesAndNumIndexEntries(t,
-		expectedIndexNames,
-		expectedTotalNumEntriesInSingleInterval,
-		expectedNumberOfIndividualIndexEntriesInSingleInterval,
-		scheduleCompleteChan,
-	), "error encountered checking the number of total entries and number of index entries")
-
-	scheduleCompleteChan <- struct{}{}
+	testutils.SucceedsSoon(t, func() error {
+		return checkNumTotalEntriesAndNumIndexEntries(t,
+			expectedIndexNames,
+			expectedTotalNumEntriesInSingleInterval,
+			expectedNumberOfIndividualIndexEntriesInSingleInterval,
+		)
+	})
 
 	// Expect number of total entries to hold 2 times the number of entries in a
 	// single interval.
@@ -185,18 +174,14 @@ func TestCaptureIndexUsageStats(t *testing.T) {
 	stubLoggingDuration := stubScheduleInterval * 2
 	sd.setLoggingDuration(stubLoggingDuration)
 
-	// Wait for channel value from end of 2nd schedule.
-	<-scheduleCompleteChan
-
 	// Check the expected number of entries from the 2nd schedule.
-	require.NoError(t, checkNumTotalEntriesAndNumIndexEntries(t,
-		expectedIndexNames,
-		expectedTotalNumEntriesAfterTwoIntervals,
-		expectedNumberOfIndividualIndexEntriesAfterTwoIntervals,
-		scheduleCompleteChan,
-	), "error encountered checking the number of total entries and number of index entries")
-
-	scheduleCompleteChan <- struct{}{}
+	testutils.SucceedsSoon(t, func() error {
+		return checkNumTotalEntriesAndNumIndexEntries(t,
+			expectedIndexNames,
+			expectedTotalNumEntriesAfterTwoIntervals,
+			expectedNumberOfIndividualIndexEntriesAfterTwoIntervals,
+		)
+	})
 
 	// Expect number of total entries to hold 3 times the number of entries in a
 	// single interval.
@@ -205,22 +190,18 @@ func TestCaptureIndexUsageStats(t *testing.T) {
 	// entries in a single interval.
 	expectedNumberOfIndividualIndexEntriesAfterThreeIntervals := expectedNumberOfIndividualIndexEntriesInSingleInterval * 3
 
-	// Wait for channel value from end of 3rd schedule.
-	<-scheduleCompleteChan
-
 	// Check the expected number of entries from the 3rd schedule.
-	require.NoError(t, checkNumTotalEntriesAndNumIndexEntries(t,
-		expectedIndexNames,
-		expectedTotalNumEntriesAfterThreeIntervals,
-		expectedNumberOfIndividualIndexEntriesAfterThreeIntervals,
-		scheduleCompleteChan,
-	), "error encountered checking the number of total entries and number of index entries")
+	testutils.SucceedsSoon(t, func() error {
+		return checkNumTotalEntriesAndNumIndexEntries(t,
+			expectedIndexNames,
+			expectedTotalNumEntriesAfterThreeIntervals,
+			expectedNumberOfIndividualIndexEntriesAfterThreeIntervals,
+		)
+	})
 
 	// Stop capturing index usage statistics.
 	telemetryCaptureIndexUsageStatsEnabled.Override(context.Background(), &settings.SV, false)
 
-	scheduleCompleteChan <- struct{}{}
-
 	// Iterate through entries, ensure that the timestamp difference between each
 	// schedule is as expected.
 	startTimestamp := int64(0)
@@ -283,7 +264,6 @@ func checkNumTotalEntriesAndNumIndexEntries(
 	expectedIndexNames []string,
 	expectedTotalEntries int,
 	expectedIndividualIndexEntries int,
-	scheduleCompleteChan chan struct{},
 ) error {
 	log.FlushFiles()
 	// Fetch log entries.
@@ -295,13 +275,11 @@ func checkNumTotalEntriesAndNumIndexEntries(
 		log.WithMarkedSensitiveData,
 	)
 	if err != nil {
-		close(scheduleCompleteChan)
 		return err
 	}
 
 	// Assert that we have the correct number of entries.
 	if expectedTotalEntries != len(entries) {
-		close(scheduleCompleteChan)
 		return errors.Newf("expected %d total entries, got %d", expectedTotalEntries, len(entries))
 	}
 
@@ -310,7 +288,7 @@ func checkNumTotalEntriesAndNumIndexEntries(
 	for _, e := range entries {
 		t.Logf("checking entry: %v", e)
 		// Check that the entry has a tag for a node ID of 1.
-		if !strings.Contains(e.Tags, `n1`) {
+		if !strings.Contains(e.Tags, `n1`) && !strings.Contains(e.Tags, `nsql1`) {
 			t.Fatalf("expected the entry's tags to include n1, but include got %s", e.Tags)
 		}
 
@@ -319,7 +297,6 @@ func checkNumTotalEntriesAndNumIndexEntries(
 		}
 		err := json.Unmarshal([]byte(e.Message), &s)
 		if err != nil {
-			close(scheduleCompleteChan)
 			t.Fatal(err)
 		}
 		countByIndex[s.IndexName] = countByIndex[s.IndexName] + 1
@@ -328,21 +305,18 @@ func checkNumTotalEntriesAndNumIndexEntries(
 	t.Logf("found index counts: %+v", countByIndex)
 
 	if expected, actual := expectedTotalEntries/expectedIndividualIndexEntries, len(countByIndex); actual != expected {
-		close(scheduleCompleteChan)
 		return errors.Newf("expected %d indexes, got %d", expected, actual)
 	}
 
 	for idxName, c := range countByIndex {
 		if c != expectedIndividualIndexEntries {
-			close(scheduleCompleteChan)
 			return errors.Newf("for index %s: expected entry count %d, got %d",
 				idxName, expectedIndividualIndexEntries, c)
 		}
 	}
 
 	for _, idxName := range expectedIndexNames {
 		if _, ok := countByIndex[idxName]; !ok {
-			close(scheduleCompleteChan)
 			return errors.Newf("no entry found for index %s", idxName)
 		}
 	}