m/test/launch: make disk size configurable

This adds a flag to the launch-cluster command to specify the disk size 
of the VMs. This takes advantage of the previously added data partition 
growing feature.

Fixes: https://github.com/monogon-dev/monogon/issues/309
Change-Id: Iecf8d8c186af16dfa9ed0418ec96c51e58900052
Reviewed-on: https://review.monogon.dev/c/monogon/+/3351
Reviewed-by: Serge Bazanski <serge@monogon.tech>
Tested-by: Jenkins CI
diff --git a/metropolis/test/launch/cli/launch-cluster/main.go b/metropolis/test/launch/cli/launch-cluster/main.go
index cceebd1..3aeb319 100644
--- a/metropolis/test/launch/cli/launch-cluster/main.go
+++ b/metropolis/test/launch/cli/launch-cluster/main.go
@@ -72,7 +72,7 @@
 	})
 }
 
-func sizeFlagMiB(p *int, name string, usage string) {
+func memoryMiBFlag(p *int, name string, usage string) {
 	flag.Func(name, usage, func(val string) error {
 		multiplier := 1
 		switch {
@@ -91,6 +91,28 @@
 	})
 }
 
+func diskBytesFlag(p *uint64, name string, usage string) {
+	flag.Func(name, usage, func(val string) error {
+		var multiplier uint64
+		switch {
+		case strings.HasSuffix(val, "M"):
+			multiplier = 1024 * 1024
+		case strings.HasSuffix(val, "G"):
+			multiplier = 1024 * 1024 * 1024
+		case strings.HasSuffix(val, "T"):
+			multiplier = 1024 * 1024 * 1024 * 1024
+		default:
+			return errors.New("must have suffix M for MiB, G for GiB or T for TiB")
+		}
+		intVal, err := strconv.ParseUint(val[:len(val)-1], 10, 64)
+		if err != nil {
+			return err
+		}
+		*p = multiplier * intVal
+		return nil
+	})
+}
+
 func main() {
 	clusterConfig := cpb.ClusterConfiguration{}
 	opts := mlaunch.ClusterOptions{
@@ -104,7 +126,8 @@
 	flagdefs.StorageSecurityPolicyVar(flag.CommandLine, &clusterConfig.StorageSecurityPolicy, "storage-security", cpb.ClusterConfiguration_STORAGE_SECURITY_POLICY_NEEDS_INSECURE, "Storage security policy to set on cluster")
 	flag.IntVar(&opts.Node.CPUs, "cpu", 1, "Number of virtual CPUs of each node")
 	flag.IntVar(&opts.Node.ThreadsPerCPU, "threads-per-cpu", 1, "Number of threads per CPU")
-	sizeFlagMiB(&opts.Node.MemoryMiB, "ram", "RAM size of each node, with suffix M for MiB or G for GiB")
+	memoryMiBFlag(&opts.Node.MemoryMiB, "ram", "RAM size of each node, with suffix M for MiB or G for GiB")
+	diskBytesFlag(&opts.Node.DiskBytes, "disk", "Disk size of each node, with suffix M for MiB, G for GiB or T for TiB")
 	nodeSetFlag(&consensusMemberList, "consensus-member", "List of nodes which get the Consensus Member role. Example: 0,3-5")
 	nodeSetFlag(&kubernetesControllerList, "kubernetes-controller", "List of nodes which get the Kubernetes Controller role. Example: 0,3-5")
 	nodeSetFlag(&kubernetesWorkerList, "kubernetes-worker", "List of nodes which get the Kubernetes Worker role. Example: 0,3-5")
diff --git a/metropolis/test/launch/cluster.go b/metropolis/test/launch/cluster.go
index 29096b8..e0a0473 100644
--- a/metropolis/test/launch/cluster.go
+++ b/metropolis/test/launch/cluster.go
@@ -72,6 +72,10 @@
 	// MemoryMiB is the RAM size in MiB of the VM.
 	MemoryMiB int
 
+	// DiskBytes contains the size of the root disk in bytes or zero if the
+	// unmodified image size is used.
+	DiskBytes uint64
+
 	// Ports contains the port mapping where to expose the internal ports of the VM to
 	// the host. See IdentityPortMap() and ConflictFreePortMap(). Ignored when
 	// ConnectToSocket is set.
@@ -146,7 +150,7 @@
 // files required to preserve its state, a level below the chosen path ld. The
 // node's socket directory is similarily created a level below sd. It may
 // return an I/O error.
-func setupRuntime(ld, sd string) (*NodeRuntime, error) {
+func setupRuntime(ld, sd string, diskBytes uint64) (*NodeRuntime, error) {
 	// Create a temporary directory to keep all the runtime files.
 	stdp, err := os.MkdirTemp(ld, "node_state*")
 	if err != nil {
@@ -154,6 +158,12 @@
 	}
 
 	// Initialize the node's storage with a prebuilt image.
+	st, err := os.Stat(xNodeImagePath)
+	if err != nil {
+		return nil, fmt.Errorf("cannot read image file: %w", err)
+	}
+	diskBytes = max(diskBytes, uint64(st.Size()))
+
 	di := filepath.Join(stdp, "image.qcow2")
 	launch.Log("Cluster: generating node QCOW2 snapshot image: %s -> %s", xNodeImagePath, di)
 
@@ -162,7 +172,7 @@
 		return nil, fmt.Errorf("while opening image for writing: %w", err)
 	}
 	defer df.Close()
-	if err := qcow2.Generate(df, qcow2.GenerateWithBackingFile(xNodeImagePath)); err != nil {
+	if err := qcow2.Generate(df, qcow2.GenerateWithBackingFile(xNodeImagePath), qcow2.GenerateWithFileSize(diskBytes)); err != nil {
 		return nil, fmt.Errorf("while creating copy-on-write node image: %w", err)
 	}
 
@@ -242,7 +252,7 @@
 
 	// If it's the node's first start, set up its runtime directories.
 	if options.Runtime == nil {
-		r, err := setupRuntime(ld, sd)
+		r, err := setupRuntime(ld, sd, options.DiskBytes)
 		if err != nil {
 			return fmt.Errorf("while setting up node runtime: %w", err)
 		}