Fix ProcessRunner issues: (Azure#16499)

AlexanderSher · web-flow · commit dc47a89af2ce · 2020-11-09T09:33:07.000-08:00
- Handle buffer overload on large output
- Handle process failed to start
diff --git a/sdk/identity/Azure.Identity/src/IProcess.cs b/sdk/identity/Azure.Identity/src/IProcess.cs
@@ -3,21 +3,22 @@
 
 using System;
 using System.Diagnostics;
-using System.IO;
 
 namespace Azure.Identity
 {
     internal interface IProcess : IDisposable
     {
         bool HasExited { get; }
         int ExitCode { get; }
-        StreamReader StandardOutput { get; }
-        StreamReader StandardError { get; }
         ProcessStartInfo StartInfo { get; set; }
 
         event EventHandler Exited;
+        event DataReceivedEventHandler OutputDataReceived;
+        event DataReceivedEventHandler ErrorDataReceived;
 
-        void Start();
+        bool Start();
         void Kill();
+        void BeginOutputReadLine();
+        void BeginErrorReadLine();
     }
 }
diff --git a/sdk/identity/Azure.Identity/src/ProcessRunner.cs b/sdk/identity/Azure.Identity/src/ProcessRunner.cs
@@ -2,6 +2,8 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Collections.Generic;
+using System.Diagnostics;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -14,6 +16,11 @@ internal sealed class ProcessRunner
         private readonly IProcess _process;
         private readonly TimeSpan _timeout;
         private readonly TaskCompletionSource<string> _tcs;
+        private readonly TaskCompletionSource<ICollection<string>> _outputTcs;
+        private readonly TaskCompletionSource<ICollection<string>> _errorTcs;
+        private readonly ICollection<string> _outputData;
+        private readonly ICollection<string> _errorData;
+
         private readonly CancellationToken _cancellationToken;
         private readonly CancellationTokenSource _timeoutCts;
         private CancellationTokenRegistration _ctRegistration;
@@ -22,6 +29,11 @@ public ProcessRunner(IProcess process, TimeSpan timeout, CancellationToken cance
         {
             _process = process;
             _timeout = timeout;
+
+            _outputData = new List<string>();
+            _errorData = new List<string>();
+            _outputTcs = new TaskCompletionSource<ICollection<string>>(TaskCreationOptions.RunContinuationsAsynchronously);
+            _errorTcs = new TaskCompletionSource<ICollection<string>>(TaskCreationOptions.RunContinuationsAsynchronously);
             _tcs = new TaskCompletionSource<string>(TaskCreationOptions.RunContinuationsAsynchronously);
 
             if (timeout.TotalMilliseconds >= 0)
@@ -56,27 +68,37 @@ private void StartProcess()
                 return;
             }
 
-            _process.Exited += (o, e) => HandleExit();
-
             _process.StartInfo.UseShellExecute = false;
             _process.StartInfo.RedirectStandardOutput = true;
             _process.StartInfo.RedirectStandardError = true;
 
+            _process.OutputDataReceived += (sender, args) => OnDataReceived(args, _outputData, _outputTcs);
+            _process.ErrorDataReceived += (sender, args) => OnDataReceived(args, _errorData, _errorTcs);
+            _process.Exited += (o, e) => _ = HandleExitAsync();
+
             _timeoutCts?.CancelAfter(_timeout);
 
-            _process.Start();
+            if (!_process.Start())
+            {
+                TrySetException(new InvalidOperationException($"Failed to start process '{_process.StartInfo.FileName}'"));
+            }
+
+            _process.BeginOutputReadLine();
+            _process.BeginErrorReadLine();
             _ctRegistration = _cancellationToken.Register(HandleCancel, false);
         }
 
-        private void HandleExit()
+        private async ValueTask HandleExitAsync()
         {
             if (_process.ExitCode == 0)
             {
-                TrySetResult(_process.StandardOutput.ReadToEnd());
+                ICollection<string> output = await _outputTcs.Task.ConfigureAwait(false);
+                TrySetResult(string.Join(Environment.NewLine, output));
             }
             else
             {
-                TrySetException(new InvalidOperationException(_process.StandardError.ReadToEnd()));
+                ICollection<string> error = await _errorTcs.Task.ConfigureAwait(false);
+                TrySetException(new InvalidOperationException(string.Join(Environment.NewLine, error)));
             }
         }
 
@@ -103,6 +125,18 @@ private void HandleCancel()
             TrySetCanceled();
         }
 
+        private static void OnDataReceived(DataReceivedEventArgs args, ICollection<string> data, TaskCompletionSource<ICollection<string>> tcs)
+        {
+            if (args.Data != null)
+            {
+                data.Add(args.Data);
+            }
+            else
+            {
+                tcs.SetResult(data);
+            }
+        }
+
         private void TrySetResult(string result)
         {
             DisposeProcess();
diff --git a/sdk/identity/Azure.Identity/src/ProcessService.cs b/sdk/identity/Azure.Identity/src/ProcessService.cs
@@ -30,8 +30,6 @@ public ProcessWrapper(ProcessStartInfo processStartInfo)
 
             public bool HasExited => _process.HasExited;
             public int ExitCode => _process.ExitCode;
-            public StreamReader StandardOutput => _process.StandardOutput;
-            public StreamReader StandardError => _process.StandardError;
 
             public ProcessStartInfo StartInfo
             {
@@ -45,8 +43,22 @@ public event EventHandler Exited
                 remove => _process.Exited -= value;
             }
 
-            public void Start() => _process.Start();
+            public event DataReceivedEventHandler OutputDataReceived
+            {
+                add => _process.OutputDataReceived += value;
+                remove => _process.OutputDataReceived -= value;
+            }
+
+            public event DataReceivedEventHandler ErrorDataReceived
+            {
+                add => _process.ErrorDataReceived += value;
+                remove => _process.ErrorDataReceived -= value;
+            }
+
+            public bool Start() => _process.Start();
             public void Kill() => _process.Kill();
+            public void BeginOutputReadLine() => _process.BeginOutputReadLine();
+            public void BeginErrorReadLine() => _process.BeginErrorReadLine();
             public void Dispose() => _process.Dispose();
         }
     }
diff --git a/sdk/identity/Azure.Identity/tests/ProcessRunnerTests.cs b/sdk/identity/Azure.Identity/tests/ProcessRunnerTests.cs
@@ -4,6 +4,7 @@
 using System;
 using System.ComponentModel;
 using System.Diagnostics;
+using System.Linq;
 using System.Runtime.InteropServices;
 using System.Text;
 using System.Threading;
@@ -42,7 +43,40 @@ public async Task ProcessRunnerRealProcessSucceeded()
             var runner = new ProcessRunner(process, TimeSpan.FromSeconds(30), default);
             var result = await Run(runner);
 
-            Assert.AreEqual(output, result.TrimEnd());
+            Assert.AreEqual(output, result);
+        }
+
+        [Test]
+        public async Task ProcessRunnerRealProcessLargeOutputSucceeded()
+        {
+            var output = string.Concat(Enumerable.Repeat("Test output", 500));
+            var command = $"echo {output}";
+            var process = CreateRealProcess(command, command);
+            var runner = new ProcessRunner(process, TimeSpan.FromSeconds(30), default);
+            var result = await Run(runner);
+
+            Assert.AreEqual(output, result);
+        }
+
+        [Test]
+        public async Task ProcessRunnerRealProcessMultipleLinesSucceeded()
+        {
+            var output = "Test output";
+            var command = string.Join("&&", Enumerable.Repeat($"echo {output}", 100));
+            var process = CreateRealProcess(command, command);
+            var runner = new ProcessRunner(process, TimeSpan.FromSeconds(30), default);
+            var result = await Run(runner);
+
+            Assert.AreEqual(string.Join(Environment.NewLine, Enumerable.Repeat(output, 100)), result);
+        }
+
+        [Test]
+        public void ProcessRunnerProcessFailsToStart()
+        {
+            var process = new TestProcess { FailedToStart = true };
+            var runner = new ProcessRunner(process, TimeSpan.FromSeconds(30), default);
+
+            Assert.CatchAsync<InvalidOperationException>(async () => await Run(runner));
         }
 
         [Test]
@@ -120,6 +154,17 @@ public void ProcessRunnerRealProcessFailedWithErrorMessage()
             Assert.AreEqual(error, exception.Message.Trim());
         }
 
+        [Test]
+        public void ProcessRunnerRealProcessFailedWithLargeErrorMessage()
+        {
+            var error = string.Concat(Enumerable.Repeat("Test error", 500));;
+            var process = CreateRealProcess($"echo {error} 1>&2 & exit 1", $">&2 echo {error} & exit 1");
+            var runner = new ProcessRunner(process, TimeSpan.FromSeconds(30), default);
+
+            var exception = Assert.CatchAsync<InvalidOperationException>(async () => await Run(runner));
+            Assert.AreEqual(error, exception.Message.Trim());
+        }
+
         [Test]
         public void ProcessRunnerFailedOnKillProcess()
         {
diff --git a/sdk/identity/Azure.Identity/tests/TestProcess.cs b/sdk/identity/Azure.Identity/tests/TestProcess.cs
@@ -3,22 +3,27 @@
 
 using System;
 using System.Diagnostics;
-using System.IO;
+using System.Linq.Expressions;
+using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
 
 namespace Azure.Identity.Tests
 {
     internal sealed class TestProcess : IProcess
     {
+        private static readonly Lazy<Func<string, DataReceivedEventArgs>> s_createDataReceivedEventArgs = new Lazy<Func<string, DataReceivedEventArgs>>(() =>
+        {
+            ConstructorInfo constructor = typeof(DataReceivedEventArgs).GetConstructors(BindingFlags.NonPublic | BindingFlags.Instance)[0];
+            ParameterExpression dataParameter = Expression.Parameter(typeof(string), "data");
+            NewExpression callConstructor = Expression.New(constructor, dataParameter);
+            return Expression.Lambda<Func<string, DataReceivedEventArgs>>(callConstructor, dataParameter).Compile();
+        });
+
         private bool _hasStarted;
         private bool _hasExited;
         private int _exitCode;
         private CancellationTokenSource _cts;
-        private StreamReader _outputStreamReader;
-        private StreamReader _errorStreamReader;
-        private StreamWriter _outputStreamWriter;
-        private StreamWriter _errorStreamWriter;
         private ProcessStartInfo _startInfo;
 
         public ProcessStartInfo StartInfo
@@ -27,17 +32,14 @@ public ProcessStartInfo StartInfo
             set => _startInfo = value;
         }
 
+        public bool FailedToStart { get; set; }
         public string Output { get; set; }
         public string Error { get; set; }
         public int? CodeOnExit { get; set; }
         public int Timeout { get; set; }
         public Exception ExceptionOnProcessKill { get; set; }
 
-        public void Dispose()
-        {
-            _outputStreamReader?.Dispose();
-            _errorStreamReader?.Dispose();
-        }
+        public void Dispose() { }
 
         public bool HasExited
         {
@@ -65,20 +67,21 @@ public int ExitCode
             }
         }
 
-        public StreamReader StandardOutput => _outputStreamReader;
-        public StreamReader StandardError => _errorStreamReader;
-
         public event EventHandler Exited;
         public event EventHandler Started;
+        public event DataReceivedEventHandler OutputDataReceived;
+        public event DataReceivedEventHandler ErrorDataReceived;
 
-        public void Start()
+        public bool Start()
         {
+            if (FailedToStart)
+            {
+                return false;
+            }
+
             _hasStarted = true;
             Started?.Invoke(this, EventArgs.Empty);
 
-            Create(out _outputStreamReader, out _outputStreamWriter);
-            Create(out _errorStreamReader, out _errorStreamWriter);
-
             if (Timeout > 0)
             {
                 _cts = new CancellationTokenSource();
@@ -88,6 +91,8 @@ public void Start()
             {
                 Task.Run(FinishProcessRun);
             }
+
+            return true;
         }
 
         private void FinishProcessRun(Task delayTask)
@@ -100,31 +105,29 @@ private void FinishProcessRun(Task delayTask)
 
         private void FinishProcessRun()
         {
-            WriteToStream(Output, _outputStreamWriter);
-            WriteToStream(Error, _errorStreamWriter);
+            Notify(Output, OutputDataReceived);
+            Notify(Error, ErrorDataReceived);
 
             _hasExited = true;
             _exitCode = CodeOnExit ?? (Error != default ? 1 : 0);
             Exited?.Invoke(this, EventArgs.Empty);
         }
 
-        private static void Create(out StreamReader reader, out StreamWriter writer)
+        private void Notify(string data, DataReceivedEventHandler handler)
         {
-            var stream = new MemoryStream();
-            reader = new StreamReader(stream);
-            writer = new StreamWriter(stream);
-        }
-
-        private static void WriteToStream(string str, StreamWriter writer)
-        {
-            if (str == default)
+            if (handler == default)
             {
                 return;
             }
 
-            writer.Write(str);
-            writer.Flush();
-            writer.BaseStream.Seek(0, SeekOrigin.Begin);
+            Task.Run(() =>
+            {
+                if (data != default)
+                {
+                    handler(this, CreateDataReceivedEventArgs(data));
+                }
+                handler(this, CreateDataReceivedEventArgs(null));
+            });
         }
 
         public void Kill()
@@ -137,5 +140,10 @@ public void Kill()
                 throw ExceptionOnProcessKill;
             }
         }
+
+        public void BeginOutputReadLine() {}
+        public void BeginErrorReadLine() {}
+
+        private static DataReceivedEventArgs CreateDataReceivedEventArgs(string data) => s_createDataReceivedEventArgs.Value(data);
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -3,21 +3,22 @@`
`3`	`3`
`4`	`4`	`using System;`
`5`	`5`	`using System.Diagnostics;`
`6`		`-using System.IO;`
`7`	`6`
`8`	`7`	`namespace Azure.Identity`
`9`	`8`	`{`
`10`	`9`	`internal interface IProcess : IDisposable`
`11`	`10`	`{`
`12`	`11`	`bool HasExited { get; }`
`13`	`12`	`int ExitCode { get; }`
`14`		`- StreamReader StandardOutput { get; }`
`15`		`- StreamReader StandardError { get; }`
`16`	`13`	`ProcessStartInfo StartInfo { get; set; }`
`17`	`14`
`18`	`15`	`event EventHandler Exited;`
	`16`	`+ event DataReceivedEventHandler OutputDataReceived;`
	`17`	`+ event DataReceivedEventHandler ErrorDataReceived;`
`19`	`18`
`20`		`- void Start();`
	`19`	`+ bool Start();`
`21`	`20`	`void Kill();`
	`21`	`+ void BeginOutputReadLine();`
	`22`	`+ void BeginErrorReadLine();`
`22`	`23`	`}`
`23`	`24`	`}`