diff --git a/src/HotChocolate/AspNetCore/src/Transport.Sockets.Client/Protocols/GraphQLOverWebSocket/GraphQLOverWebSocketProtocolHandler.cs b/src/HotChocolate/AspNetCore/src/Transport.Sockets.Client/Protocols/GraphQLOverWebSocket/GraphQLOverWebSocketProtocolHandler.cs index d00e05849a5..aee28644293 100644 --- a/src/HotChocolate/AspNetCore/src/Transport.Sockets.Client/Protocols/GraphQLOverWebSocket/GraphQLOverWebSocketProtocolHandler.cs +++ b/src/HotChocolate/AspNetCore/src/Transport.Sockets.Client/Protocols/GraphQLOverWebSocket/GraphQLOverWebSocketProtocolHandler.cs @@ -41,14 +41,15 @@ public async ValueTask ExecuteAsync( // if the user cancels this stream, we will send the server a complete request // so that we no longer receive new result messages. - cancellationToken.Register(completion.TrySendCompleteMessage); + var cancellationRegistration = cancellationToken.Register(completion.TrySendCompleteMessage); try { - return new SocketResult(observer, subscription, completion); + return new SocketResult(observer, subscription, completion, cancellationRegistration); } catch { + cancellationRegistration.Dispose(); subscription.Dispose(); observer.Dispose(); throw; diff --git a/src/HotChocolate/AspNetCore/src/Transport.Sockets.Client/SocketResult.cs b/src/HotChocolate/AspNetCore/src/Transport.Sockets.Client/SocketResult.cs index 4e05ef64c5c..0dc5f38c3b9 100644 --- a/src/HotChocolate/AspNetCore/src/Transport.Sockets.Client/SocketResult.cs +++ b/src/HotChocolate/AspNetCore/src/Transport.Sockets.Client/SocketResult.cs @@ -14,12 +14,13 @@ public sealed class SocketResult : IDisposable internal SocketResult( DataMessageObserver observer, IDisposable subscription, - IDataCompletion completion) + IDataCompletion completion, + CancellationTokenRegistration cancellationRegistration) { ArgumentNullException.ThrowIfNull(observer); ArgumentNullException.ThrowIfNull(subscription); - _enumerable = new ResultEnumerable(observer, subscription, completion); + _enumerable = new ResultEnumerable(observer, subscription, completion, cancellationRegistration); } /// @@ -44,7 +45,8 @@ public void Dispose() private sealed class ResultEnumerable( DataMessageObserver observer, IDisposable subscription, - IDataCompletion completion) + IDataCompletion completion, + CancellationTokenRegistration cancellationRegistration) : IAsyncEnumerable, IDisposable { private bool _started; @@ -89,6 +91,7 @@ public async IAsyncEnumerator GetAsyncEnumerator( public void Dispose() { completion.TrySendCompleteMessage(); + cancellationRegistration.Dispose(); subscription.Dispose(); observer.Dispose(); }