Image Customizer: Fix partition creation on Ubuntu build hosts. (#10902)

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` call seems to fix the problem.

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 {

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
+}

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