roachprod: Support AWS GP3 drives.

pkg/cmd/roachprod/vm/aws/aws.go

@@ -13,6 +13,7 @@ package aws
 import (
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"log"
 	"math/rand"
 	"os"
@@ -65,19 +66,115 @@ func init() {
 	vm.Providers[ProviderName] = p
 }
 
+// ebsDisk represent EBS disk device.
+// When marshaled to JSON format, produces JSON specification used
+// by AWS sdk to configure attached volumes.
+type ebsDisk struct {
+	VolumeType          string `json:"VolumeType"`
+	VolumeSize          int    `json:"VolumeSize"`
+	IOPs                int    `json:"Iops,omitempty"`
+	Throughput          int    `json:"Throughput,omitempty"`
+	DeleteOnTermination bool   `json:"DeleteOnTermination"`
+}
+
+// ebsVolume represents a mounted volume: name + ebsDisk
+type ebsVolume struct {
+	DeviceName string  `json:"DeviceName"`
+	Disk       ebsDisk `json:"Ebs"`
+}
+
+const ebsDefaultVolumeSizeGB = 500
+
+// Set implements flag Value interface.
+func (d *ebsDisk) Set(s string) error {
+	if err := json.Unmarshal([]byte(s), &d); err != nil {
+		return err
+	}
+
+	d.DeleteOnTermination = true
+
+	// Sanity check disk configuration.
+	// This is not strictly needed since AWS sdk would return error anyway,
+	// but we can return a nicer error message sooner.
+	if d.VolumeSize == 0 {
+		d.VolumeSize = ebsDefaultVolumeSizeGB
+	}
+
+	switch strings.ToLower(d.VolumeType) {
+	case "gp2":
+		// Nothing -- size checked above.
+	case "gp3":
+		if d.IOPs > 16000 {
+			return errors.AssertionFailedf("Iops required for gp3 disk: [3000, 16000]")
+		}
+		if d.IOPs == 0 {
+			// 30000 is a base IOPs for gp3.
+			d.IOPs = 3000
+		}
+		if d.Throughput == 0 {
+			// 125MB/s is base throughput for gp3.
+			d.Throughput = 125
+		}
+	case "io1", "io2":
+		if d.IOPs == 0 {
+			return errors.AssertionFailedf("Iops required for %s disk", d.VolumeType)
+		}
+	default:
+		return errors.Errorf("Unknown EBS volume type %s", d.VolumeType)
+	}
+	return nil
+}
+
+// Type implements flag Value interface.
+func (d *ebsDisk) Type() string {
+	return "JSON"
+}
+
+// String Implements flag Value interface.
+func (d *ebsDisk) String() string {
+	return "EBSDisk"
+}
+
+type ebsVolumeList []*ebsVolume
+
+func (vl *ebsVolumeList) newVolume() *ebsVolume {
+	return &ebsVolume{
+		DeviceName: fmt.Sprintf("/dev/sd%c", 'd'+len(*vl)),
+	}
+}
+
+// Set implements flag Value interface.
+func (vl *ebsVolumeList) Set(s string) error {
+	v := vl.newVolume()
+	if err := v.Disk.Set(s); err != nil {
+		return err
+	}
+	*vl = append(*vl, v)
+	return nil
+}
+
+// Type implements flag Value interface.
+func (vl *ebsVolumeList) Type() string {
+	return "JSON"
+}
+
+// String Implements flag Value interface.
+func (vl *ebsVolumeList) String() string {
+	return "EBSVolumeList"
+}
+
 // providerOpts implements the vm.ProviderFlags interface for aws.Provider.
 type providerOpts struct {
 	Profile string
 	Config  *awsConfig
 
-	MachineType        string
-	SSDMachineType     string
-	CPUOptions         string
-	RemoteUserName     string
-	EBSVolumeType      string
-	EBSVolumeSize      int
-	EBSProvisionedIOPs int
-	UseMultipleDisks   bool
+	MachineType      string
+	SSDMachineType   string
+	CPUOptions       string
+	RemoteUserName   string
+	DefaultEBSVolume ebsVolume
+	EBSVolumes       ebsVolumeList
+	UseMultipleDisks bool
 
 	// Use specified ImageAMI when provisioning.
 	// Overrides config.json AMI.
@@ -122,7 +219,6 @@ var defaultCreateZones = []string{
 // somewhat complicated because different EC2 regions may as well
 // be parallel universes.
 func (o *providerOpts) ConfigureCreateFlags(flags *pflag.FlagSet) {
-
 	// m5.xlarge is a 4core, 16Gb instance, approximately equal to a GCE n1-standard-4
 	flags.StringVar(&o.MachineType, ProviderName+"-machine-type", defaultMachineType,
 		"Machine type (see https://aws.amazon.com/ec2/instance-types/)")
@@ -139,13 +235,17 @@ func (o *providerOpts) ConfigureCreateFlags(flags *pflag.FlagSet) {
 	flags.StringVar(&o.RemoteUserName, ProviderName+"-user",
 		"ubuntu", "Name of the remote user to SSH as")
 
-	flags.StringVar(&o.EBSVolumeType, ProviderName+"-ebs-volume-type",
-		"gp2", "Type of the EBS volume, only used if local-ssd=false")
-	flags.IntVar(&o.EBSVolumeSize, ProviderName+"-ebs-volume-size",
-		500, "Size in GB of EBS volume, only used if local-ssd=false")
-	flags.IntVar(&o.EBSProvisionedIOPs, ProviderName+"-ebs-iops",
-		1000, "Number of IOPs to provision, only used if "+ProviderName+
-			"-ebs-volume-type=io1")
+	flags.StringVar(&o.DefaultEBSVolume.Disk.VolumeType, ProviderName+"-ebs-volume-type",
+		"", "Type of the EBS volume, only used if local-ssd=false")
+	flags.IntVar(&o.DefaultEBSVolume.Disk.VolumeSize, ProviderName+"-ebs-volume-size",
+		ebsDefaultVolumeSizeGB, "Size in GB of EBS volume, only used if local-ssd=false")
+	flags.IntVar(&o.DefaultEBSVolume.Disk.IOPs, ProviderName+"-ebs-iops",
+		0, "Number of IOPs to provision for supported disk types (io1, io2, gp3)")
+	flags.IntVar(&o.DefaultEBSVolume.Disk.Throughput, ProviderName+"-ebs-throughput",
+		0, "Additional throughput to provision, in MiB/s")
+
+	flags.VarP(&o.EBSVolumes, ProviderName+"-ebs-volume", "",
+		"Additional EBS disk to attached; specified as JSON: {VolumeType=io2,VolumeSize=213,Iops=321}")
 
 	flags.StringSliceVar(&o.CreateZones, ProviderName+"-zones", nil,
 		fmt.Sprintf("aws availability zones to use for cluster creation. If zones are formatted\n"+
@@ -723,21 +823,33 @@ func (p *Provider) runInstance(name string, zone string, opts vm.CreateOpts) err
 
 	// The local NVMe devices are automatically mapped.  Otherwise, we need to map an EBS data volume.
 	if !opts.SSDOpts.UseLocalSSD {
-		var ebsParams string
-		switch t := p.opts.EBSVolumeType; t {
-		case "gp2":
-			ebsParams = fmt.Sprintf("{VolumeSize=%d,VolumeType=%s,DeleteOnTermination=true}",
-				p.opts.EBSVolumeSize, t)
-		case "io1", "io2":
-			ebsParams = fmt.Sprintf("{VolumeSize=%d,VolumeType=%s,Iops=%d,DeleteOnTermination=true}",
-				p.opts.EBSVolumeSize, t, p.opts.EBSProvisionedIOPs)
-		default:
-			return errors.Errorf("Unknown EBS volume type %s", t)
+		if len(p.opts.EBSVolumes) == 0 && p.opts.DefaultEBSVolume.Disk.VolumeType == "" {
+			p.opts.DefaultEBSVolume.Disk.VolumeType = "gp2"
+		}
+
+		if p.opts.DefaultEBSVolume.Disk.VolumeType != "" {
+			// Add default volume to the list of volumes we'll setup.
+			v := p.opts.EBSVolumes.newVolume()
+			v.Disk = p.opts.DefaultEBSVolume.Disk
+			p.opts.EBSVolumes = append(p.opts.EBSVolumes, v)
+		}
+
+		mapping, err := json.Marshal(p.opts.EBSVolumes)
+		if err != nil {
+			return err
+		}
+
+		deviceMapping, err := ioutil.TempFile("", "aws-block-device-mapping")
+		if err != nil {
+			return err
+		}
+		defer deviceMapping.Close()
+		if _, err := deviceMapping.Write(mapping); err != nil {
+			return err
 		}
 		args = append(args,
 			"--block-device-mapping",
-			// Size is measured in GB.  gp2 type derives guaranteed iops from size.
-			"DeviceName=/dev/sdd,Ebs="+ebsParams,
+			"file://"+deviceMapping.Name(),
 		)
 	}
 	return p.runJSONCommand(args, &data)