From 0541e67709f2633fd0443fa45aac2e0510884b0b Mon Sep 17 00:00:00 2001
From: Yevgeniy Miretskiy <yevgeniy@cockroachlabs.com>
Date: Tue, 29 Sep 2020 18:44:28 -0400
Subject: [PATCH] roachprod: Support multiple stores on AWS

Add an aws specific command line option to make it possible
to create AWS instances with multiple /mnt/dataN directories in
order to support multiple stores.

Release Notes: None
Release Justification: N/A; roachprod script changes.
---
 pkg/cmd/roachprod/vm/aws/aws.go     |  9 ++++---
 pkg/cmd/roachprod/vm/aws/support.go | 37 +++++++++++++++++++----------
 2 files changed, 31 insertions(+), 15 deletions(-)

diff --git a/pkg/cmd/roachprod/vm/aws/aws.go b/pkg/cmd/roachprod/vm/aws/aws.go
index 1fccb37e50c1..d5417c0295c1 100644
--- a/pkg/cmd/roachprod/vm/aws/aws.go
+++ b/pkg/cmd/roachprod/vm/aws/aws.go
@@ -75,6 +75,7 @@ type providerOpts struct {
 	EBSVolumeType      string
 	EBSVolumeSize      int
 	EBSProvisionedIOPs int
+	UseMultipleDisks   bool
 
 	// Use specified ImageAMI when provisioning.
 	// Overrides config.json AMI.
@@ -147,6 +148,8 @@ func (o *providerOpts) ConfigureCreateFlags(flags *pflag.FlagSet) {
 			"of geo (default [%s])", strings.Join(defaultCreateZones, ",")))
 	flags.StringVar(&o.ImageAMI, ProviderName+"-image-ami",
 		"", "Override image AMI to use.  See https://awscli.amazonaws.com/v2/documentation/api/latest/reference/ec2/describe-images.html")
+	flags.BoolVar(&o.UseMultipleDisks, ProviderName+"-enable-multiple-stores",
+		false, "Enable the use of multiple stores by creating one store directory per disk.  Default is to raid0 stripe all disks.")
 }
 
 func (o *providerOpts) ConfigureClusterFlags(flags *pflag.FlagSet, _ vm.MultipleProjectsOption) {
@@ -676,7 +679,7 @@ func (p *Provider) runInstance(name string, zone string, opts vm.CreateOpts) err
 			extraMountOpts = "nobarrier"
 		}
 	}
-	filename, err := writeStartupScript(extraMountOpts)
+	filename, err := writeStartupScript(extraMountOpts, p.opts.UseMultipleDisks)
 	if err != nil {
 		return errors.Wrapf(err, "could not write AWS startup script to temp file")
 	}
@@ -713,10 +716,10 @@ func (p *Provider) runInstance(name string, zone string, opts vm.CreateOpts) err
 	if !opts.SSDOpts.UseLocalSSD {
 		var ebsParams string
 		switch t := p.opts.EBSVolumeType; t {
-		case "gp2", "io2":
+		case "gp2":
 			ebsParams = fmt.Sprintf("{VolumeSize=%d,VolumeType=%s,DeleteOnTermination=true}",
 				p.opts.EBSVolumeSize, t)
-		case "io1":
+		case "io1", "io2":
 			ebsParams = fmt.Sprintf("{VolumeSize=%d,VolumeType=%s,Iops=%d,DeleteOnTermination=true}",
 				p.opts.EBSVolumeSize, t, p.opts.EBSProvisionedIOPs)
 		default:
diff --git a/pkg/cmd/roachprod/vm/aws/support.go b/pkg/cmd/roachprod/vm/aws/support.go
index fc042d452905..b09407bb2ea5 100644
--- a/pkg/cmd/roachprod/vm/aws/support.go
+++ b/pkg/cmd/roachprod/vm/aws/support.go
@@ -42,8 +42,11 @@ sudo apt-get install -qy --no-install-recommends mdadm
 mount_opts="discard,defaults"
 {{if .ExtraMountOpts}}mount_opts="${mount_opts},{{.ExtraMountOpts}}"{{end}}
 
+use_multiple_disks='{{if .UseMultipleDisks}}true{{end}}'
+
 disks=()
-mountpoint="/mnt/data1"
+mount_prefix="/mnt/data"
+
 # On different machine types, the drives are either called nvme... or xvdd.
 for d in $(ls /dev/nvme?n1 /dev/xvdd); do
   if ! mount | grep ${d}; then
@@ -53,19 +56,28 @@ for d in $(ls /dev/nvme?n1 /dev/xvdd); do
     echo "Disk ${d} already mounted, skipping..."
   fi
 done
+
+
 if [ "${#disks[@]}" -eq "0" ]; then
+  mountpoint="${mount_prefix}1"
   echo "No disks mounted, creating ${mountpoint}"
   mkdir -p ${mountpoint}
   chmod 777 ${mountpoint}
-elif [ "${#disks[@]}" -eq "1" ]; then
-  echo "One disk mounted, creating ${mountpoint}"
-  mkdir -p ${mountpoint}
-  disk=${disks[0]}
-  mkfs.ext4 -E nodiscard ${disk}
-  mount -o ${mount_opts} ${disk} ${mountpoint}
-  chmod 777 ${mountpoint}
-  echo "${disk} ${mountpoint} ext4 ${mount_opts} 1 1" | tee -a /etc/fstab
+elif [ "${#disks[@]}" -eq "1" ] || [ -n "use_multiple_disks" ]; then
+  disknum=1
+  for disk in "${disks[@]}"
+  do
+    mountpoint="${mount_prefix}${disknum}"
+    disknum=$((disknum + 1 ))
+    echo "Creating ${mountpoint}"
+    mkdir -p ${mountpoint}
+    mkfs.ext4 -E nodiscard ${disk}
+    mount -o ${mount_opts} ${disk} ${mountpoint}
+    chmod 777 ${mountpoint}
+    echo "${disk} ${mountpoint} ext4 ${mount_opts} 1 1" | tee -a /etc/fstab
+  done
 else
+  mountpoint="${mount_prefix}1"
   echo "${#disks[@]} disks mounted, creating ${mountpoint} using RAID 0"
   mkdir -p ${mountpoint}
   raiddisk="/dev/md0"
@@ -125,12 +137,13 @@ sudo touch /mnt/data1/.roachprod-initialized
 //
 // extraMountOpts, if not empty, is appended to the default mount options. It is
 // a comma-separated list of options for the "mount -o" flag.
-func writeStartupScript(extraMountOpts string) (string, error) {
+func writeStartupScript(extraMountOpts string, useMultiple bool) (string, error) {
 	type tmplParams struct {
-		ExtraMountOpts string
+		ExtraMountOpts   string
+		UseMultipleDisks bool
 	}
 
-	args := tmplParams{ExtraMountOpts: extraMountOpts}
+	args := tmplParams{ExtraMountOpts: extraMountOpts, UseMultipleDisks: useMultiple}
 
 	tmpfile, err := ioutil.TempFile("", "aws-startup-script")
 	if err != nil {