Use byte array in message received to avoid some unnecessary encoding (…

…#2748)
hardkoded · Aug 21, 2024 · a62b86c · a62b86c
1 parent bdd47e2
commit a62b86c
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 21 deletions.
diff --git a/lib/PuppeteerSharp/Cdp/Connection.cs b/lib/PuppeteerSharp/Cdp/Connection.cs
@@ -2,6 +2,7 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Linq;
+using System.Text;
 using System.Text.Json;
 using System.Threading;
 using System.Threading.Tasks;
@@ -280,7 +281,11 @@ private async Task ProcessMessage(MessageReceivedEventArgs e)
                     return;
                 }
 
-                _logger.LogTrace("◀ Receive {Message}", response);
+                if (_logger.IsEnabled(LogLevel.Trace))
+                {
+                    _logger.LogTrace("◀ Receive {Message}", Encoding.UTF8.GetString(response));
+                }
+
                 ProcessIncomingMessage(obj);
             }
             catch (Exception ex)

diff --git a/lib/PuppeteerSharp/Transport/MessageReceivedEventArgs.cs b/lib/PuppeteerSharp/Transport/MessageReceivedEventArgs.cs
@@ -12,11 +12,11 @@ public class MessageReceivedEventArgs : EventArgs
         /// Initializes a new instance of the <see cref="PuppeteerSharp.Transport.MessageReceivedEventArgs"/> class.
         /// </summary>
         /// <param name="message">Message.</param>
-        public MessageReceivedEventArgs(string message) => Message = message;
+        public MessageReceivedEventArgs(byte[] message) => Message = message;
 
         /// <summary>
         /// Transport message.
         /// </summary>
-        public string Message { get; }
+        public byte[] Message { get; }
     }
 }
diff --git a/lib/PuppeteerSharp/Transport/WebSocketTransport.cs b/lib/PuppeteerSharp/Transport/WebSocketTransport.cs
@@ -1,5 +1,6 @@
 using System;
 using System.Diagnostics.CodeAnalysis;
+using System.IO;
 using System.Net.WebSockets;
 using System.Text;
 using System.Threading;
@@ -138,18 +139,16 @@ private static void ScheduleTransportTask(Func<CancellationToken, Task> taskFact
         /// Starts listening the socket.
         /// </summary>
         /// <returns>The start.</returns>
-        private async Task<object> GetResponseAsync(CancellationToken cancellationToken)
+        private async Task GetResponseAsync(CancellationToken cancellationToken)
         {
             var buffer = new byte[2048];
 
             while (!IsClosed)
             {
-                var endOfMessage = false;
-                var response = new StringBuilder();
-
-                while (!endOfMessage)
+                MemoryStream memoryStream = null;
+                WebSocketReceiveResult result;
+                do
                 {
-                    WebSocketReceiveResult result;
                     try
                     {
                         result = await _client.ReceiveAsync(
@@ -158,31 +157,35 @@ private async Task<object> GetResponseAsync(CancellationToken cancellationToken)
                     }
                     catch (OperationCanceledException)
                     {
-                        return null;
+                        return;
                     }
                     catch (Exception ex)
                     {
                         OnClose(ex.Message);
-                        return null;
+                        return;
                     }
 
-                    endOfMessage = result.EndOfMessage;
-
-                    if (result.MessageType == WebSocketMessageType.Text)
+                    if (result.MessageType == WebSocketMessageType.Close)
                     {
-                        response.Append(Encoding.UTF8.GetString(buffer, 0, result.Count));
+                        OnClose("WebSocket closed");
+                        return;
                     }
-                    else if (result.MessageType == WebSocketMessageType.Close)
+                    else if (result.MessageType == WebSocketMessageType.Binary)
                     {
-                        OnClose("WebSocket closed");
-                        return null;
+                        continue;
                     }
+
+                    if (memoryStream is null && !result.EndOfMessage)
+                    {
+                        memoryStream = new MemoryStream(buffer.Length);
+                    }
+
+                    memoryStream?.Write(buffer, 0, result.Count);
                 }
+                while (!result.EndOfMessage);
 
-                MessageReceived?.Invoke(this, new MessageReceivedEventArgs(response.ToString()));
+                MessageReceived?.Invoke(this, new MessageReceivedEventArgs(memoryStream is null ? buffer.AsSpan(0, result.Count).ToArray() : memoryStream.ToArray()));
             }
-
-            return null;
         }
 
         private void OnClose(string closeReason)