Merge pull request #2 from Microsoft/dev/andarno/ObjectDisposedExceptions

Throw OperationCanceledExceptions instead of ObjectDisposedExceptions

src/StreamJsonRpc.Shared/DelimitedMessageHandler.cs

@@ -109,13 +109,23 @@ namespace StreamJsonRpc
         public async Task<string> ReadAsync(CancellationToken cancellationToken)
         {
             Verify.Operation(this.ReceivingStream != null, "No receiving stream.");
+            cancellationToken.ThrowIfCancellationRequested();
             Verify.NotDisposed(this);
 
             using (var cts = CancellationTokenSource.CreateLinkedTokenSource(this.DisposalToken, cancellationToken))
             {
-                using (await this.receivingSemaphore.EnterAsync(cts.Token).ConfigureAwait(false))
+                try
                 {
-                    return await this.ReadCoreAsync(cts.Token).ConfigureAwait(false);
+                    using (await this.receivingSemaphore.EnterAsync(cts.Token).ConfigureAwait(false))
+                    {
+                        return await this.ReadCoreAsync(cts.Token).ConfigureAwait(false);
+                    }
+                }
+                catch (ObjectDisposedException)
+                {
+                    // If already canceled, throw that instead of ObjectDisposedException.
+                    cancellationToken.ThrowIfCancellationRequested();
+                    throw;
                 }
             }
         }
@@ -130,6 +140,7 @@ namespace StreamJsonRpc
         {
             Requires.NotNull(content, nameof(content));
             Verify.Operation(this.SendingStream != null, "No sending stream.");
+            cancellationToken.ThrowIfCancellationRequested();
             Verify.NotDisposed(this);
 
             // Capture Encoding as a local since it may change over the time of this method's execution.
@@ -137,9 +148,18 @@ namespace StreamJsonRpc
 
             using (var cts = CancellationTokenSource.CreateLinkedTokenSource(this.DisposalToken, cancellationToken))
             {
-                using (await this.sendingSemaphore.EnterAsync(cts.Token).ConfigureAwait(false))
+                try
                 {
-                    await this.WriteCoreAsync(content, contentEncoding, cts.Token).ConfigureAwait(false);
+                    using (await this.sendingSemaphore.EnterAsync(cts.Token).ConfigureAwait(false))
+                    {
+                        await this.WriteCoreAsync(content, contentEncoding, cts.Token).ConfigureAwait(false);
+                    }
+                }
+                catch (ObjectDisposedException)
+                {
+                    // If already canceled, throw that instead of ObjectDisposedException.
+                    cancellationToken.ThrowIfCancellationRequested();
+                    throw;
                 }
             }
         }

src/StreamJsonRpc.Tests/DelimitedMessageHandlerTests.cs

@@ -74,11 +74,93 @@ public class DelimitedMessageHandlerTests : TestBase
         Assert.Throws<ObjectDisposedException>(() => result.GetAwaiter().GetResult());
     }
 
+    /// <summary>
+    /// Verifies that when both <see cref="ObjectDisposedException"/> and <see cref="OperationCanceledException"/> are appropriate
+    /// when we first invoke the method, the <see cref="OperationCanceledException"/> is thrown.
+    /// </summary>
+    [Fact]
+    public void WriteAsync_PreferOperationCanceledException_AtEntry()
+    {
+        this.handler.Dispose();
+        Assert.Throws<OperationCanceledException>(() => this.handler.WriteAsync("content", PrecanceledToken).GetAwaiter().GetResult());
+    }
+
+    /// <summary>
+    /// Verifies that <see cref="DelimitedMessageHandler.ReadAsync(CancellationToken)"/> prefers throwing
+    /// <see cref="OperationCanceledException"/> over <see cref="ObjectDisposedException"/> when both conditions
+    /// apply while reading (at least when cancellation occurs first).
+    /// </summary>
+    [Fact]
+    public async Task WriteAsync_PreferOperationCanceledException_MidExecution()
+    {
+        var handler = new DelayedWriter(this.sendingStream, this.receivingStream, Encoding.UTF8);
+
+        var cts = new CancellationTokenSource();
+        var writeTask = handler.WriteAsync("content", cts.Token);
+
+        cts.Cancel();
+        handler.Dispose();
+
+        // Unblock writer. It should not throw anything as it is to emulate not recognizing the
+        // CancellationToken before completing its work.
+        handler.WriteBlock.Set();
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => writeTask);
+    }
+
     [Fact]
     public void ReadAsync_ThrowsObjectDisposedException()
     {
         this.handler.Dispose();
         Task result = this.handler.ReadAsync(TimeoutToken);
         Assert.Throws<ObjectDisposedException>(() => result.GetAwaiter().GetResult());
+        Assert.Throws<OperationCanceledException>(() => this.handler.ReadAsync(PrecanceledToken).GetAwaiter().GetResult());
+    }
+
+    /// <summary>
+    /// Verifies that when both <see cref="ObjectDisposedException"/> and <see cref="OperationCanceledException"/> are appropriate
+    /// when we first invoke the method, the <see cref="OperationCanceledException"/> is thrown.
+    /// </summary>
+    [Fact]
+    public void ReadAsync_PreferOperationCanceledException_AtEntry()
+    {
+        this.handler.Dispose();
+        Assert.Throws<OperationCanceledException>(() => this.handler.ReadAsync(PrecanceledToken).GetAwaiter().GetResult());
+    }
+
+    /// <summary>
+    /// Verifies that <see cref="DelimitedMessageHandler.ReadAsync(CancellationToken)"/> prefers throwing
+    /// <see cref="OperationCanceledException"/> over <see cref="ObjectDisposedException"/> when both conditions
+    /// apply while reading (at least when cancellation occurs first).
+    /// </summary>
+    [Fact]
+    public async Task ReadAsync_PreferOperationCanceledException_MidExecution()
+    {
+        var cts = new CancellationTokenSource();
+        var readTask = this.handler.ReadAsync(cts.Token);
+
+        cts.Cancel();
+        this.handler.Dispose();
+
+        await Assert.ThrowsAnyAsync<OperationCanceledException>(() => readTask);
+    }
+
+    private class DelayedWriter : DelimitedMessageHandler
+    {
+        internal readonly AsyncManualResetEvent WriteBlock = new AsyncManualResetEvent();
+
+        public DelayedWriter(Stream sendingStream, Stream receivingStream, Encoding encoding)
+            : base(sendingStream, receivingStream, encoding)
+        {
+        }
+
+        protected override Task<string> ReadCoreAsync(CancellationToken cancellationToken)
+        {
+            throw new NotImplementedException();
+        }
+
+        protected override Task WriteCoreAsync(string content, Encoding contentEncoding, CancellationToken cancellationToken)
+        {
+            return WriteBlock.WaitAsync();
+        }
     }
 }

src/StreamJsonRpc.Tests/TestBase.cs

@@ -19,6 +19,8 @@ public abstract class TestBase : IDisposable
 
     protected static readonly TimeSpan UnexpectedTimeout = Debugger.IsAttached ? Timeout.InfiniteTimeSpan : TimeSpan.FromSeconds(5);
 
+    protected static readonly CancellationToken PrecanceledToken = new CancellationToken(canceled: true);
+
     protected static CancellationToken ExpectedTimeoutToken => new CancellationTokenSource(ExpectedTimeout).Token;
 
     protected TestBase(ITestOutputHelper logger)