diff --git a/src/StreamJsonRpc/MessageHandlerBase.cs b/src/StreamJsonRpc/MessageHandlerBase.cs index 86535e79..61d46247 100644 --- a/src/StreamJsonRpc/MessageHandlerBase.cs +++ b/src/StreamJsonRpc/MessageHandlerBase.cs @@ -40,6 +40,11 @@ public abstract class MessageHandlerBase : IJsonRpcMessageHandler, IDisposableOb /// private readonly object syncObject = new object(); + /// + /// The source for the task. + /// + private readonly TaskCompletionSource completionSource = new TaskCompletionSource(); + /// /// A value indicating whether the method is in progress. /// @@ -102,6 +107,11 @@ private enum MessageHandlerState /// protected CancellationToken DisposalToken => this.disposalTokenSource.Token; + /// + /// Gets a task that completes when this instance has completed disposal. + /// + private Task Completion => this.completionSource.Task; + /// /// Reads a distinct and complete message from the transport, waiting for one if necessary. /// @@ -142,7 +152,7 @@ public async ValueTask ReadAsync(CancellationToken cancellationT { if (this.CheckIfDisposalAppropriate(MessageHandlerState.Reading)) { - this.Dispose(true); + this.DoDisposeAndCompletion(); } } } @@ -193,25 +203,17 @@ public async ValueTask WriteAsync(JsonRpcMessage content, CancellationToken canc { if (shouldDispose) { - this.Dispose(true); + this.DoDisposeAndCompletion(); } } } +#pragma warning disable VSTHRD002 // We synchronously block, but nothing here should ever require the main thread. /// /// Disposes this instance, and cancels any pending read or write operations. /// - public void Dispose() - { - if (!this.disposalTokenSource.IsCancellationRequested) - { - this.disposalTokenSource.Cancel(); - if (this.CheckIfDisposalAppropriate()) - { - this.Dispose(true); - } - } - } + public void Dispose() => this.DisposeAsync().GetAwaiter().GetResult(); +#pragma warning restore VSTHRD002 /// /// Disposes resources allocated by this instance. @@ -265,6 +267,28 @@ protected virtual void Dispose(bool disposing) /// protected abstract ValueTask FlushAsync(CancellationToken cancellationToken); + /// + /// Disposes this instance, and cancels any pending read or write operations. + /// + private Task DisposeAsync() + { + if (!this.disposalTokenSource.IsCancellationRequested) + { + this.disposalTokenSource.Cancel(); + if (this.CheckIfDisposalAppropriate()) + { + this.DoDisposeAndCompletion(); + } + else + { + // Wait for completion to actually complete, and re-throw any exceptions. + return this.Completion; + } + } + + return Task.CompletedTask; + } + private void SetState(MessageHandlerState startingOperation) { lock (this.syncObject) @@ -299,5 +323,19 @@ private bool CheckIfDisposalAppropriate(MessageHandlerState completedOperation = return shouldDispose; } } + + private void DoDisposeAndCompletion() + { + try + { + this.Dispose(true); + this.completionSource.TrySetResult(0); + } + catch (Exception ex) + { + this.completionSource.TrySetException(ex); + throw; + } + } } }