Add on_write_progress

src/connection.cpp

@@ -122,7 +122,7 @@ bool Connection::writeBytes() {
   // This reclaims memory from request parsing and response generation
   arena_.reset();
 
-  return closeConnection_;
+  return false;
 }
 
 void Connection::tsan_acquire() {

src/connection.hpp

@@ -44,6 +44,13 @@ struct Connection {
   void closeAfterSend() { closeConnection_ = true; }
   ArenaAllocator &getArena() { return arena_; }
   int64_t getId() const { return id_; }
+  int64_t outgoingBytesQueued() const {
+    int64_t result = 0;
+    for (auto s : messages_) {
+      result += s.size();
+    }
+    return result;
+  }
 
   // Note: To release connection back to server, use
   // Server::releaseBackToServer(std::move(connection_ptr))

src/connection_handler.hpp

@@ -35,11 +35,17 @@ public:
    * - If ownership is taken, handler must call Server::releaseBackToServer()
    * when done
    * @note `data` is *not* owned by the connection arena, and its lifetime ends
-   * after the call to process_data.
+   * after the call to on_data_arrived.
    * @note May be called from an arbitrary network thread.
    */
-  virtual void process_data(std::string_view data,
-                            std::unique_ptr<Connection> &conn_ptr) = 0;
+  virtual void on_data_arrived(std::string_view data,
+                               std::unique_ptr<Connection> &conn_ptr) = 0;
+
+  /**
+   * Successfully wrote data on the connection.
+   * @note May be called from an arbitrary network thread.
+   */
+  virtual void on_write_progress(std::unique_ptr<Connection> &) {}
 
   /**
    * Called when a new connection is established.

src/main.cpp

@@ -29,8 +29,8 @@ void signal_handler(int sig) {
  */
 class EchoHandler : public ConnectionHandler {
 public:
-  void process_data(std::string_view data,
-                    std::unique_ptr<Connection> &conn_ptr) override {
+  void on_data_arrived(std::string_view data,
+                       std::unique_ptr<Connection> &conn_ptr) override {
     // Echo the received data back to the client
     conn_ptr->appendMessage(data);
   }

src/server.cpp

@@ -270,7 +270,7 @@ void Server::start_network_threads() {
 
             // Call handler with unique_ptr - handler can take ownership if
             // needed
-            handler_.process_data(data, conn);
+            handler_.on_data_arrived(data, conn);
 
             // If handler took ownership (conn is now null), don't continue
             // processing
@@ -280,8 +280,22 @@ void Server::start_network_threads() {
           }
 
           if (events[i].events & EPOLLOUT) {
-            bool done = conn->writeBytes();
-            if (done) {
+            bool error = conn->writeBytes();
+            if (error) {
+              continue;
+            }
+
+            // Call handler with unique_ptr - handler can take ownership if
+            // needed
+            handler_.on_write_progress(conn);
+            // If handler took ownership (conn is now null), don't continue
+            // processing
+            if (!conn) {
+              continue;
+            }
+
+            // Check if we should close the connection according to application
+            if (!conn->hasMessages() && conn->closeConnection_) {
               continue;
             }
           }