From 35954dc41be1f974f71da56579ac0c4828c4619e Mon Sep 17 00:00:00 2001
From: Ivan Savenko <saven.ivan@gmail.com>
Date: Tue, 14 May 2024 19:40:20 +0000
Subject: [PATCH] Simple workaround to fix vcmiserver shutdown procedure

At the moment, vcmilobby *requires* async writes in order to handle
multiple connections with different speeds and at optimal performance,
without hanging if one player is too slow and can't eat all data server
sent to him at once.

However server (and potentially - client) can not handle this mode and
may shutdown either socket or entire asio service too early, before all
writes are performed, leading to weird freeze on ending scenario where
client would not receive notifications about end of game.
---
 lib/network/NetworkConnection.cpp | 32 ++++++++++++++++++++++++-------
 lib/network/NetworkConnection.h   |  4 +++-
 lib/network/NetworkInterface.h    |  1 +
 lobby/LobbyServer.cpp             |  1 +
 4 files changed, 30 insertions(+), 8 deletions(-)

diff --git a/lib/network/NetworkConnection.cpp b/lib/network/NetworkConnection.cpp
index cef7f8263..f7593e75c 100644
--- a/lib/network/NetworkConnection.cpp
+++ b/lib/network/NetworkConnection.cpp
@@ -127,6 +127,11 @@ void NetworkConnection::onPacketReceived(const boost::system::error_code & ec, u
 	startReceiving();
 }
 
+void NetworkConnection::setAsyncWritesEnabled(bool on)
+{
+	asyncWritesEnabled = on;
+}
+
 void NetworkConnection::sendPacket(const std::vector<std::byte> & message)
 {
 	std::lock_guard<std::mutex> lock(writeMutex);
@@ -134,14 +139,27 @@ void NetworkConnection::sendPacket(const std::vector<std::byte> & message)
 	uint32_t messageSize = message.size();
 	std::memcpy(headerVector.data(), &messageSize, sizeof(uint32_t));
 
-	bool messageQueueEmpty = dataToSend.empty();
-	dataToSend.push_back(headerVector);
-	if (message.size() > 0)
-		dataToSend.push_back(message);
+	// At the moment, vcmilobby *requires* async writes in order to handle multiple connections with different speeds and at optimal performance
+	// However server (and potentially - client) can not handle this mode and may shutdown either socket or entire asio service too early, before all writes are performed
+	if (asyncWritesEnabled)
+	{
 
-	if (messageQueueEmpty)
-		doSendData();
-	//else - data sending loop is still active and still sending previous messages
+		bool messageQueueEmpty = dataToSend.empty();
+		dataToSend.push_back(headerVector);
+		if (message.size() > 0)
+			dataToSend.push_back(message);
+
+		if (messageQueueEmpty)
+			doSendData();
+		//else - data sending loop is still active and still sending previous messages
+	}
+	else
+	{
+		boost::system::error_code ec;
+		boost::asio::write(*socket, boost::asio::buffer(headerVector), ec );
+		if (message.size() > 0)
+			boost::asio::write(*socket, boost::asio::buffer(message), ec );
+	}
 }
 
 void NetworkConnection::doSendData()
diff --git a/lib/network/NetworkConnection.h b/lib/network/NetworkConnection.h
index d6d022bbb..8ab9fcc8c 100644
--- a/lib/network/NetworkConnection.h
+++ b/lib/network/NetworkConnection.h
@@ -13,7 +13,7 @@
 
 VCMI_LIB_NAMESPACE_BEGIN
 
-class NetworkConnection : public INetworkConnection, public std::enable_shared_from_this<NetworkConnection>
+class NetworkConnection final : public INetworkConnection, public std::enable_shared_from_this<NetworkConnection>
 {
 	static const int messageHeaderSize = sizeof(uint32_t);
 	static const int messageMaxSize = 64 * 1024 * 1024; // arbitrary size to prevent potential massive allocation if we receive garbage input
@@ -25,6 +25,7 @@ class NetworkConnection : public INetworkConnection, public std::enable_shared_f
 
 	NetworkBuffer readBuffer;
 	INetworkConnectionListener & listener;
+	bool asyncWritesEnabled = false;
 
 	void heartbeat();
 	void onError(const std::string & message);
@@ -42,6 +43,7 @@ public:
 	void start();
 	void close() override;
 	void sendPacket(const std::vector<std::byte> & message) override;
+	void setAsyncWritesEnabled(bool on) override;
 };
 
 VCMI_LIB_NAMESPACE_END
diff --git a/lib/network/NetworkInterface.h b/lib/network/NetworkInterface.h
index 58bd11023..45e26bf9e 100644
--- a/lib/network/NetworkInterface.h
+++ b/lib/network/NetworkInterface.h
@@ -17,6 +17,7 @@ class DLL_LINKAGE INetworkConnection : boost::noncopyable
 public:
 	virtual ~INetworkConnection() = default;
 	virtual void sendPacket(const std::vector<std::byte> & message) = 0;
+	virtual void setAsyncWritesEnabled(bool on) = 0;
 	virtual void close() = 0;
 };
 
diff --git a/lobby/LobbyServer.cpp b/lobby/LobbyServer.cpp
index e1e756bb0..2a09e3d02 100644
--- a/lobby/LobbyServer.cpp
+++ b/lobby/LobbyServer.cpp
@@ -292,6 +292,7 @@ void LobbyServer::sendChatMessage(const NetworkConnectionPtr & target, const std
 
 void LobbyServer::onNewConnection(const NetworkConnectionPtr & connection)
 {
+	connection->setAsyncWritesEnabled(true);
 	// no-op - waiting for incoming data
 }