From 791b58324e422fd53a955ad1eda36f08ad40fc92 Mon Sep 17 00:00:00 2001
From: Radu Berinde <radu@cockroachlabs.com>
Date: Mon, 25 Oct 2021 14:44:39 -0700
Subject: [PATCH] roachprod: support multiple local clusters

This change adds support for multiple local clusters. Local cluster
names must be either "local" or of the form "local-foo".

When the cluster is named "local", the node directories stay in the
same place, e.g. `~/local/1`. If the cluster is named "local-foo",
node directories are like `~/local/foo-1`.

Fixes #71945.

Release note: None
---
 pkg/cmd/roachprod/cloud/cluster_cloud.go    |  9 +-
 pkg/cmd/roachprod/cloud/gc.go               |  2 +-
 pkg/cmd/roachprod/clusters_cache.go         |  2 +-
 pkg/cmd/roachprod/config/config.go          | 19 +++-
 pkg/cmd/roachprod/install/BUILD.bazel       |  1 +
 pkg/cmd/roachprod/install/cluster_synced.go | 23 +++--
 pkg/cmd/roachprod/install/cockroach.go      | 24 +++--
 pkg/cmd/roachprod/install/download.go       | 10 ++-
 pkg/cmd/roachprod/main.go                   | 71 ++++++++-------
 pkg/cmd/roachprod/vm/aws/BUILD.bazel        |  1 +
 pkg/cmd/roachprod/vm/azure/BUILD.bazel      |  1 +
 pkg/cmd/roachprod/vm/local/BUILD.bazel      |  1 +
 pkg/cmd/roachprod/vm/local/local.go         | 99 +++++++++++++++++----
 pkg/cmd/roachprod/vm/vm.go                  |  3 +
 14 files changed, 182 insertions(+), 84 deletions(-)

diff --git a/pkg/cmd/roachprod/cloud/cluster_cloud.go b/pkg/cmd/roachprod/cloud/cluster_cloud.go
index 761be65ff290..7a13ece4aa3d 100644
--- a/pkg/cmd/roachprod/cloud/cluster_cloud.go
+++ b/pkg/cmd/roachprod/cloud/cluster_cloud.go
@@ -150,16 +150,17 @@ func (c *Cluster) PrintDetails() {
 	}
 }
 
-// IsLocal TODO(peter): document
+// IsLocal returns true if c is a local cluster.
 func (c *Cluster) IsLocal() bool {
-	return c.Name == config.Local
+	return config.IsLocalClusterName(c.Name)
 }
 
 const vmNameFormat = "user-<clusterid>-<nodeid>"
 
-func namesFromVM(v vm.VM) (string, string, error) {
+// namesFromVM determines the user name and the cluster name from a VM.
+func namesFromVM(v vm.VM) (userName string, clusterName string, _ error) {
 	if v.IsLocal() {
-		return config.Local, config.Local, nil
+		return config.Local, v.LocalClusterName, nil
 	}
 	name := v.Name
 	parts := strings.Split(name, "-")
diff --git a/pkg/cmd/roachprod/cloud/gc.go b/pkg/cmd/roachprod/cloud/gc.go
index dc7357e90c80..ae41bdbe70de 100644
--- a/pkg/cmd/roachprod/cloud/gc.go
+++ b/pkg/cmd/roachprod/cloud/gc.go
@@ -276,7 +276,7 @@ func GCClusters(cloud *Cloud, dryrun bool) error {
 
 	var names []string
 	for name := range cloud.Clusters {
-		if name != config.Local {
+		if !config.IsLocalClusterName(name) {
 			names = append(names, name)
 		}
 	}
diff --git a/pkg/cmd/roachprod/clusters_cache.go b/pkg/cmd/roachprod/clusters_cache.go
index 39179aeb2c83..3d710f19a492 100644
--- a/pkg/cmd/roachprod/clusters_cache.go
+++ b/pkg/cmd/roachprod/clusters_cache.go
@@ -105,7 +105,7 @@ func loadClusters() error {
 
 		install.Clusters[sc.Name] = sc
 
-		if local.IsLocal(c.Name) {
+		if config.IsLocalClusterName(c.Name) {
 			// Add the local cluster to the local provider.
 			local.AddCluster(c)
 		}
diff --git a/pkg/cmd/roachprod/config/config.go b/pkg/cmd/roachprod/config/config.go
index 4261490d8456..e48dad912e70 100644
--- a/pkg/cmd/roachprod/config/config.go
+++ b/pkg/cmd/roachprod/config/config.go
@@ -13,6 +13,7 @@ package config
 import (
 	"log"
 	"os/user"
+	"strings"
 )
 
 var (
@@ -39,7 +40,8 @@ const (
 	// EmailDomain is used to form the full account name for GCE and Slack.
 	EmailDomain = "@cockroachlabs.com"
 
-	// Local is the name of the local cluster.
+	// Local is the prefix used to identify local clusters.
+	// It is also used as the zone for local clusters.
 	Local = "local"
 
 	// ClustersDir is the directory where we cache information about clusters.
@@ -60,3 +62,18 @@ const (
 	// listening for HTTP connections for the Admin UI.
 	DefaultAdminUIPort = 26258
 )
+
+// IsLocalClusterName returns true if the given name is for a local cluster.
+//
+// Local cluster names are either "local" or start with a "local-" prefix.
+func IsLocalClusterName(clusterName string) bool {
+	if !strings.HasPrefix(clusterName, Local) {
+		return false
+	}
+	clusterName = strings.TrimPrefix(clusterName, Local)
+	if clusterName == "" {
+		// clusterName is "local"
+		return true
+	}
+	return len(clusterName) >= 2 && strings.HasPrefix(clusterName, "-")
+}
diff --git a/pkg/cmd/roachprod/install/BUILD.bazel b/pkg/cmd/roachprod/install/BUILD.bazel
index 81c2b1a8600b..f939e31c586a 100644
--- a/pkg/cmd/roachprod/install/BUILD.bazel
+++ b/pkg/cmd/roachprod/install/BUILD.bazel
@@ -26,6 +26,7 @@ go_library(
         "//pkg/cmd/roachprod/ssh",
         "//pkg/cmd/roachprod/ui",
         "//pkg/cmd/roachprod/vm/aws",
+        "//pkg/cmd/roachprod/vm/local",
         "//pkg/util/envutil",
         "//pkg/util/httputil",
         "//pkg/util/log",
diff --git a/pkg/cmd/roachprod/install/cluster_synced.go b/pkg/cmd/roachprod/install/cluster_synced.go
index 49b945657ba3..6a7f9c31d8b4 100644
--- a/pkg/cmd/roachprod/install/cluster_synced.go
+++ b/pkg/cmd/roachprod/install/cluster_synced.go
@@ -36,6 +36,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/ssh"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/ui"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/vm/aws"
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/vm/local"
 	clog "github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
@@ -103,7 +104,11 @@ func (c *SyncedCluster) locality(index int) string {
 // TODO(tschottdorf): roachprod should cleanly encapsulate the home directory
 // which is currently the biggest culprit for awkward one-offs.
 func (c *SyncedCluster) IsLocal() bool {
-	return c.Name == config.Local
+	return config.IsLocalClusterName(c.Name)
+}
+
+func (c *SyncedCluster) localVMDir(nodeIdx int) string {
+	return local.VMDir(c.Name, nodeIdx)
 }
 
 // ServerNodes is the fully expanded, ordered list of nodes that any given
@@ -238,7 +243,7 @@ func (c *SyncedCluster) Wipe(preserveCerts bool) {
 				dirs = append(dirs, "certs*")
 			}
 			for _, dir := range dirs {
-				cmd += fmt.Sprintf(`rm -fr ${HOME}/local/%d/%s ;`, c.Nodes[i], dir)
+				cmd += fmt.Sprintf(`rm -fr %s/%s`, c.localVMDir(c.Nodes[i]), dir)
 			}
 		} else {
 			cmd = `sudo find /mnt/data* -maxdepth 1 -type f -exec rm -f {} \; &&
@@ -523,7 +528,7 @@ func (c *SyncedCluster) Run(stdout, stderr io.Writer, nodes []int, title, cmd st
 		nodeCmd := fmt.Sprintf(`export ROACHPROD=%d%s GOTRACEBACK=crash && bash -c %s`,
 			nodes[i], c.Tag, ssh.Escape1(expandedCmd))
 		if c.IsLocal() {
-			nodeCmd = fmt.Sprintf("cd ${HOME}/local/%d ; %s", nodes[i], nodeCmd)
+			nodeCmd = fmt.Sprintf("cd %s; %s", c.localVMDir(nodes[i]), nodeCmd)
 		}
 
 		if stream {
@@ -825,7 +830,7 @@ fi
 func (c *SyncedCluster) DistributeCerts() {
 	dir := ""
 	if c.IsLocal() {
-		dir = `${HOME}/local/1`
+		dir = c.localVMDir(1)
 	}
 
 	// Check to see if the certs have already been initialized.
@@ -893,7 +898,7 @@ func (c *SyncedCluster) DistributeCerts() {
 
 		var cmd string
 		if c.IsLocal() {
-			cmd = `cd ${HOME}/local/1 ; `
+			cmd = fmt.Sprintf(`cd %s ; `, c.localVMDir(1))
 		}
 		cmd += fmt.Sprintf(`
 rm -fr certs
@@ -961,7 +966,7 @@ tar cvf certs.tar certs
 		sess.SetStdin(bytes.NewReader(certsTar))
 		var cmd string
 		if c.IsLocal() {
-			cmd = fmt.Sprintf(`cd ${HOME}/local/%d ; `, nodes[i])
+			cmd = fmt.Sprintf(`cd %s ; `, c.localVMDir(nodes[i]))
 		}
 		cmd += `tar xf -`
 		if out, err := sess.CombinedOutput(cmd); err != nil {
@@ -1084,7 +1089,7 @@ func (c *SyncedCluster) Put(src, dest string) {
 				if filepath.IsAbs(dest) {
 					to = dest
 				} else {
-					to = fmt.Sprintf(os.ExpandEnv("${HOME}/local/%d/%s"), c.Nodes[i], dest)
+					to = filepath.Join(c.localVMDir(c.Nodes[i]), dest)
 				}
 				// Remove the destination if it exists, ignoring errors which we'll
 				// handle via the os.Symlink() call.
@@ -1392,7 +1397,7 @@ func (c *SyncedCluster) Get(src, dest string) {
 
 			if c.IsLocal() {
 				if !filepath.IsAbs(src) {
-					src = filepath.Join(fmt.Sprintf(os.ExpandEnv("${HOME}/local/%d"), c.Nodes[i]), src)
+					src = filepath.Join(c.localVMDir(c.Nodes[i]), src)
 				}
 
 				var copy func(src, dest string, info os.FileInfo) error
@@ -1611,7 +1616,7 @@ func (c *SyncedCluster) SSH(sshArgs, args []string) error {
 		allArgs = []string{
 			"/bin/bash", "-c",
 		}
-		cmd := fmt.Sprintf("cd ${HOME}/local/%d ; ", c.Nodes[0])
+		cmd := fmt.Sprintf("cd %s ; ", c.localVMDir(c.Nodes[0]))
 		if len(args) == 0 /* interactive */ {
 			cmd += "/bin/bash "
 		}
diff --git a/pkg/cmd/roachprod/install/cockroach.go b/pkg/cmd/roachprod/install/cockroach.go
index 5d2c37fe4aff..225e29ca0469 100644
--- a/pkg/cmd/roachprod/install/cockroach.go
+++ b/pkg/cmd/roachprod/install/cockroach.go
@@ -53,7 +53,7 @@ func cockroachNodeBinary(c *SyncedCluster, node int) string {
 		return "./" + config.Binary
 	}
 
-	path := filepath.Join(fmt.Sprintf(os.ExpandEnv("${HOME}/local/%d"), node), config.Binary)
+	path := filepath.Join(c.localVMDir(node), config.Binary)
 	if _, err := os.Stat(path); err == nil {
 		return path
 	}
@@ -236,27 +236,25 @@ func (Cockroach) NodeDir(c *SyncedCluster, index, storeIndex int) string {
 		if storeIndex != 1 {
 			panic("Cockroach.NodeDir only supports one store for local deployments")
 		}
-		return os.ExpandEnv(fmt.Sprintf("${HOME}/local/%d/data", index))
+		return filepath.Join(c.localVMDir(index), "data")
 	}
 	return fmt.Sprintf("/mnt/data%d/cockroach", storeIndex)
 }
 
 // LogDir implements the ClusterImpl.NodeDir interface.
 func (Cockroach) LogDir(c *SyncedCluster, index int) string {
-	dir := "logs"
 	if c.IsLocal() {
-		dir = os.ExpandEnv(fmt.Sprintf("${HOME}/local/%d/logs", index))
+		return filepath.Join(c.localVMDir(index), "logs")
 	}
-	return dir
+	return "logs"
 }
 
 // CertsDir implements the ClusterImpl.NodeDir interface.
 func (Cockroach) CertsDir(c *SyncedCluster, index int) string {
-	dir := "certs"
 	if c.IsLocal() {
-		dir = os.ExpandEnv(fmt.Sprintf("${HOME}/local/%d/certs", index))
+		return filepath.Join(c.localVMDir(index), "certs")
 	}
-	return dir
+	return "certs"
 }
 
 // NodeURL implements the ClusterImpl.NodeDir interface.
@@ -321,7 +319,7 @@ func (r Cockroach) SQL(c *SyncedCluster, args []string) error {
 
 		var cmd string
 		if c.IsLocal() {
-			cmd = fmt.Sprintf(`cd ${HOME}/local/%d ; `, c.Nodes[nodeIdx])
+			cmd = fmt.Sprintf(`cd %s ; `, c.localVMDir(c.Nodes[nodeIdx]))
 		}
 		cmd += cockroachNodeBinary(c, c.Nodes[nodeIdx]) + " sql --url " +
 			r.NodeURL(c, "localhost", r.NodePort(c, c.Nodes[nodeIdx])) + " " +
@@ -374,7 +372,7 @@ func (h *crdbInstallHelper) startNode(
 		sess.SetStdin(strings.NewReader(startCmd))
 		var cmd string
 		if h.c.IsLocal() {
-			cmd = fmt.Sprintf(`cd ${HOME}/local/%d ; `, nodes[nodeIdx])
+			cmd = fmt.Sprintf(`cd %s ; `, h.c.localVMDir(nodes[nodeIdx]))
 		}
 		cmd += `cat > cockroach.sh && chmod +x cockroach.sh`
 		if out, err := sess.CombinedOutput(cmd); err != nil {
@@ -394,7 +392,7 @@ func (h *crdbInstallHelper) startNode(
 
 	var cmd string
 	if h.c.IsLocal() {
-		cmd = fmt.Sprintf(`cd ${HOME}/local/%d ; `, nodes[nodeIdx])
+		cmd = fmt.Sprintf(`cd %s ; `, h.c.localVMDir(nodes[nodeIdx]))
 	}
 	cmd += "./cockroach.sh"
 	out, err := sess.CombinedOutput(cmd)
@@ -623,7 +621,7 @@ func (h *crdbInstallHelper) generateClusterSettingCmd(nodeIdx int) string {
 
 	var clusterSettingCmd string
 	if h.c.IsLocal() {
-		clusterSettingCmd = `cd ${HOME}/local/1 ; `
+		clusterSettingCmd = fmt.Sprintf(`cd %s ; `, h.c.localVMDir(1))
 	}
 
 	binary := cockroachNodeBinary(h.c, nodes[nodeIdx])
@@ -648,7 +646,7 @@ func (h *crdbInstallHelper) generateInitCmd(nodeIdx int) string {
 
 	var initCmd string
 	if h.c.IsLocal() {
-		initCmd = `cd ${HOME}/local/1 ; `
+		initCmd = fmt.Sprintf(`cd %s ; `, h.c.localVMDir(1))
 	}
 
 	path := fmt.Sprintf("%s/%s", h.c.Impl.NodeDir(h.c, nodes[nodeIdx], 1 /* storeIndex */), "cluster-bootstrapped")
diff --git a/pkg/cmd/roachprod/install/download.go b/pkg/cmd/roachprod/install/download.go
index 43f5373b9cca..9ec2b10b808b 100644
--- a/pkg/cmd/roachprod/install/download.go
+++ b/pkg/cmd/roachprod/install/download.go
@@ -17,6 +17,8 @@ import (
 	"os"
 	"path"
 	"path/filepath"
+
+	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/vm/local"
 )
 
 const (
@@ -71,11 +73,11 @@ func Download(c *SyncedCluster, sourceURLStr string, sha string, dest string) er
 		return err
 	}
 
-	// If we are local and the destination is relative, then copy
-	// the file from the download node to the other nodes
+	// If we are local and the destination is relative, then copy the file from
+	// the download node to the other nodes.
 	if c.IsLocal() && !filepath.IsAbs(dest) {
-		// ~/local/1/./bar.txt
-		src := fmt.Sprintf(os.ExpandEnv("${HOME}/local/%d/%s"), downloadNodes[0], dest)
+		// Eg ~/local/1/./bar.txt or ~/local/local-foo/1/./bar.txt
+		src := filepath.Join(local.VMDir(c.Name, downloadNodes[0]), dest)
 		cpCmd := fmt.Sprintf(`cp "%s" "%s"`, src, dest)
 		return c.Run(os.Stdout, os.Stderr, c.Nodes[1:], "copying to remaining nodes", cpCmd)
 	}
diff --git a/pkg/cmd/roachprod/main.go b/pkg/cmd/roachprod/main.go
index b3153a914a94..a8467edd4edc 100644
--- a/pkg/cmd/roachprod/main.go
+++ b/pkg/cmd/roachprod/main.go
@@ -84,6 +84,7 @@ var (
 	username          string
 	dryrun            bool
 	destroyAllMine    bool
+	destroyAllLocal   bool
 	extendLifetime    time.Duration
 	wipePreserveCerts bool
 	listDetails       bool
@@ -204,10 +205,6 @@ func verifyClusterName(clusterName string) error {
 	if clusterName == "" {
 		return fmt.Errorf("cluster name cannot be blank")
 	}
-	if clusterName == config.Local {
-		return nil
-	}
-
 	alphaNum, err := regexp.Compile(`^[a-zA-Z0-9\-]+$`)
 	if err != nil {
 		return err
@@ -216,6 +213,10 @@ func verifyClusterName(clusterName string) error {
 		return errors.Errorf("cluster name must match %s", alphaNum.String())
 	}
 
+	if config.IsLocalClusterName(clusterName) {
+		return nil
+	}
+
 	// Use the vm.Provider account names, or --username.
 	var accounts []string
 	if len(username) > 0 {
@@ -347,9 +348,8 @@ Cloud Clusters
 Local Clusters
 
   A local cluster stores the per-node data in ${HOME}/local on the machine
-  roachprod is being run on. Local clusters requires local ssh access. Unlike
-  cloud clusters there can be only a single local cluster, the local cluster is
-  always named "local", and has no expiration (unlimited lifetime).
+	roachprod is being run on. Local clusters names are always either "local" or
+	start with "local-" (e.g. "local-foo"). Local clusters do not expire.
 `,
 	Args: cobra.ExactArgs(1),
 	Run: wrap(func(cmd *cobra.Command, args []string) (retErr error) {
@@ -365,7 +365,7 @@ Local Clusters
 		createVMOpts.ClusterName = clusterName
 
 		defer func() {
-			if retErr == nil || clusterName == config.Local {
+			if retErr == nil || config.IsLocalClusterName(clusterName) {
 				return
 			}
 			if errors.HasType(retErr, (*clusterAlreadyExistsError)(nil)) {
@@ -379,7 +379,7 @@ Local Clusters
 			}
 		}()
 
-		if clusterName != config.Local {
+		if !config.IsLocalClusterName(clusterName) {
 			cld, err := cloud.ListCloud()
 			if err != nil {
 				return err
@@ -411,14 +411,8 @@ Local Clusters
 			return createErr
 		}
 
-		// Just create directories for the local cluster as there's no need for ssh.
-		if clusterName == config.Local {
-			for i := 0; i < numNodes; i++ {
-				err := os.MkdirAll(fmt.Sprintf(os.ExpandEnv("${HOME}/local/%d"), i+1), 0755)
-				if err != nil {
-					return err
-				}
-			}
+		if config.IsLocalClusterName(clusterName) {
+			// No need for ssh for local clusters.
 			return nil
 		}
 		return setupSSH(clusterName)
@@ -511,18 +505,19 @@ func cleanupFailedCreate(clusterName string) error {
 }
 
 var destroyCmd = &cobra.Command{
-	Use:   "destroy [ --all-mine | <cluster 1> [<cluster 2> ...] ]",
+	Use:   "destroy [ --all-mine | --all-local | <cluster 1> [<cluster 2> ...] ]",
 	Short: "destroy clusters",
 	Long: `Destroy one or more local or cloud-based clusters.
 
 The destroy command accepts the names of the clusters to destroy. Alternatively,
-the --all-mine flag can be provided to destroy all clusters that are owned by the
-current user.
+the --all-mine flag can be provided to destroy all (non-local) clusters that are
+owned by the current user, or the --all-local flag can be provided to destroy
+all local clusters.
 
 Destroying a cluster releases the resources for a cluster. For a cloud-based
 cluster the machine and associated disk resources are freed. For a local
-cluster, any processes started by roachprod are stopped, and the ${HOME}/local
-directory is removed.
+cluster, any processes started by roachprod are stopped, and the node
+directories inside ${HOME}/local directory are removed.
 `,
 	Args: cobra.ArbitraryArgs,
 	Run: wrap(func(cmd *cobra.Command, args []string) error {
@@ -530,12 +525,15 @@ directory is removed.
 		// We want to avoid running ListCloud() if we are only trying to destroy a
 		// local cluster.
 		var cld *cloud.Cloud
-		switch len(args) {
-		case 0:
-			if !destroyAllMine {
-				return errors.New("no cluster name provided")
-			}
 
+		switch {
+		case destroyAllMine:
+			if len(args) != 0 {
+				return errors.New("--all-mine cannot be combined with cluster names")
+			}
+			if destroyAllLocal {
+				return errors.New("--all-mine cannot be combined with --all-local")
+			}
 			destroyPattern, err := userClusterNameRegexp()
 			if err != nil {
 				return err
@@ -547,9 +545,16 @@ directory is removed.
 			clusters := cld.Clusters.FilterByName(destroyPattern)
 			clusterNames = clusters.Names()
 
+		case destroyAllLocal:
+			if len(args) != 0 {
+				return errors.New("--all-local cannot be combined with cluster names")
+			}
+
+			clusterNames = local.Clusters()
+
 		default:
-			if destroyAllMine {
-				return errors.New("--all-mine cannot be combined with cluster names")
+			if len(args) == 0 {
+				return errors.New("no cluster name provided")
 			}
 
 			for _, clusterName := range args {
@@ -566,7 +571,7 @@ directory is removed.
 			func(ctx context.Context, idx int) error {
 				name := clusterNames[idx]
 				var c *cloud.Cluster
-				if local.IsLocal(name) {
+				if config.IsLocalClusterName(name) {
 					// This is an optimization to avoid calling ListCloud.
 					c = local.GetCluster(name)
 				} else {
@@ -1277,7 +1282,7 @@ environments and will fall back to a no-op.`,
 			return err
 		}
 
-		if clusterName == config.Local {
+		if config.IsLocalClusterName(clusterName) {
 			return nil
 		}
 
@@ -1902,7 +1907,9 @@ func main() {
 	}
 
 	destroyCmd.Flags().BoolVarP(&destroyAllMine,
-		"all-mine", "m", false, "Destroy all clusters belonging to the current user")
+		"all-mine", "m", false, "Destroy all non-local clusters belonging to the current user")
+	destroyCmd.Flags().BoolVarP(&destroyAllLocal,
+		"all-local", "l", false, "Destroy all local clusters")
 
 	extendCmd.Flags().DurationVarP(&extendLifetime,
 		"lifetime", "l", 12*time.Hour, "Lifetime of the cluster")
diff --git a/pkg/cmd/roachprod/vm/aws/BUILD.bazel b/pkg/cmd/roachprod/vm/aws/BUILD.bazel
index b0af253894a7..71cc328fa157 100644
--- a/pkg/cmd/roachprod/vm/aws/BUILD.bazel
+++ b/pkg/cmd/roachprod/vm/aws/BUILD.bazel
@@ -13,6 +13,7 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachprod/vm/aws",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/cmd/roachprod/config",
         "//pkg/cmd/roachprod/vm",
         "//pkg/cmd/roachprod/vm/flagstub",
         "//pkg/util/retry",
diff --git a/pkg/cmd/roachprod/vm/azure/BUILD.bazel b/pkg/cmd/roachprod/vm/azure/BUILD.bazel
index 2a32fb4cebfa..31152594496c 100644
--- a/pkg/cmd/roachprod/vm/azure/BUILD.bazel
+++ b/pkg/cmd/roachprod/vm/azure/BUILD.bazel
@@ -13,6 +13,7 @@ go_library(
     importpath = "github.com/cockroachdb/cockroach/pkg/cmd/roachprod/vm/azure",
     visibility = ["//visibility:public"],
     deps = [
+        "//pkg/cmd/roachprod/config",
         "//pkg/cmd/roachprod/vm",
         "//pkg/cmd/roachprod/vm/flagstub",
         "//pkg/util/syncutil",
diff --git a/pkg/cmd/roachprod/vm/local/BUILD.bazel b/pkg/cmd/roachprod/vm/local/BUILD.bazel
index b2cc832b1c3d..63b3c4119c14 100644
--- a/pkg/cmd/roachprod/vm/local/BUILD.bazel
+++ b/pkg/cmd/roachprod/vm/local/BUILD.bazel
@@ -9,6 +9,7 @@ go_library(
         "//pkg/cmd/roachprod/cloud",
         "//pkg/cmd/roachprod/config",
         "//pkg/cmd/roachprod/vm",
+        "//pkg/util",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_spf13_pflag//:pflag",
diff --git a/pkg/cmd/roachprod/vm/local/local.go b/pkg/cmd/roachprod/vm/local/local.go
index 795c89ed9275..aea2d43383c4 100644
--- a/pkg/cmd/roachprod/vm/local/local.go
+++ b/pkg/cmd/roachprod/vm/local/local.go
@@ -13,11 +13,14 @@ package local
 import (
 	"fmt"
 	"os"
+	"path/filepath"
+	"strings"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/cloud"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/config"
 	"github.com/cockroachdb/cockroach/pkg/cmd/roachprod/vm"
+	"github.com/cockroachdb/cockroach/pkg/util"
 	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
 	"github.com/cockroachdb/errors"
 	"github.com/spf13/pflag"
@@ -26,6 +29,26 @@ import (
 // ProviderName is config.Local.
 const ProviderName = config.Local
 
+// VMDir returns the local directory for a given node in a cluster.
+// Node indexes start at 1.
+//
+// If the cluster name is "local", node 1 directory is:
+//   ${HOME}/local/1
+//
+// If the cluster name is "local-foo", node 1 directory is:
+//   ${HOME}/local/foo-1
+func VMDir(clusterName string, nodeIdx int) string {
+	if nodeIdx < 1 {
+		panic("invalid nodeIdx")
+	}
+	localDir := os.ExpandEnv("${HOME}/local")
+	if clusterName == config.Local {
+		return filepath.Join(localDir, fmt.Sprintf("%d", nodeIdx))
+	}
+	name := strings.TrimPrefix(clusterName, config.Local+"-")
+	return filepath.Join(localDir, fmt.Sprintf("%s-%d", name, nodeIdx))
+}
+
 // Init initializes the Local provider and registers it into vm.Providers.
 func Init(storage VMStorage) {
 	vm.Providers[ProviderName] = &Provider{
@@ -34,11 +57,6 @@ func Init(storage VMStorage) {
 	}
 }
 
-// IsLocal returns true if the given cluster name is a local cluster.
-func IsLocal(clusterName string) bool {
-	return clusterName == config.Local
-}
-
 // AddCluster adds the metadata of a local cluster; used when loading the saved
 // metadata for local clusters.
 func AddCluster(cluster *cloud.Cluster) {
@@ -53,6 +71,12 @@ func GetCluster(name string) *cloud.Cluster {
 	return p.clusters[name]
 }
 
+// Clusters returns a list of all known local clusters.
+func Clusters() []string {
+	p := vm.Providers[ProviderName].(*Provider)
+	return p.clusters.Names()
+}
+
 // VMStorage is the interface for saving metadata for local clusters.
 type VMStorage interface {
 	// SaveCluster saves the metadata for a local cluster. It is expected that
@@ -101,26 +125,62 @@ func (p *Provider) Create(names []string, opts vm.CreateOpts) error {
 		Lifetime:  time.Hour,
 		VMs:       make(vm.List, len(names)),
 	}
+
+	if !config.IsLocalClusterName(c.Name) {
+		return errors.Errorf("'%s' is not a valid local cluster name", c.Name)
+	}
+
+	// We will need to assign ports to the nodes, and they must not conflict with
+	// any other local clusters.
+	var portsTaken util.FastIntSet
+	for _, c := range p.clusters {
+		for i := range c.VMs {
+			portsTaken.Add(c.VMs[i].SQLPort)
+			portsTaken.Add(c.VMs[i].AdminUIPort)
+		}
+	}
+	sqlPort := config.DefaultSQLPort
+	adminUIPort := config.DefaultAdminUIPort
+
+	// getPort returns the first available port (starting at *port), and modifies
+	// (*port) to be the following value.
+	getPort := func(port *int) int {
+		for portsTaken.Contains(*port) {
+			(*port)++
+		}
+		result := *port
+		portsTaken.Add(result)
+		(*port)++
+		return result
+	}
+
 	for i := range names {
 		c.VMs[i] = vm.VM{
-			Name:        "localhost",
-			CreatedAt:   now,
-			Lifetime:    time.Hour,
-			PrivateIP:   "127.0.0.1",
-			Provider:    ProviderName,
-			ProviderID:  ProviderName,
-			PublicIP:    "127.0.0.1",
-			RemoteUser:  config.OSUser.Username,
-			VPC:         ProviderName,
-			MachineType: ProviderName,
-			Zone:        ProviderName,
-			SQLPort:     config.DefaultSQLPort + 2*i,
-			AdminUIPort: config.DefaultAdminUIPort + 2*i,
+			Name:             "localhost",
+			CreatedAt:        now,
+			Lifetime:         time.Hour,
+			PrivateIP:        "127.0.0.1",
+			Provider:         ProviderName,
+			ProviderID:       ProviderName,
+			PublicIP:         "127.0.0.1",
+			RemoteUser:       config.OSUser.Username,
+			VPC:              ProviderName,
+			MachineType:      ProviderName,
+			Zone:             ProviderName,
+			SQLPort:          getPort(&sqlPort),
+			AdminUIPort:      getPort(&adminUIPort),
+			LocalClusterName: c.Name,
+		}
+
+		err := os.MkdirAll(VMDir(c.Name, i+1), 0755)
+		if err != nil {
+			return err
 		}
 	}
 	if err := p.storage.SaveCluster(c); err != nil {
 		return err
 	}
+
 	p.clusters[c.Name] = c
 	return nil
 }
@@ -136,9 +196,10 @@ func (p *Provider) DeleteCluster(name string) error {
 	if c == nil {
 		return fmt.Errorf("local cluster %s does not exist", name)
 	}
+	fmt.Printf("Deleting local cluster %s\n", name)
 
 	for i := range c.VMs {
-		err := os.RemoveAll(fmt.Sprintf(os.ExpandEnv("${HOME}/local/%d"), i+1))
+		err := os.RemoveAll(VMDir(c.Name, i+1))
 		if err != nil {
 			return err
 		}
diff --git a/pkg/cmd/roachprod/vm/vm.go b/pkg/cmd/roachprod/vm/vm.go
index f16c5b62db92..c009276b3ed4 100644
--- a/pkg/cmd/roachprod/vm/vm.go
+++ b/pkg/cmd/roachprod/vm/vm.go
@@ -65,6 +65,9 @@ type VM struct {
 	// HTTP traffic for the Admin UI.
 	// Usually config.DefaultAdminUIPort, except for local clusters.
 	AdminUIPort int `json:"adminui_port"`
+
+	// LocalClusterName is only set for VMs in a local cluster.
+	LocalClusterName string `json:"local_cluster_name,omitempty"`
 }
 
 // Name generates the name for the i'th node in a cluster.