From 39ee61a9e59467c470a8953b7bcb7674c1d034d3 Mon Sep 17 00:00:00 2001
From: Mateus Melchiades <matbme@duck.com>
Date: Tue, 26 Dec 2023 17:18:12 -0300
Subject: [PATCH] fix: Wait until partitions are recognized by the kernel

---
 core/disk.go | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/core/disk.go b/core/disk.go
index 2142701f..e9de1b47 100644
--- a/core/disk.go
+++ b/core/disk.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"strconv"
+	"strings"
 )
 
 const (
@@ -107,6 +108,26 @@ func (disk *Disk) Update() error {
 	return nil
 }
 
+// waitForNewPartition is called after creating a new partition in order to
+// ensure the system is aware of it before proceeding.
+func (disk *Disk) waitForNewPartition() error {
+	for {
+		output, err := OutputCommand("lsblk -nro NAME /dev/nvme0n1 | wc -l")
+		if err != nil {
+			return err
+		}
+
+		count, err := strconv.Atoi(strings.TrimSpace(output))
+		if err != nil {
+			return err
+		}
+
+		if count-1 != len(disk.Partitions) {
+			return nil
+		}
+	}
+}
+
 func (disk *Disk) LabelDisk(label DiskLabel) error {
 	labelDiskCmd := "parted -s %s mklabel %s"
 
@@ -157,6 +178,12 @@ func (target *Disk) NewPartition(name string, fsType PartitionFs, start, end int
 		return nil, fmt.Errorf("failed to create partition: %s", err)
 	}
 
+	// Wait until kernel is aware of new partition
+	err = target.waitForNewPartition()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create partition: %s", err)
+	}
+
 	// Update partition list because we made changes to the disk
 	err = target.Update()
 	if err != nil {