Accumulate headers properly

tools/load_tester.cpp

@@ -159,17 +159,17 @@ void signal_handler(int sig) {
 struct Connection {
   static const llhttp_settings_t settings;
 
-  static std::atomic<uint64_t> requestId;
+  static std::atomic<uint64_t> nextRequestId;
 
-  char buf[1024]; // Increased size for dynamic request format
+  char buf[1024];
 
   std::string_view request;
-  uint64_t currentRequestId;
+  uint64_t requestId = 0;
   void initRequest() {
-    currentRequestId = requestId.fetch_add(1, std::memory_order_relaxed);
+    requestId = nextRequestId.fetch_add(1, std::memory_order_relaxed);
     int len = snprintf(buf, sizeof(buf),
                        "GET /ok HTTP/1.1\r\nX-Request-Id: %" PRIu64 "\r\n\r\n",
-                       currentRequestId);
+                       requestId);
     if (len == -1 || len > int(sizeof(buf))) {
       abort();
     }
@@ -203,7 +203,7 @@ struct Connection {
 
   bool readBytes() {
     for (;;) {
-      char buf[1024]; // Use a reasonable default, configurable via g_config
+      char buf[64 * (1 << 10)]; // 64 KiB
       int buf_size = std::min(int(sizeof(buf)), g_config.connection_buf_size);
       int r = read(fd, buf, buf_size);
       if (r == -1) {
@@ -214,6 +214,7 @@ struct Connection {
           return false;
         }
       }
+      // printf("read: %.*s\n", r, buf);
       if (r == 0) {
         llhttp_finish(&parser);
         return true;
@@ -233,8 +234,8 @@ struct Connection {
 
   bool writeBytes() {
     for (;;) {
-      int w;
-      w = write(fd, request.data(), request.size());
+      assert(!request.empty());
+      int w = write(fd, request.data(), request.size());
       if (w == -1) {
         if (errno == EINTR) {
           continue;
@@ -247,14 +248,12 @@ struct Connection {
         return true;
       }
       assert(w != 0);
+      // printf("write: %.*s\n", w, request.data());
       request = request.substr(w, request.size() - w);
       if (request.empty()) {
         ++requestsSent;
-        TRACE_EVENT("http", "Send request",
-                    perfetto::Flow::Global(currentRequestId));
-        if (requestsSent == g_config.requests_per_connection) {
-          return true;
-        }
+        TRACE_EVENT("http", "Send request", perfetto::Flow::Global(requestId));
+        return requestsSent == g_config.requests_per_connection;
       }
     }
   }
@@ -285,14 +284,40 @@ private:
   }
 
   uint64_t responseId = 0;
+  std::string headerValueBuffer;
+  bool expectingResponseId = false;
+
+  int on_header_field(const char *data, size_t s) {
+    std::string_view field(data, s);
+    expectingResponseId = (field == "X-Response-ID");
+    if (expectingResponseId) {
+      headerValueBuffer.clear();
+    }
+    return 0;
+  }
+
   int on_header_value(const char *data, size_t s) {
-    for (int i = 0; i < int(s); ++i) {
-      responseId = responseId * 10 + data[i] - '0';
+    if (expectingResponseId) {
+      headerValueBuffer.append(data, s);
+    }
+    return 0;
+  }
+
+  int on_header_value_complete() {
+    if (expectingResponseId) {
+      responseId = 0;
+      for (char c : headerValueBuffer) {
+        if (c >= '0' && c <= '9') {
+          responseId = responseId * 10 + (c - '0');
+        }
+      }
+      expectingResponseId = false;
     }
     return 0;
   }
 
   int on_message_complete() {
+    assert(responseId == requestId);
     TRACE_EVENT("http", "Receive response", perfetto::Flow::Global(responseId));
     responseId = 0;
     ++responsesReceived;
@@ -303,13 +328,16 @@ private:
   llhttp_t parser;
 };
 
-std::atomic<uint64_t> Connection::requestId = {};
+std::atomic<uint64_t> Connection::nextRequestId = {};
 
 const llhttp_settings_t Connection::settings = []() {
   llhttp_settings_t settings;
   llhttp_settings_init(&settings);
   settings.on_message_complete = callback<&Connection::on_message_complete>;
+  settings.on_header_field = callback<&Connection::on_header_field>;
   settings.on_header_value = callback<&Connection::on_header_value>;
+  settings.on_header_value_complete =
+      callback<&Connection::on_header_value_complete>;
   return settings;
 }();
 
@@ -514,7 +542,7 @@ int main(int argc, char *argv[]) {
       pthread_setname_np(pthread_self(),
                          ("network-" + std::to_string(i)).c_str());
       while (!g_shutdown.load(std::memory_order_relaxed)) {
-        struct epoll_event events[64]; // Use a reasonable max size
+        struct epoll_event events[256]; // Use a reasonable max size
         int batch_size = std::min(int(sizeof(events) / sizeof(events[0])),
                                   g_config.event_batch_size);
         int eventCount;
@@ -528,10 +556,6 @@ int main(int argc, char *argv[]) {
             perror("epoll_wait");
             abort(); // Keep abort for critical errors like server does
           }
-          if (eventCount == 0) {
-            // Timeout - check shutdown flag
-            break;
-          }
           break;
         }
 
@@ -624,6 +648,7 @@ int main(int argc, char *argv[]) {
 
         // Add to epoll with proper events matching server pattern
         struct epoll_event event{};
+        assert(conn->hasMessages());
         event.events = EPOLLOUT | EPOLLONESHOT | EPOLLRDHUP;
         conn->tsan_release();
         Connection *raw_conn = conn.release();