From 5c3d4ce153b207954ecd3de32efdc349065a4452 Mon Sep 17 00:00:00 2001
From: Nikhil Benesch <nikhil.benesch@gmail.com>
Date: Wed, 10 Jan 2018 17:27:49 -0500
Subject: [PATCH] sqlccl: speed up partitioning tests

Partitioning tests are bottlenecked on test cluster rebalancing. Make
rebalancing much faster by turning down the scan interval and disabling
store throttling. TestInitialPartitioning and TestRepartitioning now
take under 30s combined, when previously they could take several
minutes.

Release note: None
---
 pkg/ccl/sqlccl/partition_test.go | 52 +++++++++++++++-----------------
 1 file changed, 24 insertions(+), 28 deletions(-)

diff --git a/pkg/ccl/sqlccl/partition_test.go b/pkg/ccl/sqlccl/partition_test.go
index 3eaca8e78ffa..963ab1285f61 100644
--- a/pkg/ccl/sqlccl/partition_test.go
+++ b/pkg/ccl/sqlccl/partition_test.go
@@ -1061,18 +1061,14 @@ func verifyScansOnNode(db *gosql.DB, query string, node string) error {
 	return nil
 }
 
-func TestInitialPartitioning(t *testing.T) {
-	defer leaktest.AfterTest(t)()
-	rng, _ := randutil.NewPseudoRand()
-
+func setupPartitioningTestCluster(ctx context.Context, t testing.TB) (*sqlutils.SQLRunner, func()) {
 	cfg := config.DefaultZoneConfig()
 	cfg.NumReplicas = 1
-	defer config.TestingSetDefaultZoneConfig(cfg)()
+	resetZoneConfig := config.TestingSetDefaultZoneConfig(cfg)
 
-	ctx := context.Background()
 	tsArgs := func(attr string) base.TestServerArgs {
 		return base.TestServerArgs{
-			ScanInterval: time.Second,
+			ScanInterval: 100 * time.Millisecond,
 			StoreSpecs: []base.StoreSpec{
 				{InMemory: true, Attributes: roachpb.Attributes{Attrs: []string{attr}}},
 			},
@@ -1084,12 +1080,30 @@ func TestInitialPartitioning(t *testing.T) {
 		2: tsArgs("n3"),
 	}}
 	tc := testcluster.StartTestCluster(t, 3, tcArgs)
-	defer tc.Stopper().Stop(context.Background())
+
 	sqlDB := sqlutils.MakeSQLRunner(tc.Conns[0])
 	sqlDB.Exec(t, `CREATE DATABASE data`)
 	sqlDB.Exec(t, `USE data`)
 
+	// Disabling store throttling vastly speeds up rebalancing.
+	sqlDB.Exec(t, `SET CLUSTER SETTING server.declined_reservation_timeout = '0s'`)
+	sqlDB.Exec(t, `SET CLUSTER SETTING server.failed_reservation_timeout = '0s'`)
+
+	return sqlDB, func() {
+		tc.Stopper().Stop(context.Background())
+		resetZoneConfig()
+	}
+}
+
+func TestInitialPartitioning(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	rng, _ := randutil.NewPseudoRand()
 	testCases := allPartitioningTests(rng)
+
+	ctx := context.Background()
+	sqlDB, cleanup := setupPartitioningTestCluster(ctx, t)
+	defer cleanup()
+
 	for _, test := range testCases {
 		if len(test.scans) == 0 {
 			continue
@@ -1189,27 +1203,9 @@ func TestRepartitioning(t *testing.T) {
 		t.Fatalf("%+v", err)
 	}
 
-	cfg := config.DefaultZoneConfig()
-	cfg.NumReplicas = 1
-	defer config.TestingSetDefaultZoneConfig(cfg)()
-
 	ctx := context.Background()
-	tsArgs := func(attr string) base.TestServerArgs {
-		return base.TestServerArgs{
-			ScanInterval: time.Second,
-			StoreSpecs: []base.StoreSpec{
-				{InMemory: true, Attributes: roachpb.Attributes{Attrs: []string{attr}}},
-			},
-		}
-	}
-	tcArgs := base.TestClusterArgs{ServerArgsPerNode: map[int]base.TestServerArgs{
-		0: tsArgs("n1"),
-		1: tsArgs("n2"),
-		2: tsArgs("n3"),
-	}}
-	tc := testcluster.StartTestCluster(t, 3, tcArgs)
-	defer tc.Stopper().Stop(context.Background())
-	sqlDB := sqlutils.MakeSQLRunner(tc.Conns[0])
+	sqlDB, cleanup := setupPartitioningTestCluster(ctx, t)
+	defer cleanup()
 
 	for _, test := range testCases {
 		t.Run(fmt.Sprintf("%s/%s", test.old.name, test.new.name), func(t *testing.T) {