Prevent queueing of messages on connection after it will be closed

src/connection.cpp

@@ -98,13 +98,22 @@ void Connection::append_bytes(std::span<std::string_view> data_parts,
 
   std::unique_lock lock(mutex_);
 
+  // Prevent queueing messages after close has been requested
+  if (close_requested_) {
+    return;
+  }
+
   // Check if queue was empty to determine if we need to enable EPOLLOUT
   bool was_empty = message_queue_.empty();
 
+  // Set close_requested flag if this message requests close
+  if (close_after_send) {
+    close_requested_ = true;
+  }
+
   // Add message to queue
   // TODO this allocates while holding the connection lock
-  message_queue_.emplace_back(
-      Message{std::move(arena), data_parts, close_after_send});
+  message_queue_.emplace_back(Message{std::move(arena), data_parts});
   outgoing_bytes_queued_ += total_bytes;
 
   // If queue was empty, we need to add EPOLLOUT interest.
@@ -128,6 +137,11 @@ void Connection::send_response(void *protocol_context,
                                std::string_view response_json, Arena arena) {
   std::unique_lock lock(mutex_);
 
+  // Prevent queueing responses after close has been requested
+  if (close_requested_) {
+    return;
+  }
+
   // Store response in queue for protocol handler processing
   pending_response_queue_.emplace_back(
       PendingResponse{protocol_context, response_json, std::move(arena)});
@@ -206,10 +220,6 @@ uint32_t Connection::write_bytes() {
               part.size()};
           iov_count++;
         }
-
-        if (message.close_after_send) {
-          break;
-        }
       }
 
       if (iov_count == 0)
@@ -272,9 +282,6 @@ uint32_t Connection::write_bytes() {
           }
         }
 
-        if (front_message.close_after_send) {
-          result |= Close;
-        }
         // Move arena to thread-local vector for deferred cleanup
         g_arenas_to_free.emplace_back(std::move(front_message.arena));
         message_queue_.pop_front();
@@ -300,6 +307,9 @@ uint32_t Connection::write_bytes() {
         // sets it and we would hang.
         epoll_ctl(server->epoll_fds_[epoll_index_], EPOLL_CTL_MOD, fd_, &event);
       }
+      if (close_requested_) {
+        result |= Close;
+      }
     }
   }