[macOS] Improve select datastore script (#1892)

* modify script to set tages

* add vmname

* nitpicks

* add remove tag

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
-# 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
-}
+Select-DataStore -VMName $VMName -TagCategory $TagCategory