From dd7d533c6c279d7ccc4b9f9664dd5c9fa858437d Mon Sep 17 00:00:00 2001
From: Chris Gunn <chrisgun@microsoft.com>
Date: Wed, 30 Oct 2024 16:11:40 -0700
Subject: [PATCH] Image Customizer: Fix partition creation on Ubuntu build
 hosts.

The change #10804 (Toolkit: Add missing flock calls) seems to have
caused a problem where on Ubuntu 22.04 build hosts, the partition info
isn't populated quickly enough after partition creation and formatting.
So, the `lsblk` call might return missing information. Adding a
`partprobe` call before the `lsblk` seems to fix the problem.
---
 toolkit/tools/pkg/imagecustomizerlib/imageutils.go    |  6 ++++++
 .../tools/pkg/imagecustomizerlib/partitionutils.go    | 11 +++++++++++
 .../tools/pkg/imagecustomizerlib/shrinkfilesystems.go |  4 ++--
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/toolkit/tools/pkg/imagecustomizerlib/imageutils.go b/toolkit/tools/pkg/imagecustomizerlib/imageutils.go
index b8927491800..7ff6e306941 100644
--- a/toolkit/tools/pkg/imagecustomizerlib/imageutils.go
+++ b/toolkit/tools/pkg/imagecustomizerlib/imageutils.go
@@ -183,6 +183,12 @@ func createImageBoilerplate(imageConnection *ImageConnection, filename string, b
 		return nil, "", fmt.Errorf("failed to create partitions on disk (%s):\n%w", imageConnection.Loopback().DevicePath(), err)
 	}
 
+	// Refresh partition entries under /dev.
+	err = refreshPartitions(imageConnection.Loopback().DevicePath())
+	if err != nil {
+		return nil, "", err
+	}
+
 	// Read the disk partitions.
 	diskPartitions, err := diskutils.GetDiskPartitions(imageConnection.Loopback().DevicePath())
 	if err != nil {
diff --git a/toolkit/tools/pkg/imagecustomizerlib/partitionutils.go b/toolkit/tools/pkg/imagecustomizerlib/partitionutils.go
index 98c4695a4f9..a4a27886de3 100644
--- a/toolkit/tools/pkg/imagecustomizerlib/partitionutils.go
+++ b/toolkit/tools/pkg/imagecustomizerlib/partitionutils.go
@@ -18,6 +18,7 @@ import (
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/logger"
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/safechroot"
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/safemount"
+	"github.com/microsoft/azurelinux/toolkit/tools/internal/shell"
 	"github.com/microsoft/azurelinux/toolkit/tools/internal/sliceutils"
 )
 
@@ -460,3 +461,13 @@ func getPartitionNum(partitionLoopDevice string) (int, error) {
 
 	return num, nil
 }
+
+func refreshPartitions(diskDevPath string) error {
+	err := shell.ExecuteLiveWithErr(1 /*stderrLines*/, "flock", "--timeout", "5", diskDevPath,
+		"partprobe", "-s", diskDevPath)
+	if err != nil {
+		return fmt.Errorf("partprobe failed:\n%w", err)
+	}
+
+	return nil
+}
diff --git a/toolkit/tools/pkg/imagecustomizerlib/shrinkfilesystems.go b/toolkit/tools/pkg/imagecustomizerlib/shrinkfilesystems.go
index 798616a87a2..0f7eadb2d43 100644
--- a/toolkit/tools/pkg/imagecustomizerlib/shrinkfilesystems.go
+++ b/toolkit/tools/pkg/imagecustomizerlib/shrinkfilesystems.go
@@ -108,9 +108,9 @@ func shrinkFilesystems(imageLoopDevice string, verity []imagecustomizerapi.Verit
 		}
 
 		// Re-read the partition table
-		err = shell.ExecuteLive(true, "flock", "--timeout", "5", imageLoopDevice, "partprobe", "-s", imageLoopDevice)
+		err = refreshPartitions(imageLoopDevice)
 		if err != nil {
-			return fmt.Errorf("partprobe failed:\n%w", err)
+			return err
 		}
 	}
 	return nil