Consolidates task cancellation helpers

src/MsBuildPipeLogger.Server/AnonymousPipeLoggerServer.cs

@@ -44,12 +44,6 @@ protected override void Connect()
             // But I can only catch the pipe disposal from cancellation after the handle has been disposed
             Buffer.FillFromStream(PipeStream, CancellationToken);
 
-            // This doesn't actually disconnect, it just disposes the client handle
-            Disconnect();
-        }
-
-        protected override void Disconnect()
-        {
             // Dispose the client handle if we asked for one
             // If we don't do this we won't get notified when the stream closes, see https://stackoverflow.com/q/39682602/807064
             if (_clientHandle != null)

src/MsBuildPipeLogger.Server/CancellationTokenExtensions.cs

@@ -0,0 +1,54 @@
+using System;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace MsBuildPipeLogger
+{
+    internal static class CancellationTokenExtensions
+    {
+        public static void Try(this CancellationToken cancellationToken, Action action) =>
+            cancellationToken.Try(action, null);
+
+        public static void Try(this CancellationToken cancellationToken, Action action, Action cancelled) =>
+            cancellationToken.Try<object>(() =>
+            {
+                action();
+                return null;
+            }, () => null);
+
+        public static TResult Try<TResult>(this CancellationToken cancellationToken, Func<TResult> action) =>
+            cancellationToken.Try(action, null);
+
+        public static TResult Try<TResult>(this CancellationToken cancellationToken, Func<TResult> action, Func<TResult> cancelled)
+        {
+            if (cancellationToken.IsCancellationRequested)
+            {
+                return cancelled == null ? default(TResult) : cancelled();
+            }
+            try
+            {
+                return action();
+            }
+            catch (TaskCanceledException)
+            {
+                // Thrown if the task itself was cancelled from inside the read method
+                return cancelled == null ? default(TResult) : cancelled();
+            }
+            catch (OperationCanceledException)
+            {
+                // Thrown if the operation was cancelled (I.e., the task didn't deal with cancellation)
+                return cancelled == null ? default(TResult) : cancelled();
+            }
+            catch (AggregateException ex)
+            {
+                // Sometimes the cancellation exceptions are thrown in aggregate
+                if (!(ex.InnerException is TaskCanceledException)
+                    && !(ex.InnerException is OperationCanceledException))
+                {
+                    throw;
+                }
+                return cancelled == null ? default(TResult) : cancelled();
+            }
+        }
+    }
+}

src/MsBuildPipeLogger.Server/NamedPipeLoggerServer.cs

@@ -28,26 +28,7 @@ public NamedPipeLoggerServer(string pipeName, CancellationToken cancellationToke
         {
         }
 
-        protected override void Connect()
-        {
-            try
-            {
-                PipeStream.WaitForConnectionAsync(CancellationToken).Wait();
-            }
-            catch(Exception)
-            {
-                // If something went wrong here, just close it up
-                // This will always throw when the CancellationToken triggers pipe disposal on cancel handler
-                Disconnect();
-            }
-        }
-
-        protected override void Disconnect()
-        {
-            if (PipeStream.IsConnected)
-            {
-                PipeStream.Disconnect();
-            }
-        }
+        protected override void Connect() =>
+            CancellationToken.Try(() => PipeStream.WaitForConnectionAsync(CancellationToken).Wait());
     }
 }

src/MsBuildPipeLogger.Server/PipeBuffer.cs

@@ -113,39 +113,13 @@ public Buffer(byte[] buffer, int offset, int count)
 
             public int FillFromStream(Stream stream, CancellationToken cancellationToken)
             {
-                try
+                Count = cancellationToken.Try(() =>
                 {
-                    if (cancellationToken.IsCancellationRequested)
-                    {
-                        Count = 0;
-                    }
-                    else
-                    {
-                        _offset = 0;
-                        Task<int> readTask = stream.ReadAsync(_buffer, _offset, BufferSize, cancellationToken);
-                        readTask.Wait(cancellationToken);
-                        Count = readTask.Status == TaskStatus.Canceled ? 0 : readTask.Result;
-                    }
-                } 
-                catch(TaskCanceledException)
-                {
-                    // Thrown if the task itself was cancelled from inside the read method
-                    Count = 0;
-                }
-                catch (OperationCanceledException)
-                {
-                    // Thrown if the operation was cancelled (I.e., the task didn't deal with cancellation)
-                    Count = 0;
-                }
-                catch (AggregateException ex)
-                {
-                    // Sometimes the cancellation exceptions are thrown in aggregate
-                    if(!(ex.InnerException is TaskCanceledException)
-                        && !(ex.InnerException is OperationCanceledException))
-                    {
-                        throw;
-                    }
-                }
+                    _offset = 0;
+                    Task<int> readTask = stream.ReadAsync(_buffer, _offset, BufferSize, cancellationToken);
+                    readTask.Wait(cancellationToken);
+                    return readTask.Status == TaskStatus.Canceled ? 0 : readTask.Result;
+                }, () => 0);
                 return Count;
             }
 

src/MsBuildPipeLogger.Server/PipeLoggerServer.cs

@@ -64,8 +64,6 @@ protected PipeLoggerServer(TPipeStream pipeStream, CancellationToken cancellatio
                     // The pipe was disposed
                 }
 
-                Disconnect();
-
                 // Add a final 0 (BinaryLogRecordKind.EndOfFile) into the stream in case the BuildEventArgsReader is waiting for a read
                 Buffer.Write(new byte[1] { 0 }, 0, 1);
 
@@ -74,9 +72,7 @@ protected PipeLoggerServer(TPipeStream pipeStream, CancellationToken cancellatio
         }
 
         protected abstract void Connect();
-
-        protected abstract void Disconnect();
-
+
         /// <inheritdoc/>
         public BuildEventArgs Read()
         {