Merge pull request #3871 from AlexanderSher/master

Fix TaskCompletionSourceExtensions.RegisterForCancellation

src/Common/Core/Impl/Extensions/TaskCompletionSourceExtensions.cs

@@ -11,19 +11,17 @@ namespace Microsoft.Common.Core {
             => taskCompletionSource.RegisterForCancellation(-1, cancellationToken);
 
         public static CancellationTokenRegistration RegisterForCancellation<T>(this TaskCompletionSource<T> taskCompletionSource, int millisecondsDelay, CancellationToken cancellationToken) {
-            CancelOnTokenAction<T> action;
             if (millisecondsDelay >= 0) {
                 var cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
                 cts.CancelAfter(millisecondsDelay);
-                action = new CancelOnTokenAction<T>(taskCompletionSource, cts.Token);
-            } else {
-                action = new CancelOnTokenAction<T>(taskCompletionSource, cancellationToken);
+                cancellationToken = cts.Token;
             }
 
+            var action = new CancelOnTokenAction<T>(taskCompletionSource, cancellationToken);
             return cancellationToken.Register(action.Invoke);
         }
 
-        private class CancelOnTokenAction<T> {
+        private struct CancelOnTokenAction<T> {
             private readonly TaskCompletionSource<T> _taskCompletionSource;
             private readonly CancellationToken _cancellationToken;
 
@@ -34,11 +32,11 @@ namespace Microsoft.Common.Core {
 
             public void Invoke() {
                 if (!_taskCompletionSource.Task.IsCompleted) {
-                    Task.Run(new Action(TryCancel));
+                    ThreadPool.QueueUserWorkItem(TryCancel);
                 }
             }
 
-            private void TryCancel() => _taskCompletionSource.TrySetCanceled(_cancellationToken);
+            private void TryCancel(object state) => _taskCompletionSource.TrySetCanceled(_cancellationToken);
         }
     }
 }

src/Common/Core/Impl/OS/IProcess.cs

@@ -16,5 +16,6 @@ namespace Microsoft.Common.Core.OS {
         event EventHandler Exited;
 
         void Kill();
+        bool WaitForExit(int milliseconds);
     }
 }

src/Common/Core/Test/Extensions/TaskCompletionSourceExtensionsTest.cs

@@ -0,0 +1,36 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License. See LICENSE in the project root for license information.
+
+using System;
+using System.Diagnostics.CodeAnalysis;
+using System.Threading;
+using System.Threading.Tasks;
+using FluentAssertions;
+using Microsoft.UnitTests.Core.Threading;
+using Microsoft.UnitTests.Core.XUnit;
+using Microsoft.UnitTests.Core.FluentAssertions;
+using Xunit;
+
+namespace Microsoft.Common.Core.Test.Extensions {
+    [ExcludeFromCodeCoverage]
+    [Category.CoreExtensions]
+    public class TaskCompletionSourceExtensionsTest {
+        [Test]
+        public async Task CanceledOnToken() {
+            var tcs = new TaskCompletionSource<int>();
+            var cts = new CancellationTokenSource();
+            tcs.RegisterForCancellation(cts.Token);
+            cts.CancelAfter(200);
+            await ParallelTools.When(tcs.Task);
+            tcs.Task.Should().BeCanceled();
+        }
+
+        [Test]
+        public async Task CanceledOnTimeout() {
+            var tcs = new TaskCompletionSource<int>();
+            tcs.RegisterForCancellation(200, CancellationToken.None);
+            await ParallelTools.When(tcs.Task);
+            tcs.Task.Should().BeCanceled();
+        }
+    }
+}

src/Common/Core/Test/Microsoft.Common.Core.Test.csproj

@@ -59,6 +59,7 @@
     <Compile Include="Disposables\DisposeTokenTest.cs" />
     <Compile Include="EnumerableExtensionsTest.cs" />
     <Compile Include="Extensions\IOExtensionsTest.cs" />
+    <Compile Include="Extensions\TaskCompletionSourceExtensionsTest.cs" />
     <Compile Include="Fakes\Shell\TestApplication.cs" />
     <Compile Include="Fakes\Shell\TestCompositionCatalog.cs" />
     <Compile Include="Fakes\Shell\TestFileDialog.cs" />

src/Linux/Core/Impl/OS/UnixProcess.cs

@@ -33,5 +33,7 @@ namespace Microsoft.Common.Core.OS {
             // This is needed because broker user cannot kill process running as another user.
             _ps.Kill(_process.Id);
         }
+
+        public bool WaitForExit(int milliseconds) => _process.WaitForExit(milliseconds);
     }
 }

src/Linux/Host/Broker/Impl/Services/LinuxRHostProcessService.cs

@@ -32,7 +32,7 @@ namespace Microsoft.R.Host.Broker.Services {
             var environment = GetHostEnvironment(interpreter, profilePath, userName);
             var password = principal.FindFirst(UnixClaims.RPassword).Value;
 
-            Process process = Utility.AuthenticateAndRunAsUser(_sessionLogger, _ps, userName, password, profilePath, args, environment);
+            var process = Utility.AuthenticateAndRunAsUser(_sessionLogger, _ps, userName, password, profilePath, args, environment);
             process.WaitForExit(250);
             if (process.HasExited && process.ExitCode != 0) {
                 var message = _ps.MessageFromExitCode(process.ExitCode);
@@ -42,7 +42,7 @@ namespace Microsoft.R.Host.Broker.Services {
                 throw new Win32Exception(process.ExitCode);
             }
 
-            return new UnixProcess(_ps, process);
+            return process;
         }
 
         private string GetLoadLibraryPath(string binPath) {

src/Linux/Host/Broker/Impl/Services/LinuxSystemInfoService.cs

@@ -237,7 +237,7 @@ namespace Microsoft.R.Host.Broker.Services {
         }
 
         private IEnumerable<string> ExecuteAndGetOutput(string command, string arguments) {
-            Process proc = null;
+            IProcess proc = null;
             List<string> standardOutData = new List<string>();
             try {
                 ProcessStartInfo psi = new ProcessStartInfo() {

src/Linux/Host/Broker/Impl/Services/Utility.cs

@@ -26,8 +26,8 @@ namespace Microsoft.R.Host.Broker.Services {
         private const string RtvsResult = "rtvs-result";
         private const string RtvsError = "rtvs-error";
 
-        public static Process AuthenticateAndRunAsUser(ILogger<Session> logger, IProcessServices ps, string username, string password, string profileDir, IEnumerable<string> arguments, IDictionary<string, string> environment) {
-            Process proc = CreateRunAsUserProcess(ps, true);
+        public static IProcess AuthenticateAndRunAsUser(ILogger<Session> logger, IProcessServices ps, string username, string password, string profileDir, IEnumerable<string> arguments, IDictionary<string, string> environment) {
+            var proc = CreateRunAsUserProcess(ps, true);
             using (BinaryWriter writer = new BinaryWriter(proc.StandardInput.BaseStream, Encoding.UTF8, true)) {
                 var message = new AuthenticateAndRunMessage() {
                     Username = GetUnixUserName(username),
@@ -48,7 +48,7 @@ namespace Microsoft.R.Host.Broker.Services {
 
         public static bool AuthenticateUser(ILogger<IAuthenticationService> logger, IProcessServices ps,  string username, string password, string allowedGroup, out string profileDir) {
             bool retval = false;
-            Process proc = null;
+            IProcess proc = null;
             string userDir = string.Empty;
             try {
                 proc = CreateRunAsUserProcess(ps, false);
@@ -145,13 +145,14 @@ namespace Microsoft.R.Host.Broker.Services {
             };
         }
 
-        private static Process CreateRunAsUserProcess(IProcessServices ps, bool quietMode) {
-        ProcessStartInfo psi = new ProcessStartInfo();
-            psi.FileName = PathConstants.RunAsUserBinPath;
-            psi.Arguments = quietMode ? "-q" : "";
-            psi.RedirectStandardError = true;
-            psi.RedirectStandardInput = true;
-            psi.RedirectStandardOutput = true;
+        private static IProcess CreateRunAsUserProcess(IProcessServices ps, bool quietMode) {
+            var psi = new ProcessStartInfo {
+                FileName = PathConstants.RunAsUserBinPath,
+                Arguments = quietMode ? "-q" : "",
+                RedirectStandardError = true,
+                RedirectStandardInput = true,
+                RedirectStandardOutput = true
+            };
 
             return ps.Start(psi);
         }

src/Windows/Core/Impl/OS/DotNetProcess.cs

@@ -26,5 +26,6 @@ namespace Microsoft.Common.Core.OS {
         }
 
         public void Kill() => _process.Kill();
+        public bool WaitForExit(int milliseconds) => _process.WaitForExit(milliseconds);
     }
 }

src/Windows/Core/Impl/OS/Win32Process.cs

@@ -53,13 +53,16 @@ namespace Microsoft.Common.Core.OS {
             }, null, -1, true);
         }
 
-        public void WaitForExit(int milliseconds) {
+        public bool WaitForExit(int milliseconds) {
             using (var processWaitHandle = new ProcessWaitHandle(_processHandle)) {
                 if (processWaitHandle.WaitOne(milliseconds)) {
                     // This means the process exited while waiting.
                     SetExitState();
+                    return true;
                 }
             }
+
+            return false;
         }
 
         public void Kill() {