From 069e66b31cd93b983f8f2be9c7416176dd8ef20a Mon Sep 17 00:00:00 2001
From: Alex Hultman <alexhultman@gmail.com>
Date: Sun, 19 Sep 2021 01:19:29 +0000
Subject: [PATCH] Optimize getSendBuffer for large sends

---
 examples/EchoServer.cpp |  6 +++---
 src/AsyncSocket.h       | 33 +++++++++++++++++++++++++--------
 src/AsyncSocketData.h   |  3 +++
 src/WebSocket.h         | 33 ++++++++-------------------------
 4 files changed, 39 insertions(+), 36 deletions(-)

diff --git a/examples/EchoServer.cpp b/examples/EchoServer.cpp
index e96c0f6..ff9c625 100644
--- a/examples/EchoServer.cpp
+++ b/examples/EchoServer.cpp
@@ -19,10 +19,10 @@ int main() {
 	    .passphrase = "1234"
 	}).ws<PerSocketData>("/*", {
         /* Settings */
-        .compression = uWS::SHARED_COMPRESSOR,
-        .maxPayloadLength = 16 * 1024 * 1024,
+        .compression = uWS::DEDICATED_COMPRESSOR_4KB,
+        .maxPayloadLength = 100 * 1024 * 1024,
         .idleTimeout = 16,
-        .maxBackpressure = 1 * 1024 * 1024,
+        .maxBackpressure = 100 * 1024 * 1024,
         .closeOnBackpressureLimit = false,
         .resetIdleTimeoutOnSend = false,
         .sendPingsAutomatically = true,
diff --git a/src/AsyncSocket.h b/src/AsyncSocket.h
index e3e2771..20f023b 100644
--- a/src/AsyncSocket.h
+++ b/src/AsyncSocket.h
@@ -34,6 +34,12 @@
 
 namespace uWS {
 
+    enum SendBufferAttribute {
+        NEEDS_NOTHING,
+        NEEDS_DRAIN,
+        NEEDS_UNCORK
+    };
+
     template <bool, bool, typename> struct WebSocketContext;
 
 template <bool SSL>
@@ -102,16 +108,27 @@ protected:
     }
 
     /* Returns a suitable buffer for temporary assemblation of send data */
-    std::pair<char *, bool> getSendBuffer(size_t size) {
-        /* If we are corked and we have room, return the cork buffer itself */
+    std::pair<char *, SendBufferAttribute> getSendBuffer(size_t size) {
+        /* First step is to determine if we already have backpressure or not */
         LoopData *loopData = getLoopData();
-        if (loopData->corkedSocket == this && loopData->corkOffset + size < LoopData::CORK_BUFFER_SIZE) {
-            char *sendBuffer = loopData->corkBuffer + loopData->corkOffset;
-            loopData->corkOffset += (unsigned int) size;
-            return {sendBuffer, false};
+        BackPressure &backPressure = getAsyncSocketData()->buffer;
+        size_t existingBackpressure = backPressure.length();
+        if ((!existingBackpressure) && (isCorked() || canCork()) && (loopData->corkOffset + size < LoopData::CORK_BUFFER_SIZE)) {
+            /* Cork automatically if we can */
+            if (isCorked()) {
+                char *sendBuffer = loopData->corkBuffer + loopData->corkOffset;
+                loopData->corkOffset += (unsigned int) size;
+                return {sendBuffer, SendBufferAttribute::NEEDS_NOTHING};
+            } else {
+                cork();
+                char *sendBuffer = loopData->corkBuffer + loopData->corkOffset;
+                loopData->corkOffset += (unsigned int) size;
+                return {sendBuffer, SendBufferAttribute::NEEDS_UNCORK};
+            }
         } else {
-            /* Slow path for now, we want to always be corked if possible */
-            return {(char *) malloc(size), true};
+            /* Fallback is to use the backpressure as buffer */
+            backPressure.resize(existingBackpressure + size);
+            return {(char *) backPressure.data() + existingBackpressure, SendBufferAttribute::NEEDS_DRAIN};
         }
     }
 
diff --git a/src/AsyncSocketData.h b/src/AsyncSocketData.h
index 16dd997..6184826 100644
--- a/src/AsyncSocketData.h
+++ b/src/AsyncSocketData.h
@@ -50,6 +50,9 @@ struct BackPressure {
     void reserve(size_t length) {
         buffer.reserve(length + pendingRemoval);
     }
+    void resize(size_t length) {
+        buffer.resize(length + pendingRemoval);
+    }
     const char *data() {
         return buffer.data() + pendingRemoval;
     }
diff --git a/src/WebSocket.h b/src/WebSocket.h
index b2181ce..6e878e0 100644
--- a/src/WebSocket.h
+++ b/src/WebSocket.h
@@ -107,38 +107,21 @@ public:
             }
         }
 
-        /* Check to see if we can cork for the user */
-        bool automaticallyCorked = false;
-        if (!Super::isCorked() && Super::canCork()) {
-            automaticallyCorked = true;
-            Super::cork();
-        }
-
-        /* Get size, alloate size, write if needed */
+        /* Get size, allocate size, write if needed */
         size_t messageFrameSize = protocol::messageFrameSize(message.length());
-        auto [sendBuffer, requiresWrite] = Super::getSendBuffer(messageFrameSize);
+        auto [sendBuffer, sendBufferAttribute] = Super::getSendBuffer(messageFrameSize);
         protocol::formatMessage<isServer>(sendBuffer, message.data(), message.length(), opCode, message.length(), compress);
-        /* This is the slow path, when we couldn't cork for the user */
-        if (requiresWrite) {
-            /* We tried corking for the user but in the end we did not even fit in the cork buffer */
-            if (automaticallyCorked) {
-                Super::uncork();
-                automaticallyCorked = false;
-            }
-
-            auto[written, failed] = Super::write(sendBuffer, (int) messageFrameSize);
-
-            /* For now, we are slow here */
-            free(sendBuffer);
 
+        /* Depending on size of message we have different paths */
+        if (sendBufferAttribute == SendBufferAttribute::NEEDS_DRAIN) {
+            /* This is a drain */
+            auto[written, failed] = Super::write(nullptr, 0);
             if (failed) {
                 /* Return false for failure, skipping to reset the timeout below */
                 return BACKPRESSURE;
             }
-        }
-
-        /* Uncork here if we automatically corked for the user */
-        if (automaticallyCorked) {
+        } else if (sendBufferAttribute == SendBufferAttribute::NEEDS_UNCORK) {
+            /* Uncork if we came here uncorked */
             auto [written, failed] = Super::uncork();
             if (failed) {
                 return BACKPRESSURE;