Add optional chunked response writes with tryWrite

.gitmodules

@@ -1,6 +1,6 @@
 [submodule "uSockets"]
 	path = uSockets
-	url = https://github.com/uNetworking/uSockets.git
+	url = https://github.com/GetThatCookie/uSockets.git
 [submodule "fuzzing/libEpollFuzzer"]
 	path = fuzzing/libEpollFuzzer
 	url = https://github.com/uNetworking/libEpollFuzzer

build.c

@@ -9,7 +9,7 @@ int main(int argc, char **argv) {
     char *CXX = strncpy(calloc(1024, 1), or_else(getenv("CXX"), "g++"), 1024);
     char *EXEC_SUFFIX = strncpy(calloc(1024, 1), maybe(getenv("EXEC_SUFFIX")), 1024);
 
-    char *EXAMPLE_FILES[] = {"HttpCache", "Precompress", "EchoBody", "HelloWorldThreaded", "Http3Server", "Broadcast", "HelloWorld", "Crc32", "ServerName",
+    char *EXAMPLE_FILES[] = {"HttpCache", "Precompress", "EchoBody", "HelloWorldThreaded", "Http3Server", "Broadcast", "HelloWorld", "ChunkedResponse", "Crc32", "ServerName",
     "EchoServer", "BroadcastingEchoServer", "UpgradeSync", "UpgradeAsync", "ParameterRoutes"};
 
     strcat(CXXFLAGS, " -march=native -O3 -Wpedantic -Wall -Wextra -Wsign-conversion -Wconversion -std=c++2b -Isrc -IuSockets/src");

examples/ChunkedResponse.cpp

@@ -0,0 +1,26 @@
+#include "App.h"
+
+/* Stream a chunked response with bounded backpressure and no known total size. */
+
+const std::string payload(16 * 1024 * 1024, 'x');
+
+int main() {
+    uWS::SSLApp({
+      .key_file_name = "misc/key.pem",
+      .cert_file_name = "misc/cert.pem",
+      .passphrase = "1234"
+    }).get("/*", [](auto *res, auto * /*req*/) {
+        res->writeHeader("Content-Type", "application/octet-stream");
+        res->onAborted([]() {});
+
+        if (!res->tryWrite(payload, true)) {
+            res->onWritable([res](uintmax_t offset) {
+                return res->tryWrite(std::string_view(payload).substr((size_t) offset), true);
+            });
+        }
+    }).listen(3000, [](auto *listenSocket) {
+        if (listenSocket) {
+            std::cout << "Listening on port 3000" << std::endl;
+        }
+    }).run();
+}

src/AsyncSocket.h

@@ -24,6 +24,7 @@
  * to unsigned length for everything to/from uSockets - this would however remove the opportunity
  * to signal error with -1 (which is how the entire UNIX syscalling is built). */
 
+#include <climits>
 #include <cstring>
 #include <iostream>
 
@@ -333,6 +334,62 @@ struct AsyncSocket {
         return {length, false};
     }
 
+    /* Optional write of two buffers. Only already accepted cork data may be buffered. */
+    std::pair<int, bool> write2(const char *first, int firstLength, const char *second, int secondLength) {
+        int length = firstLength + secondLength;
+        if (us_socket_is_closed(SSL, (us_socket_t *) this)) {
+            return {length, false};
+        }
+
+        LoopData *loopData = getLoopData();
+        AsyncSocketData<SSL> *asyncSocketData = getAsyncSocketData();
+
+        if (!asyncSocketData->buffer.length()) {
+            if (isCorked() && (unsigned int) length > LoopData::CORK_BUFFER_SIZE - loopData->corkOffset) {
+                /* Flush the accepted prefix with a more-data hint, without reserving space for optional data. */
+                int corkLength = (int) loopData->corkOffset;
+                if constexpr (!SSL) {
+#ifndef LIBUS_USE_IO_URING
+                    if (corkLength && length <= INT_MAX - corkLength) {
+                        int written = us_socket_write3(0, (us_socket_t *) this, loopData->corkBuffer, corkLength, first, firstLength, second, secondLength);
+                        loopData->corkedSocket = nullptr;
+                        loopData->corkOffset = 0;
+                        if (written < corkLength) {
+                            asyncSocketData->buffer.append(loopData->corkBuffer + written, (size_t) (corkLength - written));
+                        }
+                        return {std::max(0, written - corkLength), written != corkLength + length};
+                    }
+#endif
+                }
+                loopData->corkedSocket = nullptr;
+                loopData->corkOffset = 0;
+                if (corkLength) {
+                    int written = us_socket_write(SSL, (us_socket_t *) this, loopData->corkBuffer, corkLength, 1);
+                    if (written < corkLength) {
+                        asyncSocketData->buffer.append(loopData->corkBuffer + written, (size_t) (corkLength - written));
+                        return {0, true};
+                    }
+                }
+            }
+
+            if constexpr (!SSL) {
+#ifndef LIBUS_USE_IO_URING
+                if (!isCorked()) {
+                    int written = us_socket_write2(0, (us_socket_t *) this, first, firstLength, second, secondLength);
+                    return {written, written != length};
+                }
+#endif
+            }
+        }
+
+        auto [firstWritten, failed] = write(first, firstLength, true, secondLength);
+        if (failed) {
+            return {firstWritten, true};
+        }
+        auto [secondWritten, secondFailed] = write(second, secondLength, true);
+        return {firstLength + secondWritten, secondFailed};
+    }
+
     /* Uncork this socket and flush or buffer any corked and/or passed data. It is essential to remember doing this. */
     /* It does NOT count bytes written from cork buffer (they are already accounted for in the write call responsible for its corking)! */
     std::pair<int, bool> uncork(const char *src = nullptr, int length = 0, bool optionally = false) {

src/HttpResponse.h

@@ -33,7 +33,7 @@
 
 #include "MoveOnlyFunction.h"
 
-/* todo: tryWrite is missing currently, only send smaller segments with write */
+#include <charconv>
 
 namespace uWS {
 
@@ -143,7 +143,7 @@ struct HttpResponse : public AsyncSocket<SSL> {
 
         if (httpResponseData->state & HttpResponseData<SSL>::HTTP_WRITE_CALLED) {
 
-            /* We do not have tryWrite-like functionalities, so ignore optional in this path */
+            /* Chunked tryEnd retains the non-optional semantics of end. Use tryWrite for optional chunked writes. */
 
             /* Do not allow sending 0 chunk here */
             if (data.length()) {
@@ -171,7 +171,7 @@ struct HttpResponse : public AsyncSocket<SSL> {
                 }
             }
 
-            /* tryEnd can never fail when in chunked mode, since we do not have tryWrite (yet), only write */
+            /* tryEnd can never fail when in chunked mode. */
             Super::timeout(HTTP_TIMEOUT_S);
             return true;
         } else {
@@ -500,6 +500,97 @@ struct HttpResponse : public AsyncSocket<SSL> {
         return ok;
     }
 
+    /* Try to write an HTTP/1.1 chunk without buffering unaccepted body data. Use without Content-Length.
+     * In chunked mode, offsets count body bytes accepted by tryWrite; write/end do not advance them.
+     * For each chunk, subtract its starting offset when computing the unaccepted suffix.
+     * On false, resume via onWritable with the same chunk's unaccepted suffix and finalChunk value.
+     * An empty non-final write drains accepted framing; finalChunk completes the response once the body is accepted.
+     * Do not interleave write/end/tryEnd with an incomplete chunk. */
+    bool tryWrite(std::string_view data, bool finalChunk = false) {
+        HttpResponseData<SSL> *httpResponseData = getHttpResponseData();
+        bool insideChunk = httpResponseData->state & HttpResponseData<SSL>::HTTP_WRITE_CONTINUATION_PENDING;
+
+        if (data.empty() && (!finalChunk || insideChunk)) {
+            bool ok = !Super::write(nullptr, 0, true).second && !insideChunk;
+            if (!ok) {
+                Super::timeout(HTTP_TIMEOUT_S);
+            }
+            return ok;
+        }
+
+        /* SSL retries must retain the original body write length, even inside an explicit cork. */
+        if constexpr (SSL) {
+            if (insideChunk) {
+                Super::uncork();
+            }
+        }
+        bool needsUncork = !insideChunk && !Super::isCorked() && Super::canCork();
+        if (needsUncork) {
+            Super::cork();
+        }
+        writeStatus(HTTP_200_OK);
+        ensureChunkedBodyStarted();
+
+        if (!insideChunk && data.length()) {
+            char chunkHeader[2 * sizeof(size_t) + 2];
+            char *end = std::to_chars(chunkHeader, chunkHeader + sizeof(chunkHeader) - 2, data.length(), 16).ptr;
+            *end++ = '\r';
+            *end++ = '\n';
+            /* Framing must not be optional: the caller only owns offsets into the body. */
+            Super::write(chunkHeader, (int) (end - chunkHeader), false);
+        }
+
+        const char *trailer = finalChunk ? "\r\n0\r\n\r\n" : "\r\n";
+        int trailerLength = finalChunk ? 7 : 2;
+        int writtenTrailer = 0;
+        size_t written = 0;
+        bool failed = false;
+        while (written < data.length() && !failed) {
+            /* Leave room for framing in uSockets' signed combined byte count. */
+            int length = (int) std::min<size_t>(data.length() - written, (size_t) (INT_MAX - trailerLength));
+            auto writtenFailed = Super::write2(data.data() + written, length, trailer,
+                data.length() - written == (size_t) length ? trailerLength : 0);
+            int writtenBody = std::min(writtenFailed.first, length);
+            written += (size_t) writtenBody;
+            writtenTrailer = writtenFailed.first - writtenBody;
+            failed = writtenFailed.second;
+        }
+        httpResponseData->offset += written;
+
+        bool completed = written == data.length();
+        bool ok = completed && !failed;
+        if (completed) {
+            httpResponseData->state &= ~HttpResponseData<SSL>::HTTP_WRITE_CONTINUATION_PENDING;
+            /* Finish partially written framing internally; it never advances the body offset. */
+            if (data.empty()) {
+                trailer += 2;
+                trailerLength -= 2;
+            }
+            if (writtenTrailer != trailerLength) {
+                ok = !Super::write(trailer + writtenTrailer, trailerLength - writtenTrailer, false).second && ok;
+            }
+        } else {
+            httpResponseData->state |= HttpResponseData<SSL>::HTTP_WRITE_CONTINUATION_PENDING;
+        }
+
+        if (needsUncork) {
+            ok = !Super::uncork().second && ok;
+        }
+        if (!ok || finalChunk) {
+            Super::timeout(HTTP_TIMEOUT_S);
+        }
+
+        if (finalChunk && completed) {
+            httpResponseData->markDone();
+            if (!Super::isCorked() && (httpResponseData->state & HttpResponseData<SSL>::HTTP_CONNECTION_CLOSE) && !Super::getBufferedAmount()) {
+                Super::shutdown();
+                Super::close();
+            }
+            return true;
+        }
+        return ok;
+    }
+
     /* Get the current byte write offset for this Http response */
     uintmax_t getWriteOffset() {
         HttpResponseData<SSL> *httpResponseData = getHttpResponseData();

src/HttpResponseData.h

@@ -97,7 +97,8 @@ struct HttpResponseData : AsyncSocketData<SSL>, HttpParser {
         HTTP_WRITE_CALLED = 2, // used
         HTTP_END_CALLED = 4, // used
         HTTP_RESPONSE_PENDING = 8, // used
-        HTTP_CONNECTION_CLOSE = 16 // used
+        HTTP_CONNECTION_CLOSE = 16, // used
+        HTTP_WRITE_CONTINUATION_PENDING = 32 // used
     };
 
     /* Per socket event handlers */

tests/HttpResponse.cpp

@@ -6,15 +6,16 @@
 #include <unistd.h>
 
 #include <atomic>
+#include <cassert>
 #include <cerrno>
 #include <cstring>
 #include <iostream>
 #include <string>
 #include <string_view>
 #include <thread>
 
-/* Regression tests for chunked response framing (issue #1939).
- * beginWrite() must terminate headers once; write()/end() must emit
+/* Regression tests for chunked response framing (issue #1939) and tryWrite.
+ * beginWrite() must terminate headers once; write()/tryWrite()/end() must emit
  * complete chunks (size CRLF data CRLF) without a leading extra CRLF. */
 
 static void dump(std::string_view s) {
@@ -265,8 +266,11 @@ int main() {
     uWS::Loop *loop = uWS::Loop::get();
     int port = 0;
     std::atomic<bool> backpressureWritten{false};
+    std::atomic<bool> tryWriteStarted{false};
+    std::atomic<unsigned int> writableCalls{0};
     constexpr int chunkSize = 64 * 1024;
     constexpr int maxFillChunks = 512;
+    const std::string payload = std::string(8 * 1024 * 1024, 'a') + std::string(8 * 1024 * 1024, 'b');
 
     app.get("/write-end", [](auto *res, auto * /*req*/) {
         res->onAborted([]() {});
@@ -355,6 +359,56 @@ int main() {
             res->end();
         });
         backpressureWritten.store(true);
+    }).get("/trywrite", [](auto *res, auto * /*req*/) {
+        assert(res->tryWrite("foo"));
+        assert(res->getWriteOffset() == 3);
+        assert(!res->hasResponded());
+        assert(res->tryWrite("bar", true));
+        assert(res->getWriteOffset() == 6);
+        assert(res->hasResponded());
+    }).get("/begin-trywrite", [](auto *res, auto * /*req*/) {
+        res->beginWrite();
+        assert(res->tryWrite("foo", true));
+    }).get("/empty-trywrite", [](auto *res, auto * /*req*/) {
+        assert(res->tryWrite(""));
+        assert(!res->hasResponded());
+        assert(res->tryWrite("", true));
+        assert(res->getWriteOffset() == 0);
+        assert(res->hasResponded());
+    }).get("/trywrite-end", [](auto *res, auto * /*req*/) {
+        assert(res->tryWrite("foo"));
+        res->end("bar");
+    }).get("/trywrite-backpressure/:final", [&](auto *res, auto *req) {
+        res->onAborted([]() { std::abort(); });
+        size_t chunkEnd = req->getParameter(0) == "1" ? payload.size() : payload.size() / 2;
+        auto stream = [res, &payload, chunkEnd](uintmax_t offset) mutable {
+            while (true) {
+                assert(offset <= chunkEnd);
+                bool finalChunk = chunkEnd == payload.size();
+                if (!res->tryWrite(std::string_view(payload).substr((size_t) offset, chunkEnd - (size_t) offset), finalChunk)) {
+                    assert(!res->hasResponded());
+                    return false;
+                }
+                if (finalChunk) {
+                    return true;
+                }
+                assert(res->getWriteOffset() == chunkEnd);
+                assert(!res->hasResponded());
+                chunkEnd = payload.size();
+                offset = res->getWriteOffset();
+            }
+        };
+        assert(!stream(0));
+        assert(res->getWriteOffset() < chunkEnd);
+        uintmax_t offset = res->getWriteOffset();
+        assert(!res->tryWrite(""));
+        assert(!res->tryWrite("", true));
+        assert(res->getWriteOffset() == offset && !res->hasResponded());
+        res->onWritable([stream = std::move(stream), &writableCalls](uintmax_t offset) mutable {
+            writableCalls++;
+            return stream(offset);
+        });
+        tryWriteStarted.store(true);
     }).listen(0, [&](us_listen_socket_t *listenSocket) {
         if (!listenSocket) {
             std::cerr << "Failed to listen" << std::endl;
@@ -368,7 +422,7 @@ int main() {
         return 1;
     }
 
-    std::thread client([loop, &app, port, &backpressureWritten]() {
+    std::thread client([loop, &app, port, &backpressureWritten, &tryWriteStarted, &writableCalls, &payload]() {
         expectChunked(port, "/write-end", "3\r\nfoo\r\n0\r\n\r\n");
         expectChunked(port, "/begin-write-end", "3\r\nfoo\r\n0\r\n\r\n");
         expectChunked(port, "/begin-end", "0\r\n\r\n");
@@ -378,6 +432,10 @@ int main() {
         expectChunked(port, "/write-end-data", "3\r\nfoo\r\n3\r\nbar\r\n0\r\n\r\n");
         expectChunked(port, "/empty-write", "0\r\n\r\n");
         expectChunked(port, "/hex-size", "b\r\nhello world\r\n0\r\n\r\n");
+        expectChunked(port, "/trywrite", "3\r\nfoo\r\n3\r\nbar\r\n0\r\n\r\n");
+        expectChunked(port, "/begin-trywrite", "3\r\nfoo\r\n0\r\n\r\n");
+        expectChunked(port, "/empty-trywrite", "0\r\n\r\n");
+        expectChunked(port, "/trywrite-end", "3\r\nfoo\r\n3\r\nbar\r\n0\r\n\r\n");
 
         std::string response = requestAfterBackpressure(port, "/backpressure", backpressureWritten);
         if (response.find("Transfer-Encoding: chunked\r\n") == std::string::npos) {
@@ -398,6 +456,15 @@ int main() {
         }
         std::cout << "OK /backpressure (" << decoded.size() << " decoded bytes)" << std::endl;
 
+        for (const char *path : {"/trywrite-backpressure/0", "/trywrite-backpressure/1"}) {
+            tryWriteStarted.store(false);
+            writableCalls.store(0);
+            std::string streamed = requestAfterBackpressure(port, path, tryWriteStarted);
+            assert(decodeChunked(bodyOf(streamed)) == payload);
+            assert(writableCalls.load() > 1);
+            std::cout << "OK " << path << " (" << writableCalls.load() << " writable callbacks)" << std::endl;
+        }
+
         loop->defer([&app]() {
             app.close();
         });