From e56cf41a0135e77f56e936094dcc9cf725ace6e9 Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Sun, 24 Aug 2025 11:04:26 -0400
Subject: [PATCH] Maintain outgoing bytes queued count

---
 src/connection.cpp |  2 ++
 src/connection.hpp | 22 ++++++++++++++++------
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/src/connection.cpp b/src/connection.cpp
index 31e4ecc..b2353ae 100644
--- a/src/connection.cpp
+++ b/src/connection.cpp
@@ -43,6 +43,7 @@ void Connection::appendMessage(std::string_view s, bool copy_to_arena) {
   } else {
     messages_.push_back(s);
   }
+  outgoing_bytes_queued_ += s.size();
 }
 
 int Connection::readBytes(char *buf, size_t buffer_size) {
@@ -104,6 +105,7 @@ bool Connection::writeBytes() {
 
     // Handle partial writes by updating string_view data/size
     size_t bytes_written = static_cast<size_t>(w);
+    outgoing_bytes_queued_ -= bytes_written;
     while (bytes_written > 0 && !messages_.empty()) {
       auto &front = messages_.front();
 
diff --git a/src/connection.hpp b/src/connection.hpp
index 7c7c8b6..d88bdf4 100644
--- a/src/connection.hpp
+++ b/src/connection.hpp
@@ -178,7 +178,7 @@ struct Connection {
   /**
    * @brief Get the number of bytes queued for transmission.
    *
-   * Calculates and returns the total number of bytes in all messages currently
+   * Returns the total number of bytes in all messages currently
    * queued for transmission to the client. This includes all data added via
    * appendMessage() that has not yet been sent over the network.
    *
@@ -187,8 +187,8 @@ struct Connection {
    * @warning Thread Safety: Only call from the thread that currently owns this
    * connection. Concurrent access to the message queue is not thread-safe.
    *
-   * @note Performance: This method iterates through all queued messages to
-   * calculate the total, so avoid calling it frequently in hot paths.
+   * @note Performance: This method uses an O(1) counter for fast retrieval
+   * in release builds. In debug builds, validates counter accuracy.
    *
    * @note The count decreases as the server sends data via writeBytes() and
    * removes completed messages from the queue.
@@ -210,11 +210,17 @@ struct Connection {
    * ```
    */
   int64_t outgoingBytesQueued() const {
-    int64_t result = 0;
+#ifndef NDEBUG
+    // Debug build: validate counter accuracy
+    int64_t computed_total = 0;
     for (auto s : messages_) {
-      result += s.size();
+      computed_total += s.size();
     }
-    return result;
+    assert(
+        outgoing_bytes_queued_ == computed_total &&
+        "outgoing_bytes_queued_ counter is out of sync with actual queue size");
+#endif
+    return outgoing_bytes_queued_;
   }
 
   /**
@@ -282,6 +288,7 @@ struct Connection {
    */
   void reset() {
     assert(messages_.empty());
+    outgoing_bytes_queued_ = 0;
     arena_.reset();
     messages_ =
         std::deque<std::string_view, ArenaStlAllocator<std::string_view>>{
@@ -345,6 +352,9 @@ private:
   std::deque<std::string_view, ArenaStlAllocator<std::string_view>> messages_{
       ArenaStlAllocator<std::string_view>{&arena_}};
 
+  // Counter tracking total bytes queued for transmission
+  int64_t outgoing_bytes_queued_{0};
+
   // Whether or not to close the connection after completing writing the
   // response
   bool closeConnection_{false};