From 0049b60faea2b14f77a263b326a415c94cf4bad3 Mon Sep 17 00:00:00 2001
From: Dan Engelbrecht <dan.engelbrecht@epicgames.com>
Date: Thu, 11 May 2023 10:52:14 +0200
Subject: Close down http server gracefully when exiting even while requests
 are still being processed (#290)

* Close down http server gracefully when exiting even while requests are still being processed
---
 src/zenhttp/httpasio.cpp | 58 +++++++++++++++++++++++++++++++++++++++---------
 1 file changed, 48 insertions(+), 10 deletions(-)

(limited to 'src/zenhttp/httpasio.cpp')

diff --git a/src/zenhttp/httpasio.cpp b/src/zenhttp/httpasio.cpp
index 79b2c0a3d..7149caf28 100644
--- a/src/zenhttp/httpasio.cpp
+++ b/src/zenhttp/httpasio.cpp
@@ -329,7 +329,7 @@ private:
 	void EnqueueRead();
 	void OnDataReceived(const asio::error_code& Ec, std::size_t ByteCount);
 	void OnResponseDataSent(const asio::error_code& Ec, std::size_t ByteCount, bool Pop = false);
-	void OnError();
+	void CloseConnection();
 
 	HttpAsioServerImpl&					   m_Server;
 	asio::streambuf						   m_RequestBuffer;
@@ -368,9 +368,17 @@ HttpServerConnection::HandleNewRequest()
 void
 HttpServerConnection::TerminateConnection()
 {
+	if (m_RequestState == RequestState::kDone || m_RequestState == RequestState::kTerminated)
+	{
+		return;
+	}
+
 	m_RequestState = RequestState::kTerminated;
+	ZEN_ASSERT(m_Socket);
 
+	// Terminating, we don't care about any errors when closing socket
 	std::error_code Ec;
+	m_Socket->shutdown(asio::socket_base::shutdown_both, Ec);
 	m_Socket->close(Ec);
 }
 
@@ -407,7 +415,7 @@ HttpServerConnection::OnDataReceived(const asio::error_code& Ec, [[maybe_unused]
 		else
 		{
 			ZEN_WARN("on data received ERROR, connection '{}' reason '{}'", m_ConnectionId, Ec.message());
-			return OnError();
+			return TerminateConnection();
 		}
 	}
 
@@ -424,7 +432,7 @@ HttpServerConnection::OnDataReceived(const asio::error_code& Ec, [[maybe_unused]
 		size_t Result = m_RequestData.ConsumeData((const char*)InputBuffer.data(), InputBuffer.size());
 		if (Result == ~0ull)
 		{
-			return OnError();
+			return TerminateConnection();
 		}
 
 		m_RequestBuffer.consume(Result);
@@ -449,7 +457,7 @@ HttpServerConnection::OnResponseDataSent(const asio::error_code& Ec, [[maybe_unu
 	if (Ec)
 	{
 		ZEN_WARN("on data sent ERROR, connection '{}' reason '{}'", m_ConnectionId, Ec.message());
-		OnError();
+		TerminateConnection();
 	}
 	else
 	{
@@ -461,9 +469,7 @@ HttpServerConnection::OnResponseDataSent(const asio::error_code& Ec, [[maybe_unu
 
 		if (!m_RequestData.IsKeepAlive())
 		{
-			m_RequestState = RequestState::kDone;
-
-			m_Socket->close();
+			CloseConnection();
 		}
 		else
 		{
@@ -479,9 +485,26 @@ HttpServerConnection::OnResponseDataSent(const asio::error_code& Ec, [[maybe_unu
 }
 
 void
-HttpServerConnection::OnError()
+HttpServerConnection::CloseConnection()
 {
-	m_Socket->close();
+	if (m_RequestState == RequestState::kDone || m_RequestState == RequestState::kTerminated)
+	{
+		return;
+	}
+	ZEN_ASSERT(m_Socket);
+	m_RequestState = RequestState::kDone;
+
+	std::error_code Ec;
+	m_Socket->shutdown(asio::socket_base::shutdown_receive, Ec);
+	if (Ec)
+	{
+		ZEN_WARN("socket shutdown ERROR, reason '{}'", Ec.message());
+	}
+	m_Socket->close(Ec);
+	if (Ec)
+	{
+		ZEN_WARN("socket close ERROR, reason '{}'", Ec.message());
+	}
 }
 
 void
@@ -1053,7 +1076,12 @@ struct HttpAcceptor
 			}
 			else
 			{
-				m_Acceptor.close();
+				std::error_code CloseEc;
+				m_Acceptor.close(CloseEc);
+				if (CloseEc)
+				{
+					ZEN_WARN("acceptor close ERROR, reason '{}'", CloseEc.message());
+				}
 			}
 		});
 	}
@@ -1295,6 +1323,15 @@ HttpAsioServer::HttpAsioServer() : m_Impl(std::make_unique<asio_http::HttpAsioSe
 }
 
 HttpAsioServer::~HttpAsioServer()
+{
+	if (m_Impl)
+	{
+		ZEN_ERROR("~HttpAsioServer() called without calling Close() first");
+	}
+}
+
+void
+HttpAsioServer::Close()
 {
 	try
 	{
@@ -1304,6 +1341,7 @@ HttpAsioServer::~HttpAsioServer()
 	{
 		ZEN_WARN("Caught exception stopping http asio server: {}", ex.what());
 	}
+	m_Impl.reset();
 }
 
 void
-- 
cgit v1.2.3