Improve disposal logic in AsyncWebsocketMessageResultEnumerator to prevent multiple disposals (#527)

jsquire · christothes · stephentoub · web-flow · commit 1e5da07555da · 2025-07-10T12:21:21.000-07:00
* Improve disposal logic in AsyncWebsocketMessageResultEnumerator to prevent multiple disposals (#476) * Improve disposal logic in AsyncWebsocketMessageResultEnumerator to prevent multiple disposals * fb * Update src/Custom/RealtimeConversation/Internal/AsyncWebsocketMessageEnumerator.cs Co-authored-by: Stephen Toub <stoub@microsoft.com> --------- Co-authored-by: Stephen Toub <stoub@microsoft.com> * Update src/Custom/Realtime/Internal/AsyncWebsocketMessageEnumerator.cs --------- Co-authored-by: Christopher Scott <chriscott@hotmail.com> Co-authored-by: Stephen Toub <stoub@microsoft.com>
diff --git a/src/Custom/Realtime/Internal/AsyncWebsocketMessageEnumerator.cs b/src/Custom/Realtime/Internal/AsyncWebsocketMessageEnumerator.cs
@@ -13,7 +13,7 @@ internal partial class AsyncWebsocketMessageResultEnumerator : IAsyncEnumerator<
     public ClientResult Current { get; private set; }
     private readonly CancellationToken _cancellationToken;
     private readonly WebSocket _webSocket;
-    private readonly byte[] _receiveBuffer;
+    private byte[] _receiveBuffer;
 
     public AsyncWebsocketMessageResultEnumerator(WebSocket webSocket, CancellationToken cancellationToken)
     {
@@ -26,12 +26,11 @@ public AsyncWebsocketMessageResultEnumerator(WebSocket webSocket, CancellationTo
 
     public ValueTask DisposeAsync()
     {
-        _webSocket?.Dispose();
-        if (_receiveBuffer is not null)
+        if (Interlocked.Exchange(ref _receiveBuffer, null) is byte[] toReturn)
         {
-            ArrayPool<byte>.Shared.Return(_receiveBuffer);
+            ArrayPool<byte>.Shared.Return(toReturn);
         }
-        return new ValueTask(Task.CompletedTask);
+        return default;
     }
 
     public async ValueTask<bool> MoveNextAsync()

Original file line number	Diff line number	Diff line change
`@@ -13,7 +13,7 @@ internal partial class AsyncWebsocketMessageResultEnumerator : IAsyncEnumerator<`
`13`	`13`	`public ClientResult Current { get; private set; }`
`14`	`14`	`private readonly CancellationToken _cancellationToken;`
`15`	`15`	`private readonly WebSocket _webSocket;`
`16`		`- private readonly byte[] _receiveBuffer;`
	`16`	`+ private byte[] _receiveBuffer;`
`17`	`17`
`18`	`18`	`public AsyncWebsocketMessageResultEnumerator(WebSocket webSocket, CancellationToken cancellationToken)`
`19`	`19`	`{`
`@@ -26,12 +26,11 @@ public AsyncWebsocketMessageResultEnumerator(WebSocket webSocket, CancellationTo`
`26`	`26`
`27`	`27`	`public ValueTask DisposeAsync()`
`28`	`28`	`{`
`29`		`- _webSocket?.Dispose();`
`30`		`- if (_receiveBuffer is not null)`
	`29`	`+ if (Interlocked.Exchange(ref _receiveBuffer, null) is byte[] toReturn)`
`31`	`30`	`{`
`32`		`- ArrayPool<byte>.Shared.Return(_receiveBuffer);`
	`31`	`+ ArrayPool<byte>.Shared.Return(toReturn);`
`33`	`32`	`}`
`34`		`- return new ValueTask(Task.CompletedTask);`
	`33`	`+ return default;`
`35`	`34`	`}`
`36`	`35`
`37`	`36`	`public async ValueTask<bool> MoveNextAsync()`