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;
+ }
+ }
}
}