From 72205ee9730761198561deea798b2970c7bbe13f Mon Sep 17 00:00:00 2001
From: Drew Kimball <drewk@cockroachlabs.com>
Date: Mon, 26 Jun 2023 15:42:02 +0000
Subject: [PATCH] multitenant: deflake TestConsumption

This patch adds retries to the `TestConsumption` test for retrieving the
total tenant RU consumption in order to ensure that the delta includes the
bucket flush that was triggered by the test query, instead of background
activity.

Fixes #94286
Fixes #106572

Release note: None
---
 .../tenantcostclient/tenant_side_test.go      | 38 ++++++++++++-------
 1 file changed, 25 insertions(+), 13 deletions(-)

diff --git a/pkg/ccl/multitenantccl/tenantcostclient/tenant_side_test.go b/pkg/ccl/multitenantccl/tenantcostclient/tenant_side_test.go
index 3dfe0ff0eb4d..169038b23896 100644
--- a/pkg/ccl/multitenantccl/tenantcostclient/tenant_side_test.go
+++ b/pkg/ccl/multitenantccl/tenantcostclient/tenant_side_test.go
@@ -876,8 +876,9 @@ func TestConsumption(t *testing.T) {
 	hostServer, _, _ := serverutils.StartServer(t, base.TestServerArgs{})
 	defer hostServer.Stopper().Stop(context.Background())
 
+	const targetPeriod = time.Millisecond
 	st := cluster.MakeTestingClusterSettings()
-	tenantcostclient.TargetPeriodSetting.Override(context.Background(), &st.SV, time.Millisecond)
+	tenantcostclient.TargetPeriodSetting.Override(context.Background(), &st.SV, targetPeriod)
 	tenantcostclient.CPUUsageAllowance.Override(context.Background(), &st.SV, 0)
 
 	testProvider := newTestProvider()
@@ -905,21 +906,32 @@ func TestConsumption(t *testing.T) {
 		r.Exec(t, "INSERT INTO t (v) SELECT repeat('1234567890', 1024) FROM generate_series(1, 10) AS g(i)")
 		const expectedBytes = 10 * 10 * 1024
 
-		afterWrite := testProvider.waitForConsumption(t)
-		delta := afterWrite
-		delta.Sub(&beforeWrite)
-		if delta.WriteBatches < 1 || delta.WriteRequests < 2 || delta.WriteBytes < expectedBytes*2 {
-			t.Errorf("usage after write: %s", delta.String())
-		}
+		// Try a few times because background activity can trigger bucket
+		// requests before the test query does.
+		testutils.SucceedsSoon(t, func() error {
+			afterWrite := testProvider.waitForConsumption(t)
+			delta := afterWrite
+			delta.Sub(&beforeWrite)
+			if delta.WriteBatches < 1 || delta.WriteRequests < 2 || delta.WriteBytes < expectedBytes*2 {
+				return errors.Newf("usage after write: %s", delta.String())
+			}
+			return nil
+		})
 
+		beforeRead := testProvider.waitForConsumption(t)
 		r.QueryStr(t, "SELECT min(v) FROM t")
 
-		afterRead := testProvider.waitForConsumption(t)
-		delta = afterRead
-		delta.Sub(&afterWrite)
-		if delta.ReadBatches < 1 || delta.ReadRequests < 1 || delta.ReadBytes < expectedBytes {
-			t.Errorf("usage after read: %s", delta.String())
-		}
+		// Try a few times because background activity can trigger bucket
+		// requests before the test query does.
+		testutils.SucceedsSoon(t, func() error {
+			afterRead := testProvider.waitForConsumption(t)
+			delta := afterRead
+			delta.Sub(&beforeRead)
+			if delta.ReadBatches < 1 || delta.ReadRequests < 1 || delta.ReadBytes < expectedBytes {
+				return errors.Newf("usage after read: %s", delta.String())
+			}
+			return nil
+		})
 		r.Exec(t, "DELETE FROM t WHERE true")
 	}
 	// Make sure some CPU usage is reported.