roachtest/indexes: support aws and tune for latency

This commit updates the `indexes` suite of roachtests to support a multi-region
AWS deployment alongside its multi-region GCE deployment.

The commit then tunes the workload to optimize for a stable p50 and p99 latency
all the way up to 50 secondary indexes. The concurrency is fairly low with low
numbers of indexes, but that's not what the test is trying to stress.

pkg/cmd/roachtest/cluster.go

@@ -885,16 +885,9 @@ func (s *clusterSpec) args() []string {
 
 	switch cloud {
 	case aws:
-		if s.Zones != "" {
-			fmt.Fprintf(os.Stderr, "zones spec not yet supported on AWS: %s\n", s.Zones)
-			os.Exit(1)
-		}
-		if s.Geo {
-			fmt.Fprintf(os.Stderr, "geo-distributed clusters not yet supported on AWS\n")
-			os.Exit(1)
-		}
-
 		args = append(args, "--clouds=aws")
+	case gce:
+		args = append(args, "--clouds=gce")
 	case azure:
 		args = append(args, "--clouds=azure")
 	}

pkg/cmd/roachtest/indexes.go

@@ -15,11 +15,17 @@ import (
 	"fmt"
 	"strconv"
 	"strings"
+	"time"
+
+	"github.com/cockroachdb/cockroach/pkg/util/retry"
 )
 
 func registerNIndexes(r *testRegistry, secondaryIndexes int) {
 	const nodes = 6
-	geoZones := []string{"us-west1-b", "us-east1-b", "us-central1-a"}
+	geoZones := []string{"us-east1-b", "us-west1-b", "europe-west2-b"}
+	if cloud == aws {
+		geoZones = []string{"us-east-2b", "us-west-1a", "eu-west-1a"}
+	}
 	geoZonesStr := strings.Join(geoZones, ",")
 	r.Add(testSpec{
 		Name:    fmt.Sprintf("indexes/%d/nodes=%d/multi-region", secondaryIndexes, nodes),
@@ -36,6 +42,7 @@ func registerNIndexes(r *testRegistry, secondaryIndexes int) {
 			c.Put(ctx, cockroach, "./cockroach", roachNodes)
 			c.Put(ctx, workload, "./workload", loadNode)
 			c.Start(ctx, t, roachNodes)
+			conn := c.Conn(ctx, 1)
 
 			t.Status("running workload")
 			m := newMonitor(ctx, c, roachNodes)
@@ -47,15 +54,55 @@ func registerNIndexes(r *testRegistry, secondaryIndexes int) {
 				// Set lease preferences so that all leases for the table are
 				// located in the availability zone with the load generator.
 				if !local {
-					leasePrefs := fmt.Sprintf(`ALTER TABLE indexes.indexes
-						                       CONFIGURE ZONE USING
-						                       constraints = COPY FROM PARENT,
-						                       lease_preferences = '[[+zone=%s]]'`, firstAZ)
-					c.Run(ctx, c.Node(1), `./cockroach sql --insecure -e "`+leasePrefs+`"`)
+					t.l.Printf("setting lease preferences")
+					if _, err := conn.ExecContext(ctx, fmt.Sprintf(`
+						ALTER TABLE indexes.indexes
+						CONFIGURE ZONE USING
+						constraints = COPY FROM PARENT,
+						lease_preferences = '[[+zone=%s]]'`,
+						firstAZ,
+					)); err != nil {
+						return err
+					}
+
+					// Add a small delay to allow ranges to rebalance.
+					t.l.Printf("waiting for rebalancing")
+					time.Sleep(time.Duration(5*secondaryIndexes) * time.Second)
+
+					// Wait for leases to adhere to preferences, if they aren't
+					// already.
+					for r := retry.StartWithCtx(ctx, retry.Options{}); r.Next(); {
+						t.l.Printf("checking lease preferences")
+						var ok bool
+						if err := conn.QueryRowContext(ctx, `
+							SELECT lease_holder <= $1
+							FROM crdb_internal.ranges
+							WHERE table_name = 'indexes'`,
+							nodes/3,
+						).Scan(&ok); err != nil {
+							return err
+						} else if ok {
+							break
+						}
+
+						t.l.Printf("leases still rebalancing...")
+					}
+				}
+
+				// Set the DistSender concurrency setting high enough so that no
+				// requests get throttled. Add 2x headroom on top of this.
+				conc := 16 * len(gatewayNodes)
+				parallelWrites := (secondaryIndexes + 1) * conc
+				distSenderConc := 2 * parallelWrites
+				if _, err := conn.ExecContext(ctx, `
+					SET CLUSTER SETTING kv.dist_sender.concurrency_limit = $1`,
+					distSenderConc,
+				); err != nil {
+					return err
 				}
 
-				payload := " --payload=256"
-				concurrency := ifLocal("", " --concurrency="+strconv.Itoa(nodes*32))
+				payload := " --payload=64"
+				concurrency := ifLocal("", " --concurrency="+strconv.Itoa(conc))
 				duration := " --duration=" + ifLocal("10s", "10m")
 				runCmd := fmt.Sprintf("./workload run indexes --histograms="+perfArtifactsDir+"/stats.json"+
 					payload+concurrency+duration+" {pgurl%s}", gatewayNodes)