Time-constrain user commands in XHarness workloads

src/Microsoft.DotNet.Helix/Sdk/CreateXHarnessAndroidWorkItems.cs

@@ -105,7 +105,7 @@ private async Task<ITaskItem> PrepareWorkItem(IZipArchiveManager zipArchiveManag
                 apkName = apkName.Replace(".zip", ".apk");
             }
 
-            string command = GetHelixCommand(appPackage, apkName, androidPackageName, testTimeout, expectedExitCode);
+            string command = GetHelixCommand(appPackage, apkName, androidPackageName, workItemTimeout, testTimeout, expectedExitCode);
 
             string workItemZip = await CreatePayloadArchive(
                 zipArchiveManager,
@@ -150,7 +150,13 @@ private string GetDefaultCommand(ITaskItem appPackage, int expectedExitCode)
                 $"{ devOutArg } { instrumentationArg } { exitCodeArg } { extraArguments } { passthroughArgs }";
         }
 
-        private string GetHelixCommand(ITaskItem appPackage, string apkName, string androidPackageName, TimeSpan xHarnessTimeout, int expectedExitCode)
+        private string GetHelixCommand(
+            ITaskItem appPackage,
+            string apkName,
+            string androidPackageName,
+            TimeSpan workItemTimeout,
+            TimeSpan xHarnessTimeout,
+            int expectedExitCode)
         {
             appPackage.TryGetMetadata(MetadataNames.AndroidInstrumentationName, out string androidInstrumentationName);
             appPackage.TryGetMetadata(MetadataNames.DeviceOutputPath, out string deviceOutputPath);
@@ -163,6 +169,7 @@ private string GetHelixCommand(ITaskItem appPackage, string apkName, string andr
             string dash = IsPosixShell ? "--" : "-";
             string xharnessRunCommand = $"{xharnessHelixWrapperScript} " +
                 $"{dash}app \"{apkName}\" " +
+                $"{dash}command_timeout {(int)workItemTimeout.TotalSeconds} " +
                 $"{dash}timeout \"{xHarnessTimeout}\" " +
                 $"{dash}package_name \"{androidPackageName}\" " +
                 (expectedExitCode != 0 ? $" {dash}expected_exit_code \"{expectedExitCode}\" " : string.Empty) +

src/Microsoft.DotNet.Helix/Sdk/CreateXHarnessAppleWorkItems.cs

@@ -1,6 +1,7 @@
 using System;
 using System.IO;
 using System.Linq;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Arcade.Common;
 using Microsoft.Build.Framework;
@@ -168,11 +169,11 @@ private async Task<ITaskItem> PrepareWorkItem(
             if (customCommands == null)
             {
                 // When no user commands are specified, we add the default `apple test ...` command
-                customCommands = GetDefaultCommand(target, includesTestRunner, resetSimulator);
+                customCommands = GetDefaultCommand(includesTestRunner, resetSimulator);
             }
 
             string appName = isAlreadyArchived ? $"{fileSystem.GetFileNameWithoutExtension(appFolderPath)}.app" : fileSystem.GetFileName(appFolderPath);
-            string helixCommand = GetHelixCommand(appName, target, testTimeout, launchTimeout, includesTestRunner, expectedExitCode, resetSimulator);
+            string helixCommand = GetHelixCommand(appName, target, workItemTimeout, testTimeout, launchTimeout, includesTestRunner, expectedExitCode, resetSimulator);
             string payloadArchivePath = await CreatePayloadArchive(
                 zipArchiveManager,
                 fileSystem,
@@ -196,7 +197,7 @@ private bool ValidateAppBundlePath(
             return isAlreadyArchived ? fileSystem.FileExists(appBundlePath) : fileSystem.DirectoryExists(appBundlePath);
         }
 
-        private string GetDefaultCommand(string target, bool includesTestRunner, bool resetSimulator) =>
+        private string GetDefaultCommand(bool includesTestRunner, bool resetSimulator) =>
             $"xharness apple {(includesTestRunner ? "test" : "run")} " +
             "--app \"$app\" " +
             "--output-directory \"$output_directory\" " +
@@ -213,6 +214,7 @@ private string GetDefaultCommand(string target, bool includesTestRunner, bool re
         private string GetHelixCommand(
             string appName,
             string target,
+            TimeSpan workItemTimeout,
             TimeSpan testTimeout,
             TimeSpan launchTimeout,
             bool includesTestRunner,
@@ -222,6 +224,7 @@ private string GetHelixCommand(
             $"chmod +x {EntryPointScript} && ./{EntryPointScript} " +
             $"--app \"{appName}\" " +
             $"--target \"{target}\" " +
+            $"--command-timeout {(int)workItemTimeout.TotalSeconds} " +
             $"--timeout \"{testTimeout}\" " +
             $"--launch-timeout \"{launchTimeout}\" " +
             (includesTestRunner ? "--includes-test-runner " : string.Empty) +

src/Microsoft.DotNet.Helix/Sdk/XharnessTaskBase.cs

@@ -14,6 +14,7 @@ public abstract class XHarnessTaskBase : MSBuildTaskBase
     {
         private static readonly TimeSpan s_defaultWorkItemTimeout = TimeSpan.FromMinutes(20);
         private static readonly TimeSpan s_defaultTestTimeout = TimeSpan.FromMinutes(12);
+        protected static readonly TimeSpan s_telemetryBuffer = TimeSpan.FromSeconds(20); // extra time to send the XHarness telemetry
 
         public class MetadataName
         {
@@ -102,6 +103,9 @@ protected Build.Utilities.TaskItem CreateTaskItem(string workItemName, string pa
         {
             Log.LogMessage($"Creating work item with properties Identity: {workItemName}, Payload: {payloadArchivePath}, Command: {command}");
 
+            // Leave some time at the end of the work item to send the telemetry (in case it times out)
+            timeout += s_telemetryBuffer;
+
             return new(workItemName, new Dictionary<string, string>()
             {
                 { "Identity", workItemName },

src/Microsoft.DotNet.Helix/Sdk/tools/xharness-runner/xharness-helix-job.android.ps1

@@ -12,6 +12,8 @@ param (
     [string]$app,
     [Parameter(Mandatory)]
     [string]$timeout,
+    [Parameter(Mandatory)]
+    [int]$command_timeout, # in seconds
     [Parameter()]
     [string]$package_name = $null,
     [Parameter()]
@@ -24,9 +26,6 @@ param (
 
 $ErrorActionPreference="Stop"
 
-[Diagnostics.CodeAnalysis.SuppressMessageAttribute("PSUseDeclaredVarsMoreThanAssignments", "")] # Variable used in sourced script
-$output_directory=$Env:HELIX_WORKITEM_UPLOAD_ROOT
-
 # The xharness alias
 function xharness() {
     dotnet exec $Env:XHARNESS_CLI_PATH @args
@@ -35,11 +34,52 @@ function xharness() {
 $ErrorActionPreference="Continue"
 
 # Act out the actual commands
-. "$PSScriptRoot\command.ps1"
+# We have to time constrain them to create buffer for the end of this script
+$code = @{}
+$job = [PowerShell]::Create().AddScript({
+    param(
+        $current_dir,
+        $output_directory,
+        $app,
+        $timeout,
+        $command_timeout,
+        $package_name,
+        $expected_exit_code,
+        $device_output_path,
+        $instrumentation,
+        $result)
+    $result.ExitCode = 0
+    . "$current_dir\command.ps1"
+    $result.ExitCode = $LASTEXITCODE
+}).AddArgument($PSScriptRoot).AddArgument($Env:HELIX_WORKITEM_UPLOAD_ROOT).AddArgument($app).AddArgument($timeout).AddArgument($command_timeout).AddArgument($package_name).AddArgument($expected_exit_code).AddArgument($device_output_path).AddArgument($instrumentation).AddArgument($code)
 
-$ErrorActionPreference="Continue"
+$output = New-Object 'System.Management.Automation.PSDataCollection[psobject]'
+$task = $job.BeginInvoke($output, $output);
+$timer = [Diagnostics.Stopwatch]::StartNew();
+
+do
+{
+    if ($task.IsCompleted) {
+        $exit_code = $code.ExitCode;
+        break;
+    }
 
-$exit_code=$LASTEXITCODE
+    if ($timer.Elapsed.TotalSeconds -gt $command_timeout) {
+        Write-Error "User command timed out after $command_timeout seconds!";
+        $job.Stop();
+        $exit_code = -3;
+        break;
+    }
+
+    Start-Sleep -Milliseconds 250;
+} while (1 -eq 1)
+
+# This prints the output of the $job (the actual user commands)
+Write-Output $output;
+
+$job.Dispose();
+
+$ErrorActionPreference="Continue"
 
 $retry=$false
 $reboot=$false

src/Microsoft.DotNet.Helix/Sdk/tools/xharness-runner/xharness-helix-job.android.sh

@@ -12,6 +12,7 @@ echo "XHarness Helix Job Wrapper calling '$@'"
 set -x
 
 app=''
+command_timeout=20
 timeout=''
 package_name=''
 expected_exit_code=0
@@ -26,6 +27,10 @@ while [[ $# -gt 0 ]]; do
         app="$2"
         shift
         ;;
+      --command_timeout)
+        command_timeout="$2"
+        shift
+        ;;
       --timeout)
         timeout="$2"
         shift
@@ -67,8 +72,8 @@ function xharness() {
     dotnet exec $XHARNESS_CLI_PATH "$@"
 }
 
-# Act out the actual commands
-source command.sh
+# Act out the actual commands (and time constrain them to create buffer for the end of this script)
+source command.sh & PID=$! ; (sleep $command_timeout && kill $PID 2> /dev/null & ) ; wait $PID
 
 exit_code=$?
 

src/Microsoft.DotNet.Helix/Sdk/tools/xharness-runner/xharness-runner.apple.sh

@@ -10,6 +10,7 @@ app=''
 target=''
 timeout=''
 launch_timeout=''
+command_timeout=20
 xcode_version=''
 app_arguments=''
 expected_exit_code=0
@@ -31,6 +32,10 @@ while [[ $# -gt 0 ]]; do
         timeout="$2"
         shift
         ;;
+      --command-timeout)
+        command_timeout="$2"
+        shift
+        ;;
       --launch-timeout)
         launch_timeout="$2"
         shift
@@ -139,8 +144,8 @@ function xharness() {
     dotnet exec "$XHARNESS_CLI_PATH" "$@"
 }
 
-# Act out the actual commands
-source command.sh
+# Act out the actual commands (and time constrain them to create buffer for the end of this script)
+source command.sh & PID=$! ; (sleep $command_timeout && kill $PID 2> /dev/null & ) ; wait $PID
 exit_code=$?
 
 # Exit code values - https://github.com/dotnet/xharness/blob/main/src/Microsoft.DotNet.XHarness.Common/CLI/ExitCode.cs