roachprod: improve snapshot support for AWS and GCE (#170519)

roachprod: improve snapshot support for AWS and GCE

Author: trunk-io[bot]

pkg/cmd/roachprod/cli/commands.go

@@ -2130,7 +2130,7 @@ func buildSnapshotListCmd() *cobra.Command {
 				return err
 			}
 			for _, snapshot := range snapshots {
-				config.Logger.Printf("found snapshot %s (id: %s)", snapshot.Name, snapshot.ID)
+				config.Logger.Printf("found snapshot %s (id: %s, status: %s)", snapshot.Name, snapshot.ID, snapshot.Status)
 			}
 			return nil
 		}),

pkg/roachprod/roachprod.go

@@ -7,6 +7,7 @@ package roachprod
 
 import (
 	"bytes"
+	"cmp"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -26,6 +27,7 @@ import (
 	"strconv"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/build"
@@ -2415,14 +2417,10 @@ func CreateSnapshot(
 	// probably the predecessor one. Also ensure that any running CRDB processes
 	// have been stopped since we're taking raw disk snapshots cluster-wide.
 
-	// Count total volumes across all nodes for the progress spinner.
-	totalVolumes := 0
-	for _, node := range nodes {
-		totalVolumes += len(c.VMs[node-1].NonBootAttachedVolumes)
-	}
-	spinner := ui.NewDefaultCountingSpinner(l, "creating snapshots", totalVolumes)
+	spinner := ui.NewDefaultCountingSpinner(l, "creating snapshots", len(nodes))
 	defer spinner.Start()()
 
+	var nodesSnapped atomic.Int32
 	volumesSnapshotMu := struct {
 		syncutil.Mutex
 		snapshots []vm.VolumeSnapshot
@@ -2488,14 +2486,14 @@ func CreateSnapshot(
 						}
 						volumesSnapshotMu.Lock()
 						volumesSnapshotMu.snapshots = append(volumesSnapshotMu.snapshots, volumeSnapshot)
-						spinner.CountStatus(len(volumesSnapshotMu.snapshots))
 						volumesSnapshotMu.Unlock()
 						return nil
 					})
 				}
 				if err := g.Wait(); err != nil {
 					return err
 				}
+				spinner.CountStatus(int(nodesSnapped.Add(1)))
 				return nil
 			}); err != nil {
 				res.Err = err
@@ -2649,6 +2647,15 @@ func ApplySnapshots(
 					return err
 				}
 
+				// Sort by storeIdx so volumes are attached in order.
+				// AttachVolume assigns device names sequentially
+				// (/dev/sdd, /dev/sde, ...) based on attachment order,
+				// so attaching in storeIdx order ensures the device-to-
+				// store mapping is consistent with the original cluster.
+				slices.SortFunc(createdVolumes, func(a, b createdVolume) int {
+					return cmp.Compare(a.storeIdx, b.storeIdx)
+				})
+
 				// Attach and mount volumes sequentially since GCE
 				// does not support concurrent instance modifications.
 				for _, cv := range createdVolumes {
@@ -2665,7 +2672,7 @@ func ApplySnapshots(
 						l.Printf(buf.String())
 						return err
 					}
-					l.Printf("mounted %s at %s on %s", cv.volume.ProviderResourceID, mountDir, cVM.ProviderID)
+					l.Printf("mounted %s at %s on %s", cv.volume.ProviderResourceID, mountDir, cVM.Name)
 				}
 				// Save the cluster cache once after all volumes are
 				// attached and mounted for this node.

pkg/roachprod/vm/aws/aws.go

@@ -2658,6 +2658,7 @@ func (p *Provider) CreateVolume(
 		args = append(args, "--volume-type", vco.Type)
 	case "":
 		// Use the default.
+		args = append(args, "--volume-type", defaultEBSVolumeType)
 	default:
 		return vol, errors.Newf("Invalid volume type %q", vco.Type)
 	}
@@ -2754,8 +2755,73 @@ func (p *Provider) DeleteVolume(l *logger.Logger, volume vm.Volume, v *vm.VM) er
 	return nil
 }
 
-func (p *Provider) ListVolumes(l *logger.Logger, vm *vm.VM) ([]vm.Volume, error) {
-	return vm.NonBootAttachedVolumes, nil
+func (p *Provider) ListVolumes(l *logger.Logger, v *vm.VM) ([]vm.Volume, error) {
+	if v.ProviderID == "" {
+		return nil, nil
+	}
+	region := v.Zone[:len(v.Zone)-1]
+
+	// Describe this instance to get block device mappings and the root device
+	// name, which we need to identify (and exclude) the boot volume.
+	var descResp DescribeInstancesOutput
+	descArgs := []string{
+		"ec2", "describe-instances",
+		"--region", region,
+		"--instance-ids", v.ProviderID,
+	}
+	if err := p.runJSONCommand(l, descArgs, &descResp); err != nil {
+		return nil, err
+	}
+	var instance *DescribeInstancesOutputInstance
+	for _, res := range descResp.Reservations {
+		for i := range res.Instances {
+			if res.Instances[i].InstanceID == v.ProviderID {
+				instance = &res.Instances[i]
+				break
+			}
+		}
+	}
+	if instance == nil {
+		l.Printf("WARNING: instance %s not found in describe-instances response for region %s", v.ProviderID, region)
+		return nil, nil
+	}
+
+	// Collect non-boot volume device names from the block device
+	// mappings. The root device is excluded so we only return data volumes.
+	deviceByVolumeID := make(map[string]string)
+	for _, bdm := range instance.BlockDeviceMappings {
+		if bdm.DeviceName != instance.RootDeviceName {
+			deviceByVolumeID[bdm.Disk.VolumeID] = bdm.DeviceName
+		}
+	}
+	if len(deviceByVolumeID) == 0 {
+		return nil, nil
+	}
+
+	// Describe the non-boot volumes to get their full metadata.
+	volsByInstance, err := p.getVolumesForInstances(
+		context.Background(), l, region, []string{v.ProviderID},
+	)
+	if err != nil {
+		return nil, err
+	}
+
+	var volumes []vm.Volume
+	for _, vol := range volsByInstance[v.ProviderID] {
+		if _, ok := deviceByVolumeID[vol.ProviderResourceID]; ok {
+			volumes = append(volumes, vol)
+		}
+	}
+
+	// Sort by device name to ensure deterministic ordering that matches
+	// the mount point assignment (/mnt/data1, /mnt/data2, etc.).
+	slices.SortFunc(volumes, func(a, b vm.Volume) int {
+		return strings.Compare(
+			deviceByVolumeID[a.ProviderResourceID],
+			deviceByVolumeID[b.ProviderResourceID],
+		)
+	})
+	return volumes, nil
 }
 
 type snapshotOutput struct {
@@ -2797,21 +2863,11 @@ func (p *Provider) CreateVolumeSnapshot(
 		return vm.VolumeSnapshot{}, err
 	}
 
-	// Wait for the snapshot to complete before returning. AWS snapshots
-	// are asynchronous and cannot be used to create volumes while pending.
-	waitArgs := []string{
-		"ec2", "wait", "snapshot-completed",
-		"--region", region,
-		"--snapshot-ids", so.SnapshotID,
-	}
-	if _, err := p.runCommand(l, waitArgs); err != nil {
-		return vm.VolumeSnapshot{}, errors.Wrapf(err, "waiting for snapshot %s to complete", so.SnapshotID)
-	}
-
 	return vm.VolumeSnapshot{
 		ID:     so.SnapshotID,
 		Name:   vsco.Name,
 		Region: region,
+		Status: so.State,
 	}, nil
 }
 
@@ -2875,6 +2931,7 @@ func (p *Provider) ListVolumeSnapshots(
 					ID:     so.SnapshotID,
 					Name:   name,
 					Region: r,
+					Status: so.State,
 				})
 			}
 			mu.Lock()

pkg/roachprod/vm/gce/gcloud.go

@@ -743,8 +743,9 @@ func (p *Provider) CreateVolumeSnapshot(
 		return vm.VolumeSnapshot{}, err
 	}
 	return vm.VolumeSnapshot{
-		ID:   createJsonResponse.ID,
-		Name: createJsonResponse.Name,
+		ID:     createJsonResponse.ID,
+		Name:   createJsonResponse.Name,
+		Status: createJsonResponse.Status,
 	}, nil
 }
 
@@ -756,7 +757,7 @@ func (p *Provider) ListVolumeSnapshots(
 		"--project", p.GetProject(),
 		"snapshots",
 		"list",
-		"--format", "json(name,id)",
+		"--format", "json(name,id,status)",
 	}
 	var filters []string
 	// Only list snapshots that are fully created. Without this filter,
@@ -787,8 +788,9 @@ func (p *Provider) ListVolumeSnapshots(
 			continue
 		}
 		snapshots = append(snapshots, vm.VolumeSnapshot{
-			ID:   snapshotJson.ID,
-			Name: snapshotJson.Name,
+			ID:     snapshotJson.ID,
+			Name:   snapshotJson.Name,
+			Status: snapshotJson.Status,
 		})
 	}
 	sort.Sort(vm.VolumeSnapshots(snapshots))

pkg/roachprod/vm/vm.go

@@ -456,6 +456,17 @@ type VolumeSnapshot struct {
 	// Region is set and used by the AWS provider to scope snapshot
 	// operations to the correct region. Other providers may leave it empty.
 	Region string
+	// Status is the current state of the snapshot. The value is
+	// provider-specific (e.g. "completed" on AWS, "READY" on GCE).
+	// An empty string means the state is unknown.
+	Status string
+}
+
+// IsReady returns true when the snapshot has finished being created and
+// is usable. See VolumeSnapshot.Status for provider-specific values.
+func (v VolumeSnapshot) IsReady() bool {
+	return v.Status == "READY" || // GCE
+		v.Status == "completed" // AWS
 }
 
 type VolumeSnapshots []VolumeSnapshot
@@ -616,8 +627,10 @@ type Provider interface {
 	DeleteVolume(l *logger.Logger, volume Volume, vm *VM) error
 	// AttachVolume attaches the given volume to the given VM.
 	AttachVolume(l *logger.Logger, volume Volume, vm *VM) (string, error)
-	// CreateVolumeSnapshot creates a snapshot of the given volume, using the
-	// given options.
+	// CreateVolumeSnapshot creates a snapshot of the given volume. Some
+	// providers may return before the snapshot is fully ready. Callers
+	// that need a completed snapshot should poll ListVolumeSnapshots
+	// and check the Status field.
 	CreateVolumeSnapshot(l *logger.Logger, volume Volume, vsco VolumeSnapshotCreateOpts) (VolumeSnapshot, error)
 	// ListVolumeSnapshots lists the individual volume snapshots that satisfy
 	// the search criteria.