Implement measures to fix flakey tests (#3850)

src/Azure.Functions.Cli/Common/Executable.cs

@@ -3,13 +3,14 @@
 using System.ComponentModel;
 using System.Diagnostics;
 using System.IO;
+using System.Runtime.InteropServices;
 using System.Threading.Tasks;
 using Azure.Functions.Cli.Extensions;
 using static Azure.Functions.Cli.Common.OutputTheme;
 
 namespace Azure.Functions.Cli.Common
 {
-    internal class Executable
+    internal class Executable : IAsyncDisposable
     {
         private readonly string _arguments;
         private readonly string _exeName;
@@ -18,6 +19,8 @@ internal class Executable
         private readonly bool _visibleProcess;
         private readonly string _workingDirectory;
         private readonly IDictionary<string, string> _environmentVariables;
+        private JobObjectRegistry _jobObjectRegistry;
+        private bool _disposed;
 
         public Executable(
             string exeName,
@@ -28,7 +31,7 @@ public Executable(
             string workingDirectory = null,
             IDictionary<string, string> environmentVariables = null)
         {
-            _exeName = exeName;
+            _exeName = exeName ?? throw new ArgumentNullException(nameof(exeName));
             _arguments = arguments;
             _streamOutput = streamOutput;
             _shareConsole = shareConsole;
@@ -98,44 +101,71 @@ public async Task<int> RunAsync(Action<string> outputCallback = null, Action<str
             try
             {
                 Process.Start();
+
+                if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
+                {
+                    // Ensure child processes are cleaned up
+                    _jobObjectRegistry = new JobObjectRegistry();
+                    _jobObjectRegistry.Register(Process);
+                }
+
                 if (_streamOutput)
                 {
                     Process.BeginOutputReadLine();
                     Process.BeginErrorReadLine();
                 }
-                
+
                 if (!string.IsNullOrEmpty(stdIn))
                 {
                     Process.StandardInput.WriteLine(stdIn);
                     Process.StandardInput.Close();
                 }
-            }
-            catch (Win32Exception ex)
-            {
-                if (ex.Message == "The system cannot find the file specified")
+
+                if (timeout is null)
+                {
+                    return await exitCodeTask.ConfigureAwait(false);
+                }
+                else
                 {
-                    throw new FileNotFoundException(ex.Message, ex);
+                    return await exitCodeTask.WaitAsync(timeout.Value).ConfigureAwait(false);
                 }
-                throw ex;
             }
+            catch (TimeoutException)
+            {
+                throw new TimeoutException($"Process {_exeName} didn't exit within the specified timeout.");
+            }
+            catch (Win32Exception ex) when (ex.Message.Contains("cannot find the file specified"))
+            {
+                throw new FileNotFoundException(ex.Message, ex);
+            }
+        }
 
-            if (timeout == null)
+        public async ValueTask DisposeAsync()
+        {
+            if (_disposed)
             {
-                return await exitCodeTask;
+                return;
             }
-            else
+
+            if (Process is not null)
             {
-                await Task.WhenAny(exitCodeTask, Task.Delay(timeout.Value));
-                if (exitCodeTask.IsCompleted)
+                try
                 {
-                    return exitCodeTask.Result;
+                    if (!Process.HasExited)
+                    {
+                        Process.Kill();
+                        await Process.WaitForExitAsync();
+                    }
                 }
-                else
+                finally
                 {
-                    Process.Kill();
-                    throw new Exception("Process didn't exit within specified timeout");
+                    Process.Dispose();
                 }
             }
+
+            _jobObjectRegistry?.Dispose();
+
+            _disposed = true;
         }
     }
 }

src/Azure.Functions.Cli/Common/JobObjectRegistry.cs

@@ -0,0 +1,146 @@
+// Copyright (c) .NET Foundation. All rights reserved.
+// Licensed under the MIT License. See License.txt in the project root for license information.
+
+using System;
+using System.Diagnostics;
+using System.Runtime.InteropServices;
+
+namespace Azure.Functions.Cli.Common
+{
+    // Taken from: https://github.com/Azure/azure-functions-host/blob/69111926ee920d4ba10829c8fa34303bb8165a42/src/WebJobs.Script/Workers/ProcessManagement/JobObjectRegistry.cs
+    // This kills child func.exe even if tests are killed from VS mid-run.
+
+    // Registers processes on windows with a job object to ensure disposal after parent exit.
+    internal class JobObjectRegistry : IDisposable
+    {
+        private IntPtr _handle;
+        private bool _disposed = false;
+
+        public JobObjectRegistry()
+        {
+            _handle = CreateJobObject(null, null);
+
+            var info = new JOBOBJECT_BASIC_LIMIT_INFORMATION
+            {
+                LimitFlags = 0x2000
+            };
+
+            var extendedInfo = new JOBOBJECT_EXTENDED_LIMIT_INFORMATION
+            {
+                BasicLimitInformation = info
+            };
+
+            int length = Marshal.SizeOf(typeof(JOBOBJECT_EXTENDED_LIMIT_INFORMATION));
+            IntPtr extendedInfoPtr = Marshal.AllocHGlobal(length);
+            Marshal.StructureToPtr(extendedInfo, extendedInfoPtr, false);
+
+            if (!SetInformationJobObject(_handle, JobObjectInfoType.ExtendedLimitInformation, extendedInfoPtr, (uint)length))
+            {
+                throw new Exception(string.Format("Unable to set information.  Error: {0}", Marshal.GetLastWin32Error()));
+            }
+        }
+
+        public bool Register(Process proc)
+        {
+            return AssignProcessToJobObject(_handle, proc.Handle);
+        }
+
+        [DllImport("kernel32.dll", CharSet = CharSet.Unicode, SetLastError = true)]
+        private static extern IntPtr CreateJobObject(object a, string lpName);
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        private static extern bool SetInformationJobObject(IntPtr hJob, JobObjectInfoType infoType, IntPtr lpJobObjectInfo, uint cbJobObjectInfoLength);
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        private static extern bool AssignProcessToJobObject(IntPtr job, IntPtr process);
+
+        [DllImport("kernel32.dll", SetLastError = true)]
+        private static extern bool CloseHandle(IntPtr job);
+
+        public void Dispose()
+        {
+            Dispose(true);
+            GC.SuppressFinalize(this);
+        }
+
+        private void Dispose(bool disposing)
+        {
+            if (_disposed)
+            {
+                return;
+            }
+
+            if (disposing)
+            {
+                // Dispose of managed resources.
+            }
+
+            Close();
+            _disposed = true;
+        }
+
+        public void Close()
+        {
+            if (_handle != IntPtr.Zero)
+            {
+                CloseHandle(_handle);
+            }
+            _handle = IntPtr.Zero;
+        }
+    }
+
+    public enum JobObjectInfoType
+    {
+        AssociateCompletionPortInformation = 7,
+        BasicLimitInformation = 2,
+        BasicUIRestrictions = 4,
+        EndOfJobTimeInformation = 6,
+        ExtendedLimitInformation = 9,
+        SecurityLimitInformation = 5,
+        GroupInformation = 11
+    }
+
+    [StructLayout(LayoutKind.Sequential)]
+    internal struct IO_COUNTERS
+    {
+        public ulong ReadOperationCount;
+        public ulong WriteOperationCount;
+        public ulong OtherOperationCount;
+        public ulong ReadTransferCount;
+        public ulong WriteTransferCount;
+        public ulong OtherTransferCount;
+    }
+
+    [StructLayout(LayoutKind.Sequential)]
+    internal struct JOBOBJECT_BASIC_LIMIT_INFORMATION
+    {
+        public long PerProcessUserTimeLimit;
+        public long PerJobUserTimeLimit;
+        public uint LimitFlags;
+        public UIntPtr MinimumWorkingSetSize;
+        public UIntPtr MaximumWorkingSetSize;
+        public uint ActiveProcessLimit;
+        public UIntPtr Affinity;
+        public uint PriorityClass;
+        public uint SchedulingClass;
+    }
+
+    [StructLayout(LayoutKind.Sequential)]
+    public struct SECURITY_ATTRIBUTES
+    {
+        public uint nLength;
+        public IntPtr lpSecurityDescriptor;
+        public int bInheritHandle;
+    }
+
+    [StructLayout(LayoutKind.Sequential)]
+    internal struct JOBOBJECT_EXTENDED_LIMIT_INFORMATION
+    {
+        public JOBOBJECT_BASIC_LIMIT_INFORMATION BasicLimitInformation;
+        public IO_COUNTERS IoInfo;
+        public UIntPtr ProcessMemoryLimit;
+        public UIntPtr JobMemoryLimit;
+        public UIntPtr PeakProcessMemoryUsed;
+        public UIntPtr PeakJobMemoryUsed;
+    }
+}