Throw OCE from canceled requests even when remote disconnects before sending cancellation acknowledgment

src/StreamJsonRpc.Tests/JsonRpcWithFatalExceptionsTests.cs

@@ -150,7 +150,7 @@ public async Task StreamsStayOpenOnOperationCanceled()
     }
 
     [Fact]
-    public async Task CancelMayStillReturnErrorFromServer()
+    public async Task CancelExceptionPreferredOverConnectionLost()
     {
         using (var cts = new CancellationTokenSource())
         {
@@ -159,7 +159,10 @@ public async Task CancelMayStillReturnErrorFromServer()
             cts.Cancel();
             this.server.AllowServerMethodToReturn.Set();
 
-            await Assert.ThrowsAnyAsync<ConnectionLostException>(() => invokeTask);
+            // When the remote hangs up while the local side has an outstanding request,
+            // we expect ConnectionLostException to be thrown locally unless the request was already canceled anyway.
+            var ex = await Assert.ThrowsAnyAsync<OperationCanceledException>(() => invokeTask);
+            Assert.Equal(cts.Token, ex.CancellationToken);
             Assert.Equal(Server.ThrowAfterCancellationMessage, this.serverRpc.FaultException.Message);
             Assert.Equal(1, this.serverRpc.IsFatalExceptionCount);
         }
@@ -170,6 +173,50 @@ public async Task CancelMayStillReturnErrorFromServer()
         Assert.False(this.serverRpc.IsDisposed);
     }
 
+    [Fact]
+    public async Task DisposedClientResultsInCancelled()
+    {
+        using (var cts = new CancellationTokenSource())
+        {
+            var invokeTask = this.clientRpc.InvokeWithCancellationAsync<string>(nameof(Server.AsyncMethodFaultsAfterCancellation), new[] { "a" }, cts.Token);
+            await this.server.ServerMethodReached.WaitAsync(this.TimeoutToken);
+            this.clientRpc.Dispose();
+            this.server.AllowServerMethodToReturn.Set();
+
+            // Connection was closed before error was sent from the server
+            await Assert.ThrowsAnyAsync<OperationCanceledException>(() => invokeTask);
+            Assert.Equal(0, this.serverRpc.IsFatalExceptionCount);
+        }
+
+        Assert.True(((IDisposableObservable)this.clientMessageHandler).IsDisposed);
+        Assert.True(((IDisposableObservable)this.serverMessageHandler).IsDisposed);
+        Assert.True(this.clientRpc.IsDisposed);
+        Assert.False(this.serverRpc.IsDisposed);
+    }
+
+    [Fact]
+    public async Task UnexpectedDisconnectResultsInConnectionLostException()
+    {
+        using (var cts = new CancellationTokenSource())
+        {
+            var invokeTask = this.clientRpc.InvokeWithCancellationAsync<string>(nameof(Server.AsyncMethodFaultsAfterCancellation), new[] { "a" }, cts.Token);
+            await this.server.ServerMethodReached.WaitAsync(this.TimeoutToken);
+
+            // Simulate an unexpected lost connection
+            ((IDisposable)this.serverMessageHandler).Dispose();
+            this.server.AllowServerMethodToReturn.Set();
+
+            // Connection was closed before error was sent from the server
+            await Assert.ThrowsAnyAsync<ConnectionLostException>(() => invokeTask);
+            Assert.Equal(0, this.serverRpc.IsFatalExceptionCount);
+        }
+
+        Assert.True(((IDisposableObservable)this.clientMessageHandler).IsDisposed);
+        Assert.True(((IDisposableObservable)this.serverMessageHandler).IsDisposed);
+        Assert.False(this.clientRpc.IsDisposed);
+        Assert.False(this.serverRpc.IsDisposed);
+    }
+
     [Fact]
     public async Task AggregateExceptionIsNotRemovedFromAsyncMethod()
     {

src/StreamJsonRpc/JsonRpc.cs

@@ -1076,7 +1076,15 @@ protected async Task<TResult> InvokeCoreAsync<TResult>(long? id, string targetNa
                                     this.TraceSource.TraceEvent(TraceEventType.Warning, (int)TraceEvents.RequestAbandonedByRemote, "Aborting pending request \"{0}\" because the connection was lost.", id);
                                 }
 
-                                tcs.TrySetException(new ConnectionLostException());
+                                if (cancellationToken.IsCancellationRequested || this.IsDisposed)
+                                {
+                                    // Consider lost connection to be result of task canceled or disposed and set state to canceled
+                                    tcs.TrySetCanceled(cancellationToken.IsCancellationRequested ? cancellationToken : CancellationToken.None);
+                                }
+                                else
+                                {
+                                    tcs.TrySetException(new ConnectionLostException());
+                                }
                             }
                             else if (response is JsonRpcError error)
                             {
@@ -1129,7 +1137,7 @@ protected async Task<TResult> InvokeCoreAsync<TResult>(long? id, string targetNa
                     }
                 }
             }
-            catch (OperationCanceledException ex) when (this.DisconnectedToken.IsCancellationRequested && !cancellationToken.IsCancellationRequested)
+            catch (OperationCanceledException ex) when (this.DisconnectedToken.IsCancellationRequested && !cancellationToken.IsCancellationRequested && !this.IsDisposed)
             {
                 throw new ConnectionLostException(Resources.ConnectionDropped, ex);
             }