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

@@ -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, "-")

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)
 		}
 	}

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)
 		}

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, "-")
+}

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",

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 "
 		}

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")

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)
 	}