[macOS] Improve select datastore script (#1892)

* modify script to set tages

* add vmname

* nitpicks

* add remove tag

images.CI/macos/azure-pipelines/image-generation.yml

@@ -52,7 +52,8 @@ jobs:
     inputs:
       targetType: 'filePath'
       filePath: ./images.CI/macos/select-datastore.ps1
-      arguments: -VIServer "$(vcenter-server-v2)" `
+      arguments: -VMName "$(VirtualMachineName)" `
+                 -VIServer "$(vcenter-server-v2)" `
                  -VIUserName "$(vcenter-username-v2)" `
                  -VIPassword "$(vcenter-password-v2)"
 
@@ -121,8 +122,8 @@ jobs:
     condition: always()
 
   - task: PowerShell@2
-    displayName: 'Move vm to cold storage'
+    displayName: 'Move vm to cold storage and clear datastore tag'
-    condition: succeededOrFailed()
+    condition: always()
     inputs:
       targetType: 'filePath'
       filePath: ./images.CI/macos/move-vm.ps1

images.CI/macos/move-vm.ps1

@@ -48,6 +48,13 @@ Import-Module $PSScriptRoot\helpers.psm1 -DisableNameChecking
 # Connection to a vCenter Server system
 Connect-VCServer
 
+# Clear previously assigned tag with VM Name
+try {
+    Remove-Tag $VMName -Confirm:$false
+} catch {
+    Write-Host "Tag with $VMName doesn't exist"
+}
+
 $vm = Get-VM $VMName
 
 if ($env:AGENT_JOBSTATUS -eq 'Failed') {

images.CI/macos/select-datastore.ps1

@@ -20,6 +20,10 @@ vCenter password (Example "12345678")
 
 [CmdletBinding()]
 param(
+    [Parameter(Mandatory)]
+    [ValidateNotNullOrEmpty()]
+    [string]$VMName,
+
     [Parameter(Mandatory)]
     [ValidateNotNullOrEmpty()]
     [string]$VIServer,
@@ -30,35 +34,72 @@ param(
 
     [Parameter(Mandatory)]
     [ValidateNotNullOrEmpty()]
-    [string]$VIPassword
+    [string]$VIPassword,
+
+    [string]$TagCategory = "Busy"
 )
 
 # Import helpers module
 Import-Module $PSScriptRoot\helpers.psm1 -DisableNameChecking
 
+function Select-DataStore {
+    param (
+        [string]$VMName,
+        [string]$TagCategory,
+        [string]$TemplateDatastore = "ds-local-Datastore-*",
+        [int]$ThresholdInGb = 400,
+        [int]$VMCount = 2,
+        [int]$Retries = 5
+    )
+
+    # 1. Name starts with ds-local-Datastore
+    # 2. FreespaceGB > 400 Gb
+    # 3. VM count on a datastore < 2
+
+    Write-Host "Start Datastore selection process..."
+    $allDatastores = Get-Datastore -Name $templateDatastore | Where-Object { $_.State -eq "Available" }
+    $buildDatastore = $allDatastores | Where-Object { $_.FreeSpaceGB -ge $thresholdInGb } | Where-Object {
+        $vmOnDatastore = @((Get-ChildItem -Path $_.DatastoreBrowserPath).Name -notmatch "^\.").Count
+        $vmOnDatastore -lt $vmCount
+    } | Select-Object -ExpandProperty Name -First 1
+
+    $tag = Get-Tag -Category $TagCategory -Name $VMName -ErrorAction Ignore
+    if (-not $tag)
+    {
+        $tag = New-Tag -Name $VMName -Category $TagCategory
+    }
+
+    New-TagAssignment -Tag $tag -Entity $buildDatastore | Out-Null
+
+    # Wait for 60 seconds to check if any other tags are assigned to the same datastore
+    Start-Sleep -Seconds 60
+    # Take only first 2 tags, all the others will go to the next round
+    $tagAssignments = (Get-TagAssignment -Entity $buildDatastore).Tag.Name | Select-Object -First 2
+    $isAllow = $tagAssignments -contains $VMName
+
+    if ($isAllow)
+    {
+        Write-Host "Datastore selected successfully"
+        Write-Host "##vso[task.setvariable variable=buildDatastore;issecret=true]$buildDatastore"
+        return
+    }
+
+    # Remove the tag if datastore wasn't selected
+    Remove-Tag $tag -Confirm:$false
+
+    $retries--
+    if ($retries -le 0)
+    {
+        Write-Host "##vso[task.LogIssue type=error;]No datastores found for the condition"
+        exit 1
+    }
+
+    Write-Host "Datastore select failed, $retries left"
+    Select-DataStore -VMName $VMName -TagCategory $TagCategory -Retries $retries
+}
+
 # Connection to a vCenter Server system
 Connect-VCServer
 
 # Get a target datastore for current deployment
-# 1. Name starts with ds-local-Datastore
+Select-DataStore -VMName $VMName -TagCategory $TagCategory
-# 2. FreespaceGB > 400 Gb
-# 3. VM count on a datastore < 2
-$templateDatastore = "ds-local-Datastore-*"
-$thresholdInGb = 400
-$vmCount = 2
-$allDatastores = Get-Datastore -Name $templateDatastore | Where-Object { $_.State -eq "Available" }
-$buildDatastore = $allDatastores | Where-Object { $_.FreeSpaceGB -ge $thresholdInGb } | Where-Object {
-            $vmOnDatastore = @((Get-ChildItem -Path $_.DatastoreBrowserPath).Name -notmatch "^\.").Count
-            $vmOnDatastore -lt $vmCount
-        } | Select-Object -ExpandProperty Name -First 1
-
-if ($buildDatastore)
-{
-    Write-Host "Datastore selected successfully"
-    Write-Host "##vso[task.setvariable variable=buildDatastore;issecret=true]$buildDatastore"
-}
-else
-{
-    Write-Host "##vso[task.LogIssue type=error;]No datastores found for the condition"
-    exit 1
-}