fix: Wait for partition to become available

* cleanup: Add documentation for GetPartition

* fix: Wait until partition is available

* Add verbose flag

* Replace testing image with dev

* Remove verbose flag

* Add timeout to WaitUntilAvailable

core/disk.go

@@ -259,13 +259,19 @@ func (target *Disk) NewPartition(name string, fsType PartitionFs, start, end int
 	return newPartition, nil
 }
 
+// GetPartition attempts to locate a partition by its number. For instance, partition 3
+// will normally point to `/dev/sda3`, but this might not be the case if partitions have
+// been deleted (see [Issue #44]). This function first checks if the desired partition is
+// in the correct place, else it searches all partitions in target for the correct one.
+//
+// [Issue #44]: https://github.com/Vanilla-OS/Albius/issues/44
 func (target *Disk) GetPartition(partNum int) *Partition {
-    // Happy path
+    // Happy path: No partitions are missing
     if target.Partitions[partNum-1].Number == partNum {
         return &target.Partitions[partNum-1]
     }
 
-    // There are missing partition numbers, find partition manually
+    // Missing partition numbers, find correct partition manually
     for _, part := range target.Partitions {
         if part.Number == partNum {
             return &part

core/partition.go

@@ -2,8 +2,10 @@ package albius
 
 import (
 	"fmt"
+	"os"
 	"slices"
 	"strings"
+	"time"
 )
 
 const (
@@ -291,3 +293,27 @@ func (part *Partition) SetLabel(label string) error {
 
 	return nil
 }
+
+// WaitUntilAvailable polls the specified partition until it is available.
+// This is particularly useful to make sure a recently created or modified
+// partition is recognized by the system.
+func (part *Partition) WaitUntilAvailable() {
+	maxTimeout := 1000
+	timeout := 0
+
+	for {
+		_, err := os.Stat(part.Path)
+		if !os.IsNotExist(err) {
+			if uuid, err := part.GetUUID(); err != nil && uuid != "" {
+				return
+			}
+		}
+		time.Sleep(50 * time.Millisecond)
+
+		timeout += 1
+		if timeout == maxTimeout {
+			// We can't recover from this, so just panic
+			panic(fmt.Sprintf("Timed out waiting for partition %s", part.Path))
+		}
+	}
+}

core/recipe.go

@@ -123,9 +123,10 @@ func runSetupOperation(diskLabel, operation string, args []interface{}) error {
 			}
 			// lsblk seems to take a few milliseconds to update the partition's
 			// UUID, so we loop until it gives us one
-			uuid := ""
-			for uuid == "" {
-				uuid, _ = part.GetUUID()
+			part.WaitUntilAvailable()
+			uuid, err := part.GetUUID()
+			if err != nil {
+				return err
 			}
 			err = LuksOpen(part, fmt.Sprintf("luks-%s", uuid), luksPassword)
 			if err != nil {
@@ -282,9 +283,10 @@ func runSetupOperation(diskLabel, operation string, args []interface{}) error {
 		}
 		// lsblk seems to take a few milliseconds to update the partition's
 		// UUID, so we loop until it gives us one
-		uuid := ""
-		for uuid == "" {
-			uuid, _ = part.GetUUID()
+		part.WaitUntilAvailable()
+		uuid, err := part.GetUUID()
+		if err != nil {
+			return err
 		}
 		err = LuksOpen(part, fmt.Sprintf("luks-%s", uuid), password)
 		if err != nil {
@@ -310,6 +312,8 @@ func runSetupOperation(diskLabel, operation string, args []interface{}) error {
 	 */
 	case "pvcreate":
 		part := args[0].(string)
+		dummyPart := Partition{Path: part}
+		dummyPart.WaitUntilAvailable()
 		err := lvm.Pvcreate(part)
 		if err != nil {
 			return err
@@ -360,6 +364,8 @@ func runSetupOperation(diskLabel, operation string, args []interface{}) error {
 		pvs := []string{}
 		if len(args) > 1 {
 			for _, pv := range args[1].([]interface{}) {
+				dummyPart := Partition{Path: pv.(string)}
+				dummyPart.WaitUntilAvailable()
 				pvs = append(pvs, pv.(string))
 			}
 		}
@@ -585,9 +591,10 @@ func runSetupOperation(diskLabel, operation string, args []interface{}) error {
 		}
 		// lsblk seems to take a few milliseconds to update the partition's
 		// UUID, so we loop until it gives us one
-		uuid := ""
-		for uuid == "" {
-			uuid, _ = dummyPart.GetUUID()
+		dummyPart.WaitUntilAvailable()
+		uuid, err := dummyPart.GetUUID()
+		if err != nil {
+			return err
 		}
 		err = LuksOpen(&dummyPart, fmt.Sprintf("luks-%s", uuid), password)
 		if err != nil {

utils/create_test_env.sh

@@ -2,4 +2,4 @@
 
 set -e
 
-distrobox-create -r -I -Y -ap "golang libbtrfs-dev libdevmapper-dev libgpgme-dev build-essential pkg-config lvm2 parted udev" -i ghcr.io/vanilla-os/vso:main albius_test
+distrobox-create -r -I -Y -ap "golang libbtrfs-dev libdevmapper-dev libgpgme-dev build-essential pkg-config lvm2 parted udev" -i ghcr.io/vanilla-os/dev:main albius_test