Prevent destroying ASIO server connection until callbacks complete

1 files changed, 54 insertions, 61 deletions

diff --git a/zenhttp/httpasio.cpp b/zenhttp/httpasio.cpp
index ea84e7e87..4e4646e3b 100644
--- a/zenhttp/httpasio.cpp
+++ b/zenhttp/httpasio.cpp
@@ -6,6 +6,7 @@
 #include <zenhttp/httpserver.h>
 
 #include <deque>
+#include <memory>
 #include <memory_resource>
 
 ZEN_THIRD_PARTY_INCLUDES_START
@@ -169,7 +170,7 @@ private:
 	IoBuffer				 m_BodyBuffer;
 	uint64_t				 m_BodyPosition = 0;
 	http_parser				 m_Parser;
-	char					 m_HeaderBuffer[512];
+	char					 m_HeaderBuffer[1024];
 
 	void AppendInputBytes(const char* Data, size_t Bytes);
 	void AppendCurrentHeader();
@@ -281,7 +282,7 @@ private:
 
 //////////////////////////////////////////////////////////////////////////
 
-struct HttpServerConnection
+struct HttpServerConnection : std::enable_shared_from_this<HttpServerConnection>
 {
 	HttpServerConnection(HttpAsioServerImpl& Server, std::unique_ptr<asio::ip::tcp::socket>&& Socket);
 	~HttpServerConnection();
@@ -290,6 +291,8 @@ struct HttpServerConnection
 	void TerminateConnection();
 	void HandleRequest();
 
+	std::shared_ptr<HttpServerConnection> AsSharedPtr() { return shared_from_this(); }
+
 private:
 	enum class RequestState
 	{
@@ -369,50 +372,32 @@ HttpServerConnection::EnqueueRead()
 	asio::async_read(*m_Socket.get(),
 					 m_RequestBuffer,
 					 asio::transfer_at_least(16),
-					 [this](const asio::error_code& Ec, std::size_t ByteCount) {
-						 if (Ec)
-						 {
-							 if (m_RequestState == RequestState::kDone || m_RequestState == RequestState::kInitialRead)
-							 {
-								 // Expected, just silently handle the condition
-								 //
-								 // if we get an I/O error on the boundary between two messages
-								 // it should be fine to just not say anything
-
-								 ZEN_TRACE("(expected) socket read error: conn#{} '{}'", m_ConnectionId, Ec.message());
-							 }
-							 else
-							 {
-								 ZEN_WARN("unexpected socket read error: conn#{} {}", m_ConnectionId, Ec.message());
-							 }
-
-							 delete this;
-						 }
-						 else
-						 {
-							 ZEN_TRACE("read: conn#:{} seq#:{} t:{} bytes:{}",
-									   m_ConnectionId,
-									   m_RequestCounter.load(std::memory_order_relaxed),
-									   GetCurrentThreadId(),
-									   ByteCount);
-
-							 OnDataReceived(Ec, ByteCount);
-						 }
-					 });
+					 [Conn = AsSharedPtr()](const asio::error_code& Ec, std::size_t ByteCount) { Conn->OnDataReceived(Ec, ByteCount); });
 }
 
 void
 HttpServerConnection::OnDataReceived(const asio::error_code& Ec, std::size_t ByteCount)
 {
-	ZEN_UNUSED(ByteCount);
-
 	if (Ec)
 	{
-		ZEN_ERROR("OnDataReceived Error: {}", Ec.message());
-
-		return OnError();
+		if (m_RequestState == RequestState::kDone || m_RequestState == RequestState::kInitialRead)
+		{
+			ZEN_TRACE("on data received ERROR (EXPECTED), connection '{}' reason '{}'", m_ConnectionId, Ec.message());
+			return;
+		}
+		else
+		{
+			ZEN_ERROR("on data received ERROR, connection '{}' reason '{}'", m_ConnectionId, Ec.message());
+			return OnError();
+		}
 	}
 
+	ZEN_TRACE("on data received, connection '{}', request '{}', thread '{}', bytes '{}'",
+			  m_ConnectionId,
+			  m_RequestCounter.load(std::memory_order_relaxed),
+			  GetCurrentThreadId(),
+			  NiceBytes(ByteCount));
+
 	while (m_RequestBuffer.size())
 	{
 		const asio::const_buffer& InputBuffer = m_RequestBuffer.data();
@@ -438,24 +423,34 @@ HttpServerConnection::OnDataReceived(const asio::error_code& Ec, std::size_t Byt
 void
 HttpServerConnection::OnResponseDataSent(const asio::error_code& Ec, std::size_t ByteCount, bool Pop)
 {
-	ZEN_UNUSED(ByteCount);
 	if (Ec)
 	{
-		ZEN_ERROR("OnResponseDataSent Error: {}", Ec.message());
-		return OnError();
+		ZEN_ERROR("on data sent ERROR, connection '{}' reason '{}'", m_ConnectionId, Ec.message());
+		OnError();
 	}
 	else
 	{
+		ZEN_TRACE("on data sent, connection '{}', request '{}', thread '{}', bytes '{}'",
+				  m_ConnectionId,
+				  m_RequestCounter.load(std::memory_order_relaxed),
+				  GetCurrentThreadId(),
+				  NiceBytes(ByteCount));
+
 		if (!m_RequestData.IsKeepAlive())
 		{
 			m_RequestState = RequestState::kDone;
 
 			m_Socket->close();
 		}
-		else if (Pop)
+		else
 		{
-			RwLock::ExclusiveLockScope _(m_ResponsesLock);
-			m_Responses.pop_front();
+			if (Pop)
+			{
+				RwLock::ExclusiveLockScope _(m_ResponsesLock);
+				m_Responses.pop_front();
+			}
+
+			m_RequestCounter.fetch_add(1);
 		}
 	}
 }
@@ -478,7 +473,7 @@ HttpServerConnection::HandleRequest()
 
 		if (Ec)
 		{
-			ZEN_WARN("socket shutdown reported error: {}", Ec.message());
+			ZEN_WARN("socket shutdown ERROR, reason '{}'", Ec.message());
 		}
 	}
 	else
@@ -486,14 +481,11 @@ HttpServerConnection::HandleRequest()
 		m_RequestState = RequestState::kWriting;
 	}
 
-	const uint32_t RequestNum = m_RequestCounter.load(std::memory_order_relaxed);
-	m_RequestCounter.fetch_add(1);
-
 	if (HttpService* Service = m_Server.RouteRequest(m_RequestData.Url()))
 	{
 		HttpAsioServerRequest Request(m_RequestData, *Service, m_RequestData.Body());
 
-		ZEN_TRACE("Handling request: Conn#{} Req#{}", m_ConnectionId, RequestNum);
+		ZEN_TRACE("handle request, connection '{}' request '{}'", m_ConnectionId, m_RequestCounter.load(std::memory_order_relaxed));
 
 		Service->HandleRequest(Request);
 
@@ -525,9 +517,8 @@ HttpServerConnection::HandleRequest()
 			asio::async_write(*m_Socket.get(),
 							  ResponseBuffers,
 							  asio::transfer_exactly(ResponseLength),
-							  [this, RequestNum](const asio::error_code& Ec, std::size_t ByteCount) {
-								  ZEN_TRACE("Response sent: Conn#{} Req#{} ({})", m_ConnectionId, RequestNum, NiceBytes(ByteCount));
-								  OnResponseDataSent(Ec, ByteCount, true);
+							  [Conn = AsSharedPtr()](const asio::error_code& Ec, std::size_t ByteCount) {
+								  Conn->OnResponseDataSent(Ec, ByteCount, true);
 							  });
 
 			return;
@@ -548,9 +539,10 @@ HttpServerConnection::HandleRequest()
 				"\r\n"sv;
 		}
 
-		asio::async_write(*m_Socket.get(), asio::buffer(Response), [this](const asio::error_code& Ec, std::size_t ByteCount) {
-			OnResponseDataSent(Ec, ByteCount);
-		});
+		asio::async_write(
+			*m_Socket.get(),
+			asio::buffer(Response),
+			[Conn = AsSharedPtr()](const asio::error_code& Ec, std::size_t ByteCount) { Conn->OnResponseDataSent(Ec, ByteCount); });
 	}
 	else
 	{
@@ -572,9 +564,10 @@ HttpServerConnection::HandleRequest()
 				"No suitable route found"sv;
 		}
 
-		asio::async_write(*m_Socket.get(), asio::buffer(Response), [this](const asio::error_code& Ec, std::size_t ByteCount) {
-			OnResponseDataSent(Ec, ByteCount);
-		});
+		asio::async_write(
+			*m_Socket.get(),
+			asio::buffer(Response),
+			[Conn = AsSharedPtr()](const asio::error_code& Ec, std::size_t ByteCount) { Conn->OnResponseDataSent(Ec, ByteCount); });
 	}
 }
 
@@ -900,14 +893,14 @@ struct HttpAcceptor
 			else
 			{
 				// New connection established, pass socket ownership into connection object
-				// and initiate request handling loop
+				// and initiate request handling loop. The connection lifetime is
+				// managed by the async read/write loop by passing the shared
+				// reference to the callbacks.
 
 				Socket->set_option(asio::ip::tcp::no_delay(true));
 
-				HttpServerConnection* Conn = new HttpServerConnection(m_Server, std::move(Socket));
+				auto Conn = std::make_shared<HttpServerConnection>(m_Server, std::move(Socket));
 				Conn->HandleNewRequest();
-
-				// note: the connection object is responsible for deleting itself
 			}
 
 			if (!m_IsStopped.load())