diff --git a/src/connection.cpp b/src/connection.cpp index 2004357..55550e1 100644 --- a/src/connection.cpp +++ b/src/connection.cpp @@ -89,7 +89,7 @@ void Connection::close() { // Called from I/O thread only void Connection::append_bytes(std::span data_parts, - Arena arena, bool close_after_send) { + Arena arena, ConnectionShutdown shutdown_mode) { // Calculate total bytes for this message. Don't need to hold the lock yet. size_t total_bytes = 0; for (const auto &part : data_parts) { @@ -98,17 +98,17 @@ void Connection::append_bytes(std::span data_parts, std::unique_lock lock(mutex_); - // Prevent queueing messages after close has been requested - if (close_requested_) { + // Prevent queueing messages after shutdown has been requested + if (shutdown_requested_ != ConnectionShutdown::None) { 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; + // Set shutdown mode if requested + if (shutdown_mode != ConnectionShutdown::None) { + shutdown_requested_ = shutdown_mode; } // Add message to queue @@ -137,8 +137,8 @@ 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_) { + // Prevent queueing responses after shutdown has been requested + if (shutdown_requested_ != ConnectionShutdown::None) { return; } @@ -307,7 +307,11 @@ 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_) { + // Handle shutdown modes after all messages are sent + if (shutdown_requested_ == ConnectionShutdown::WriteOnly) { + // Shutdown write side but keep connection alive for reading + shutdown(fd_, SHUT_WR); + } else if (shutdown_requested_ == ConnectionShutdown::Full) { result |= Close; } } diff --git a/src/connection.hpp b/src/connection.hpp index d57acbd..dd85ea7 100644 --- a/src/connection.hpp +++ b/src/connection.hpp @@ -20,6 +20,15 @@ // Forward declaration struct Server; +/** + * Shutdown modes for connection termination. + */ +enum class ConnectionShutdown { + None, // Normal operation - no shutdown requested + WriteOnly, // shutdown(SHUT_WR) after sending queued data + Full // close() after sending queued data +}; + /** * Base interface for sending messages to a connection. * This restricted interface is safe for use by pipeline threads, @@ -109,15 +118,14 @@ struct Connection : MessageSender { * * @param data_parts Span of string_views pointing to arena-allocated data * @param arena Arena that owns all the memory referenced by data_parts - * @param close_after_send Whether to close connection after sending all - * queued data + * @param shutdown_mode Shutdown mode to apply after sending all queued data * * @note Thread Safety: Must be called from I/O thread only. * @note Ordering: Bytes are sent in the order calls are made. * @note The memory referenced by the data_parts span, must outlive @p arena. - * @note Close Request: To request connection close without sending data, - * pass empty data_parts span with close_after_send=true. This ensures - * all previously queued messages are sent before closing. + * @note Shutdown Request: To request connection shutdown without sending + * data, pass empty data_parts span with desired shutdown_mode. This ensures + * all previously queued messages are sent before shutdown. * * Example usage (from ConnectionHandler::on_preprocess_writes): * ```cpp @@ -125,11 +133,12 @@ struct Connection : MessageSender { * auto parts = arena.allocate_span(2); * parts[0] = build_header(arena); * parts[1] = build_body(arena); - * conn.append_bytes({parts, 2}, std::move(arena), false); + * conn.append_bytes({parts, 2}, std::move(arena), ConnectionShutdown::None); * ``` */ - void append_bytes(std::span data_parts, Arena arena, - bool close_after_send); + void + append_bytes(std::span data_parts, Arena arena, + ConnectionShutdown shutdown_mode = ConnectionShutdown::None); void send_response(void *protocol_context, std::string_view response_json, Arena arena) override; @@ -346,8 +355,9 @@ private: // dropped while server accesses existing elements. int64_t outgoing_bytes_queued_{0}; // Counter of queued bytes std::deque - pending_response_queue_; // Responses awaiting protocol processing - bool close_requested_{false}; // True if close_after_send has been requested + pending_response_queue_; // Responses awaiting protocol processing + ConnectionShutdown shutdown_requested_{ + ConnectionShutdown::None}; // Shutdown mode requested // Set to a negative number in `close` int fd_; diff --git a/src/http_handler.cpp b/src/http_handler.cpp index 954f98e..c688c95 100644 --- a/src/http_handler.cpp +++ b/src/http_handler.cpp @@ -522,7 +522,9 @@ void HttpConnectionState::send_ordered_response( // Send through append_bytes which handles write interest conn.append_bytes(response_data.data, std::move(response_data.arena), - response_data.connection_close); + response_data.connection_close + ? ConnectionShutdown::Full + : ConnectionShutdown::None); next_sequence_to_send++; iter = ready_responses.erase(iter); } diff --git a/tests/test_server.cpp b/tests/test_server.cpp index f38b9a7..35739d4 100644 --- a/tests/test_server.cpp +++ b/tests/test_server.cpp @@ -39,7 +39,7 @@ TEST_CASE("Echo test") { // MessageSender) auto *conn_ptr = static_cast(conn.get()); conn_ptr->append_bytes(std::exchange(handler.reply, {}), - std::move(handler.arena), false); + std::move(handler.arena), ConnectionShutdown::None); } else { REQUIRE(false); }