handle timeout on driver stop. (#4055)

microsoft · Dec 2, 2024 · da17fd0 · da17fd0
1 parent b1176c7
commit da17fd0
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 50 deletions.
diff --git a/scripts/execute_ebpf_cicd_tests.ps1 b/scripts/execute_ebpf_cicd_tests.ps1
@@ -59,12 +59,22 @@ $Job = Start-Job -ScriptBlock {
     # currently one VM runs per runner.
     $TestVMName = $VMList[0].Name
 
-    # Run Kernel tests on test VM.
-    Write-Log "Running kernel tests on $TestVMName"
-    Run-KernelTestsOnVM -VMName $TestVMName -Config $Config
+    try {
+        # Run Kernel tests on test VM.
+        Write-Log "Running kernel tests on $TestVMName"
+        Run-KernelTestsOnVM -VMName $TestVMName -Config $Config
+
+        # Stop eBPF components on test VMs.
+        Stop-eBPFComponentsOnVM -VMName $TestVMName
+    } catch [System.Management.Automation.RemoteException] {
+        # Next, generate kernel dump.
+        Write-Log $_.Exception.Message
+        Write-Log $_.ScriptStackTrace
+        if ($_.CategoryInfo.Reason -eq "TimeoutException") {
+            Generate-KernelDumpOnVM($TestVMName)
+        }
+    }
 
-    # Stop eBPF components on test VMs.
-    Stop-eBPFComponentsOnVM -VMName $TestVMName
 
     Pop-Location
 } -ArgumentList (

diff --git a/scripts/install_ebpf.psm1 b/scripts/install_ebpf.psm1
@@ -122,6 +122,37 @@ function Start-WPRTrace
     }
 }
 
+function Stop-DriverWithTimeout {
+    param([parameter(Mandatory=$true)][string] $DriverName,
+          [parameter(Mandatory=$false)][int] $Timeout = 60)
+
+    Write-Log "Stopping driver $DriverName ..."
+    $Service = Get-Service $DriverName
+    if ($Service.Status -eq "Running") {
+        # Start the job to stop the driver.
+        $Job = Start-Job -ScriptBlock {
+            param($DriverName)
+            Stop-Service -Name $DriverName -Force
+        } -ArgumentList $DriverName
+
+        # Wait for the job to complete with a timeout
+        Write-Log "Waiting for $DriverName to stop in $Timeout seconds..."
+        if (Wait-Job -Job $Job -Timeout $Timeout) {
+            Write-Output "$DriverName stopped successfully within $Timeout seconds."
+        } else {
+            # If timeout occurs, stop the job and handle the timeout scenario
+            Stop-Job -Job $Job
+            Remove-Job -Job $Job
+            throw [System.TimeoutException]::new("Failed to stop $DriverName driver in $Timeout seconds.")
+        }
+        # Cleanup the job
+        Remove-Job -Job $Job -Force
+        Write-Log "$DriverName driver stopped." -ForegroundColor Green
+    } else {
+        Write-Log "$DriverName driver is not running." -ForegroundColor Green
+    }
+}
+
 # This function specifically tests that all eBPF drivers and services can be stopped.
 function Stop-eBPFComponents {
     # First, stop user mode service, so that EbpfCore does not hang on stop.
@@ -137,14 +168,8 @@ function Stop-eBPFComponents {
     }
     # Stop the drivers and services.
     $EbpfDrivers.GetEnumerator() | ForEach-Object {
-        try {
-            if ($_.Value.IsDriver) {
-                Write-Log "Stopping $($_.Key) driver..."
-                Stop-Service $_.Name -ErrorAction Stop 2>&1 | Write-Log
-                Write-Log "$($_.Key) driver stopped." -ForegroundColor Green
-            }
-        } catch {
-            throw "Failed to stop $($_.Key) driver: $_."
+        if ($_.Value.IsDriver) {
+            Stop-DriverWithTimeout -DriverName $_.Key
         }
     }
 }

diff --git a/scripts/vm_run_tests.psm1 b/scripts/vm_run_tests.psm1
@@ -686,43 +686,33 @@ function Run-KernelTestsOnVM
     param([Parameter(Mandatory = $true)] [string] $VMName,
           [Parameter(Mandatory = $true)] [PSCustomObject] $Config)
 
-    try {
-
-        # Run CICD tests on test VM.
-        Invoke-CICDTestsOnVM `
-            -VMName $VMName `
-            -TestMode $TestMode `
-            -Options $Options
-
-        # The required behavior is selected by the $TestMode
-        # parameter.
-        if (($TestMode -eq "CI/CD") -or ($TestMode -eq "Regression")) {
-
-            # Run XDP Tests.
-            Invoke-XDPTestsOnVM `
-                -Interfaces $Config.Interfaces `
-                -VMName $VMName
-
-            # Run Connect Redirect Tests.
-            Invoke-ConnectRedirectTestsOnVM `
-                -Interfaces $Config.Interfaces `
-                -ConnectRedirectTestConfig $Config.ConnectRedirectTest `
-                -UserType "Administrator" `
-                -VMName $VMName
-
-            Invoke-ConnectRedirectTestsOnVM `
-                -Interfaces $Config.Interfaces `
-                -ConnectRedirectTestConfig $Config.ConnectRedirectTest `
-                -UserType "StandardUser" `
-                -VMName $VMName
-        }
-    } catch [System.Management.Automation.RemoteException] {
-        # Next, generate kernel dump.
-        Write-Log $_.Exception.Message
-        Write-Log $_.ScriptStackTrace
-        if ($_.CategoryInfo.Reason -eq "TimeoutException") {
-            Generate-KernelDumpOnVM($VMName)
-        }
+    # Run CICD tests on test VM.
+    Invoke-CICDTestsOnVM `
+        -VMName $VMName `
+        -TestMode $TestMode `
+        -Options $Options
+
+    # The required behavior is selected by the $TestMode
+    # parameter.
+    if (($TestMode -eq "CI/CD") -or ($TestMode -eq "Regression")) {
+
+        # Run XDP Tests.
+        Invoke-XDPTestsOnVM `
+            -Interfaces $Config.Interfaces `
+            -VMName $VMName
+
+        # Run Connect Redirect Tests.
+        Invoke-ConnectRedirectTestsOnVM `
+            -Interfaces $Config.Interfaces `
+            -ConnectRedirectTestConfig $Config.ConnectRedirectTest `
+            -UserType "Administrator" `
+            -VMName $VMName
+
+        Invoke-ConnectRedirectTestsOnVM `
+            -Interfaces $Config.Interfaces `
+            -ConnectRedirectTestConfig $Config.ConnectRedirectTest `
+            -UserType "StandardUser" `
+            -VMName $VMName
     }
 }