feat: add system startup task to expand os volume

staging/cse/windows/configfunc.ps1

@@ -23,11 +23,14 @@ function Resize-OSDrive
         $osDisk = Get-Partition -DriveLetter $osDrive | Get-Disk
         if ($osDisk.Size - $osDisk.AllocatedSize -gt 1GB)
         {
+            Write-Log "Expanding the OS volume"
             # Create a diskpart script (text file) that will select the OS volume, extend it and exit.
             $diskpartScriptPath = [String]::Format("{0}\\diskpart_extendOSVol.script", $env:temp)
             [String]::Format("select volume {0}`nextend`nexit", $osDrive) | Out-File -Encoding "UTF8" $diskpartScriptPath
             Invoke-Executable -Executable "diskpart.exe" -ArgList @("/s", $diskpartScriptPath) -ExitCode $global:WINDOWS_CSE_ERROR_RESIZE_OS_DRIVE
             Remove-Item -Path $diskpartScriptPath -Force
+        } else {
+            Write-Log "No need to expand the OS volume due to ScheduledTask executed before CSE."
         }
     } catch {
         Set-ExitCode -ExitCode $global:WINDOWS_CSE_ERROR_RESIZE_OS_DRIVE -ErrorMessage "Failed to resize os drive. Error: $_"

vhdbuilder/packer/configure-windows-vhd.ps1

@@ -242,13 +242,40 @@ function Get-ToolsToVHD {
     # Rely on the completion of Get-FilesToCacheOnVHD
     $cacheDir = "c:\akse-cache\tools"
 
-    $toolsDir = "c:\aks-tools"
-    if (!(Test-Path -Path $toolsDir)) {
-        New-Item -ItemType Directory -Path $toolsDir | Out-Null
+    if (!(Test-Path -Path $global:aksToolsDir)) {
+        New-Item -ItemType Directory -Path $global:aksToolsDir -Force | Out-Null
     }
 
     Write-Log "Getting DU (Windows Disk Usage)"
-    Expand-Archive -Path "$cacheDir\DU.zip" -DestinationPath "$toolsDir\DU" -Force
+    Expand-Archive -Path "$cacheDir\DU.zip" -DestinationPath "$global:aksToolsDir\DU" -Force
+}
+
+function Register-ExpandVolumeTask {
+    if (!(Test-Path -Path $global:aksToolsDir)) {
+        New-Item -ItemType Directory -Path $global:aksToolsDir -Force | Out-Null
+    }
+
+    # Leverage existing folder 'c:\aks-tools' to store the task scripts
+    $taskScript = @'
+        $osDrive = ((Get-WmiObject Win32_OperatingSystem -ErrorAction Stop).SystemDrive).TrimEnd(":")
+        $diskpartScriptPath = "c:\aks-tools\diskpart.script"
+        [String]::Format("select volume {0}`nextend`nexit", $osDrive) | Out-File -Encoding "UTF8" $diskpartScriptPath -Force
+        Start-Process -FilePath diskpart.exe -ArgumentList "/s $diskpartScriptPath" -Wait
+        # Run once and remove the task. Sequence: taks invokes ps1, ps1 invokes diskpart.
+        Unregister-ScheduledTask -TaskName "aks-expand-volume" -Confirm:$false
+        Remove-Item -Path "c:\aks-tools\expand-volume.ps1" -Force
+        Remove-Item -Path $diskpartScriptPath -Force
+'@
+
+    $taskScriptPath = Join-Path $global:aksToolsDir "expand-volume.ps1"
+    $taskScript| Set-Content -Path $taskScriptPath -Force
+
+    $action = New-ScheduledTaskAction -Execute "powershell.exe" -Argument "-File `"$taskScriptPath`""
+    $principal = New-ScheduledTaskPrincipal -UserId SYSTEM -LogonType ServiceAccount -RunLevel Highest
+    $trigger = New-JobTrigger -AtStartup
+    $definition = New-ScheduledTask -Action $action -Principal $principal -Trigger $trigger -Description "aks-expand-volume"
+    Register-ScheduledTask -TaskName "aks-expand-volume" -InputObject $definition
+    Write-Log "Registered ScheduledTask aks-expand-volume"
 }
 
 function Get-PrivatePackagesToCacheOnVHD {
@@ -900,10 +927,13 @@ try{
             Get-FilesToCacheOnVHD
             Get-ToolsToVHD # Rely on the completion of Get-FilesToCacheOnVHD
             Get-PrivatePackagesToCacheOnVHD
+            Log-ReofferUpdate
+        }
+        "3" {
+            Register-ExpandVolumeTask
             Remove-Item -Path c:\windows-vhd-configuration.ps1
             Cleanup-TemporaryFiles
             (New-Guid).Guid | Out-File -FilePath 'c:\vhd-id.txt'
-            Log-ReofferUpdate
         }
         default {
             Write-Log "Unable to determine provisiong phase... exiting"

vhdbuilder/packer/generate-windows-vhd-configuration.ps1

@@ -11,6 +11,9 @@ if (-not ($validSKU -contains $windowsSKU)) {
 # We use the same temp dir for all temp tools that will be used for vhd build
 $global:aksTempDir = "c:\akstemp"
 
+# We use the same dir for all tools that will be used in AKS Windows nodes
+$global:aksToolsDir = "c:\aks-tools"
+
 # We need to guarantee that the node provisioning will not fail because the vhd is full before resize-osdisk is called in AKS Windows CSE script.
 $global:lowestFreeSpace = 2*1024*1024*1024 # 2GB
 

vhdbuilder/packer/test/windows-vhd-content-test.ps1

@@ -544,6 +544,17 @@ function Test-ToolsToCacheOnVHD {
     }
 }
 
+function Test-ExpandVolumeTask {
+    $osDrive = ((Get-WmiObject Win32_OperatingSystem -ErrorAction Stop).SystemDrive).TrimEnd(":")
+    $osDisk = Get-Partition -DriveLetter $osDrive | Get-Disk
+    $osDiskSize = $osDisk.Size 
+    $osDiskAllocatedSize = $osDisk.AllocatedSize
+    if ($osDiskSize -ne $osDiskAllocatedSize) {
+        Write-ErrorWithTimestamp "The OS disk size $osDiskSize is not equal to the allocated size $osDiskAllocatedSize"
+        exit 1
+    }
+}
+
 function Test-SSHDConfig {
     # user must be the name in `TEST_VM_ADMIN_USERNAME="azureuser"` in vhdbuilder/packer/test/run-test.sh
     $result=$(sshd -T -C user=azureuser)
@@ -566,4 +577,5 @@ Test-AzureExtensions
 Test-ExcludeUDPSourcePort
 Test-WindowsDefenderPlatformUpdate
 Test-ToolsToCacheOnVHD
+Test-ExpandVolumeTask
 Remove-Item -Path c:\windows-vhd-configuration.ps1

vhdbuilder/packer/windows-vhd-builder-sig.json

@@ -107,6 +107,18 @@
             "restart_timeout": "10m",
             "type": "windows-restart"
         },
+        {
+            "elevated_user": "packer",
+            "elevated_password": "{{.WinRMPassword}}",
+            "environment_vars": [
+                "ProvisioningPhase=3",
+                "WindowsSKU={{user `windows_sku`}}"
+            ],
+            "type": "powershell",
+            "scripts": [
+                "vhdbuilder/packer/configure-windows-vhd.ps1"
+            ]
+        },
         {
             "type": "file",
             "direction": "upload",