From cf641c4db4239434303456f11c54fc192e26d00e Mon Sep 17 00:00:00 2001
From: Nick Banks <nibanks@microsoft.com>
Date: Thu, 11 Jan 2024 10:39:24 -0500
Subject: [PATCH] More fixes and refactoring

---
 scripts/secnetperf-helpers.psm1 | 25 +++++++++++++++++-------
 scripts/secnetperf.ps1          | 34 ++++++---------------------------
 2 files changed, 24 insertions(+), 35 deletions(-)

diff --git a/scripts/secnetperf-helpers.psm1 b/scripts/secnetperf-helpers.psm1
index 7d76387886..8839cbe9ae 100644
--- a/scripts/secnetperf-helpers.psm1
+++ b/scripts/secnetperf-helpers.psm1
@@ -10,26 +10,37 @@ function Write-GHError($msg) {
 
 # Waits for a remote job to be ready based on looking for a particular string in
 # the output.
-function Wait-ForRemoteReady {
-    param ($Job, $Matcher)
+function Start-RemoteServer {
+    param ($Command)
+    $Job = Invoke-Command -Session $Session -ScriptBlock { $Using:Command } -AsJob
     $StopWatch =  [system.diagnostics.stopwatch]::StartNew()
+    $Started = $false
     while ($StopWatch.ElapsedMilliseconds -lt 10000) {
         $CurrentResults = Receive-Job -Job $Job -Keep -ErrorAction Continue
         if (![string]::IsNullOrWhiteSpace($CurrentResults)) {
-            $DidMatch = $CurrentResults -match $Matcher
+            $DidMatch = $CurrentResults -match "Started!"
             if ($DidMatch) {
-                return $true
+                $Started = $true
+                break;
             }
         }
         Start-Sleep -Seconds 0.1 | Out-Null
     }
-    return $false
+    if (!$Started) {
+        Stop-Job -Job $Job
+        $RemoteResult = Receive-Job -Job $Job -ErrorAction Stop
+        $RemoteResult = $RemoteResult -join "`n"
+        Write-GHError "Server failed to start! Output:"
+        Write-Output $RemoteResult
+        return $null
+    }
+    return $Job
 }
 
 # Sends a special UDP packet to tell the remote secnetperf to shutdown, and then
 # waits for the job to complete. Finally, it returns the console output of the
 # job.
-function Wait-ForRemote {
+function Stop-RemoteServer {
     param ($Job, $ErrorAction = "Stop")
     # Ping side-channel socket on 9999 to tell the app to die
     $Socket = New-Object System.Net.Sockets.UDPClient
@@ -47,7 +58,7 @@ function Wait-ForRemote {
 
     Stop-Job -Job $Job | Out-Null
     $RetVal = Receive-Job -Job $Job -ErrorAction $ErrorAction
-    return $RetVal -join "`n"
+    $RetVal = $RetVal -join "`n"
 }
 
 # Invokes all the secnetperf tests.
diff --git a/scripts/secnetperf.ps1 b/scripts/secnetperf.ps1
index 83bc494d9f..6303f7add3 100644
--- a/scripts/secnetperf.ps1
+++ b/scripts/secnetperf.ps1
@@ -138,19 +138,8 @@ if (!$isWindows) {
 
 # Run secnetperf on the server.
 Write-Output "Starting secnetperf server..."
-$Job = Invoke-Command -Session $Session -ScriptBlock {
-    & "$Using:RemoteDir/$Using:SecNetPerfPath -exec:maxtput"
-} -AsJob
-
-# Wait for the server to start.
-Write-Output "Waiting for server to start..."
-$ReadyToStart = Wait-ForRemoteReady -Job $Job -Matcher "Started!"
-if (!$ReadyToStart) {
-    Stop-Job -Job $Job
-    $RemoteResult = Receive-Job -Job $Job -ErrorAction Stop
-    $RemoteResult = $RemoteResult -join "`n"
-    Write-GHError "Server failed to start! Output:"
-    Write-Output $RemoteResult
+$Job = Start-RemoteServer "$RemoteDir/$SecNetPerfPath -exec:maxtput"
+if ($null -eq $Job) {
     throw "Server failed to start!"
 }
 
@@ -213,23 +202,12 @@ $SQL += Invoke-SecnetperfTest $maxtputIds $maxtput $exe $json $LogProfile
 # Start and restart the SecNetPerf server without maxtput.
 Write-Host "Restarting server without maxtput..."
 Write-Host "`nStopping server. Server Output:"
-$RemoteResults = Wait-ForRemote $Job
+$RemoteResults = Stop-RemoteServer $Job
 Write-Host $RemoteResults.ToString()
 
 Write-Host "Starting server back up again..."
-$Job = Invoke-Command -Session $Session -ScriptBlock {
-    & "$Using:RemoteDir/$Using:SecNetPerfPath -exec:lowlat"
-} -AsJob
-
-# Wait for the server to start.
-Write-Output "Waiting for server to start..."
-$ReadyToStart = Wait-ForRemoteReady -Job $Job -Matcher "Started!"
-if (!$ReadyToStart) {
-    Stop-Job -Job $RemoteJob
-    $RemoteResult = Receive-Job -Job $Job -ErrorAction $ErrorAction
-    $RemoteResult = $RemoteResult -join "`n"
-    Write-GHError "Server failed to start! Output:"
-    Write-Output $RemoteResult
+$Job = Start-RemoteServer "$RemoteDir/$SecNetPerfPath -exec:lowlat"
+if ($null -eq $Job) {
     throw "Server failed to start!"
 }
 
@@ -243,7 +221,7 @@ $SQL += Invoke-SecnetperfTest $lowlatIds $lowlat $exe $json $LogProfile
 
 # Kill the server process.
 Write-Output "`nStopping server. Server Output:"
-$RemoteResults = Wait-ForRemote $Job
+$RemoteResults = Stop-RemoteServer $Job
 Write-Output $RemoteResults.ToString()
 
 # if ($LogProfile -ne "" -and $LogProfile -ne "NULL") { # TODO: Linux back slash works?