diff --git a/sdk/identity/Azure.Identity/src/DataReceivedEventArgsWrapper.cs b/sdk/identity/Azure.Identity/src/DataReceivedEventArgsWrapper.cs new file mode 100644 index 000000000000..849ff4412892 --- /dev/null +++ b/sdk/identity/Azure.Identity/src/DataReceivedEventArgsWrapper.cs @@ -0,0 +1,24 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +using System.Diagnostics; + +namespace Azure.Identity +{ + internal class DataReceivedEventArgsWrapper + { + public DataReceivedEventArgsWrapper(string data) + { + Data = data; + } + + public DataReceivedEventArgsWrapper(DataReceivedEventArgs args) + { + Data = args?.Data; + } + + public string Data { get; } + } + + internal delegate void DataReceivedEventWrapperHandler(object sender, DataReceivedEventArgsWrapper e); +} diff --git a/sdk/identity/Azure.Identity/src/IProcess.cs b/sdk/identity/Azure.Identity/src/IProcess.cs index 06e012fd2a89..b743626a2b2e 100644 --- a/sdk/identity/Azure.Identity/src/IProcess.cs +++ b/sdk/identity/Azure.Identity/src/IProcess.cs @@ -16,6 +16,8 @@ internal interface IProcess : IDisposable ProcessStartInfo StartInfo { get; set; } event EventHandler Exited; + event DataReceivedEventWrapperHandler OutputDataReceived; + event DataReceivedEventWrapperHandler ErrorDataReceived; void Start(); void Kill(); diff --git a/sdk/identity/Azure.Identity/src/ProcessRunner.cs b/sdk/identity/Azure.Identity/src/ProcessRunner.cs index 2b283c01c6ae..7466d9118f82 100644 --- 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.Diagnostics; +using System.Text; using System.Threading; using System.Threading.Tasks; @@ -13,16 +15,24 @@ internal sealed class ProcessRunner { private readonly IProcess _process; private readonly TimeSpan _timeout; - private readonly TaskCompletionSource _tcs; + private readonly TaskCompletionSource _processExitedCompletionSource; private readonly CancellationToken _cancellationToken; private readonly CancellationTokenSource _timeoutCts; private CancellationTokenRegistration _ctRegistration; + private readonly StringBuilder _stdOutBuilder; + private readonly StringBuilder _stdErrBuilder; + private readonly TaskCompletionSource _stdOutCompletionSource; + private readonly TaskCompletionSource _stdErrCompletionSource; public ProcessRunner(IProcess process, TimeSpan timeout, CancellationToken cancellationToken) { _process = process; _timeout = timeout; - _tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + _processExitedCompletionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + _stdOutBuilder = new StringBuilder(); + _stdErrBuilder = new StringBuilder(); + _stdOutCompletionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + _stdErrCompletionSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); if (timeout.TotalMilliseconds >= 0) { @@ -35,29 +45,53 @@ public ProcessRunner(IProcess process, TimeSpan timeout, CancellationToken cance } } - public Task RunAsync() + public async Task RunAsync() { StartProcess(); - return _tcs.Task; + + try + { + await _processExitedCompletionSource.Task.ConfigureAwait(false); + var output = await _stdOutCompletionSource.Task.ConfigureAwait(false); + await _stdErrCompletionSource.Task.ConfigureAwait(false); + return output; + } + finally + { + DisposeProcess(); + } } public string Run() { StartProcess(); #pragma warning disable AZC0102 // Do not use GetAwaiter().GetResult(). - return _tcs.Task.GetAwaiter().GetResult(); + try + { + _processExitedCompletionSource.Task.GetAwaiter().GetResult(); + var output = _stdOutCompletionSource.Task.GetAwaiter().GetResult(); + _stdErrCompletionSource.Task.GetAwaiter().GetResult(); + return output; + } + finally + { + DisposeProcess(); + } #pragma warning restore AZC0102 // Do not use GetAwaiter().GetResult(). } private void StartProcess() { - if (TrySetCanceled() || _tcs.Task.IsCompleted) + if (TrySetCanceled() || _processExitedCompletionSource.Task.IsCompleted) { return; } _process.Exited += (o, e) => HandleExit(); + _process.OutputDataReceived += HandleStdOutDataReceived; + _process.ErrorDataReceived += HandleStdErrDataReceived; + _process.StartInfo.UseShellExecute = false; _process.StartInfo.RedirectStandardOutput = true; _process.StartInfo.RedirectStandardError = true; @@ -72,17 +106,48 @@ private void HandleExit() { if (_process.ExitCode == 0) { - TrySetResult(_process.StandardOutput.ReadToEnd()); + _processExitedCompletionSource.TrySetResult(true); } else { - TrySetException(new InvalidOperationException(_process.StandardError.ReadToEnd())); + _processExitedCompletionSource.TrySetResult(false); + } + } + private void HandleStdErrDataReceived(object sender, DataReceivedEventArgsWrapper e) + { + if (e.Data is null) + { + var result = _stdErrBuilder.ToString(); + if (result.Length == 0) + { + _stdErrCompletionSource.TrySetResult(result); + } + else + { + _stdErrCompletionSource.TrySetException(new InvalidOperationException(result)); + } + } + else + { + _stdErrBuilder.Append(e.Data); + } + } + + private void HandleStdOutDataReceived(object sender, DataReceivedEventArgsWrapper e) + { + if (e.Data is null) + { + _stdOutCompletionSource.TrySetResult(_stdOutBuilder.ToString()); + } + else + { + _stdOutBuilder.Append(e.Data); } } private void HandleCancel() { - if (_tcs.Task.IsCompleted) + if (_processExitedCompletionSource.Task.IsCompleted) { return; } @@ -95,7 +160,7 @@ private void HandleCancel() } catch (Exception ex) { - TrySetException(ex); + _processExitedCompletionSource.TrySetException(ex); return; } } @@ -103,31 +168,23 @@ private void HandleCancel() TrySetCanceled(); } - private void TrySetResult(string result) - { - DisposeProcess(); - _tcs.TrySetResult(result); - } - private bool TrySetCanceled() { if (_cancellationToken.IsCancellationRequested) { - DisposeProcess(); - _tcs.TrySetCanceled(_cancellationToken); + _stdOutCompletionSource.TrySetCanceled(_cancellationToken); + _stdErrCompletionSource.TrySetCanceled(_cancellationToken); + _processExitedCompletionSource.TrySetCanceled(_cancellationToken); } return _cancellationToken.IsCancellationRequested; } - private void TrySetException(Exception exception) - { - DisposeProcess(); - _tcs.TrySetException(exception); - } - private void DisposeProcess() { + _process.OutputDataReceived -= HandleStdOutDataReceived; + _process.ErrorDataReceived -= HandleStdErrDataReceived; + _process.Dispose(); _ctRegistration.Dispose(); _timeoutCts?.Dispose(); diff --git a/sdk/identity/Azure.Identity/src/ProcessService.cs b/sdk/identity/Azure.Identity/src/ProcessService.cs index 15bdecb714ba..5581121b26ca 100644 --- a/sdk/identity/Azure.Identity/src/ProcessService.cs +++ b/sdk/identity/Azure.Identity/src/ProcessService.cs @@ -26,6 +26,9 @@ public ProcessWrapper(ProcessStartInfo processStartInfo) StartInfo = processStartInfo, EnableRaisingEvents = true }; + + _process.OutputDataReceived += HandleOutputDataReceived; + _process.ErrorDataReceived += HandleErrorDataReceived; } public bool HasExited => _process.HasExited; @@ -45,7 +48,20 @@ public event EventHandler Exited remove => _process.Exited -= value; } - public void Start() => _process.Start(); + public event DataReceivedEventWrapperHandler OutputDataReceived; + public event DataReceivedEventWrapperHandler ErrorDataReceived; + + private void HandleErrorDataReceived(object sender, DataReceivedEventArgs e) => ErrorDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper(e)); + private void HandleOutputDataReceived(object sender, DataReceivedEventArgs e) => OutputDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper(e)); + + public void Start() + { + _process.Start(); + + _process.BeginOutputReadLine(); + _process.BeginErrorReadLine(); + } + public void Kill() => _process.Kill(); public void Dispose() => _process.Dispose(); } diff --git a/sdk/identity/Azure.Identity/tests/ProcessRunnerTests.cs b/sdk/identity/Azure.Identity/tests/ProcessRunnerTests.cs index 2ebb7c3835ac..8ab2780a1b5f 100644 --- 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; @@ -45,6 +46,17 @@ public async Task ProcessRunnerRealProcessSucceeded() Assert.AreEqual(output, result.TrimEnd()); } + [Test] + public async Task ProcessRunnerRealProcessLargeOutputSucceeded() + { + var output = string.Concat(Enumerable.Repeat("ab", 3000)); + var process = CreateRealProcess($"echo {output}", $"echo {output}"); + var runner = new ProcessRunner(process, TimeSpan.FromSeconds(30), default); + var result = await Run(runner); + + Assert.AreEqual(output, result.TrimEnd()); + } + [Test] public void ProcessRunnerCanceledByTimeout() { diff --git a/sdk/identity/Azure.Identity/tests/TestProcess.cs b/sdk/identity/Azure.Identity/tests/TestProcess.cs index bb5ea9ac838b..f8edd0d0a3ec 100644 --- a/sdk/identity/Azure.Identity/tests/TestProcess.cs +++ b/sdk/identity/Azure.Identity/tests/TestProcess.cs @@ -70,6 +70,8 @@ public int ExitCode public event EventHandler Exited; public event EventHandler Started; + public event DataReceivedEventWrapperHandler OutputDataReceived; + public event DataReceivedEventWrapperHandler ErrorDataReceived; public void Start() { @@ -101,8 +103,23 @@ private void FinishProcessRun(Task delayTask) private void FinishProcessRun() { WriteToStream(Output, _outputStreamWriter); + + if (Output != default) + { + OutputDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper(Output)); + } + WriteToStream(Error, _errorStreamWriter); + if (Error != default) + { + ErrorDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper(Error)); + } + + // signal completion + OutputDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper((string)null)); + ErrorDataReceived?.Invoke(this, new DataReceivedEventArgsWrapper((string)null)); + _hasExited = true; _exitCode = CodeOnExit ?? (Error != default ? 1 : 0); Exited?.Invoke(this, EventArgs.Empty);