From b63e619e902b8ff3ce3e156fbe7eaff05cea7031 Mon Sep 17 00:00:00 2001
From: Andrew Arnott <andrewarnott@live.com>
Date: Sun, 19 Jan 2025 18:26:29 -0700
Subject: [PATCH] Override `WebSocketStream` APM style async methods

Closes #823
---
 src/Nerdbank.Streams/WebSocketStream.cs | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/Nerdbank.Streams/WebSocketStream.cs b/src/Nerdbank.Streams/WebSocketStream.cs
index daf8d18b..f8582313 100644
--- a/src/Nerdbank.Streams/WebSocketStream.cs
+++ b/src/Nerdbank.Streams/WebSocketStream.cs
@@ -9,6 +9,7 @@ namespace Nerdbank.Streams
     using System.Threading;
     using System.Threading.Tasks;
     using Microsoft;
+    using Microsoft.VisualStudio.Threading;
 
     /// <summary>
     /// Exposes a <see cref="WebSocket"/> as a <see cref="Stream"/>.
@@ -98,7 +99,21 @@ public override Task WriteAsync(byte[] buffer, int offset, int count, Cancellati
             return this.webSocket.SendAsync(new ArraySegment<byte>(buffer, offset, count), WebSocketMessageType.Binary, true, cancellationToken);
         }
 
-#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits
+        /// <inheritdoc />
+        public override IAsyncResult BeginRead(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
+            => this.ReadAsync(buffer, offset, count, CancellationToken.None).ToApm(callback, state);
+
+        /// <inheritdoc />
+        public override IAsyncResult BeginWrite(byte[] buffer, int offset, int count, AsyncCallback? callback, object? state)
+            => this.WriteAsync(buffer, offset, count, CancellationToken.None).ToApm(callback, state);
+
+#pragma warning disable VSTHRD002 // Avoid problematic synchronous waits - This should not deadlock because the implementation always uses .ConfigureAwait(false).
+        /// <inheritdoc />
+        public override int EndRead(IAsyncResult asyncResult)
+            => ((Task<int>)asyncResult).GetAwaiter().GetResult();
+
+        public override void EndWrite(IAsyncResult asyncResult)
+            => ((Task)asyncResult).GetAwaiter().GetResult();
 
         /// <inheritdoc />
         public override int Read(byte[] buffer, int offset, int count) => this.ReadAsync(buffer, offset, count, CancellationToken.None).GetAwaiter().GetResult();