From d336a3ea9e581372e225ee64113fe7fd7e080d45 Mon Sep 17 00:00:00 2001
From: Mattes D <github@xoft.cz>
Date: Sat, 14 Feb 2015 13:55:54 +0100
Subject: Fixed TCP link shutdown.

The shutdown is postponed until there's no more outgoing data in the LibEvent buffers.
---
 src/OSSupport/TCPLinkImpl.cpp | 65 +++++++++++++++++++++++++++++++++++++------
 1 file changed, 56 insertions(+), 9 deletions(-)

(limited to 'src/OSSupport/TCPLinkImpl.cpp')

diff --git a/src/OSSupport/TCPLinkImpl.cpp b/src/OSSupport/TCPLinkImpl.cpp
index 88fb57838..c6f1978ad 100644
--- a/src/OSSupport/TCPLinkImpl.cpp
+++ b/src/OSSupport/TCPLinkImpl.cpp
@@ -7,6 +7,7 @@
 #include "TCPLinkImpl.h"
 #include "NetworkSingleton.h"
 #include "ServerHandleImpl.h"
+#include "event2/buffer.h"
 
 
 
@@ -17,7 +18,10 @@
 
 cTCPLinkImpl::cTCPLinkImpl(cTCPLink::cCallbacksPtr a_LinkCallbacks):
 	super(a_LinkCallbacks),
-	m_BufferEvent(bufferevent_socket_new(cNetworkSingleton::Get().GetEventBase(), -1, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE))
+	m_BufferEvent(bufferevent_socket_new(cNetworkSingleton::Get().GetEventBase(), -1, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE)),
+	m_LocalPort(0),
+	m_RemotePort(0),
+	m_ShouldShutdown(false)
 {
 	LOGD("Created new cTCPLinkImpl at %p with BufferEvent at %p", this, m_BufferEvent);
 }
@@ -29,7 +33,10 @@ cTCPLinkImpl::cTCPLinkImpl(cTCPLink::cCallbacksPtr a_LinkCallbacks):
 cTCPLinkImpl::cTCPLinkImpl(evutil_socket_t a_Socket, cTCPLink::cCallbacksPtr a_LinkCallbacks, cServerHandleImplPtr a_Server, const sockaddr * a_Address, socklen_t a_AddrLen):
 	super(a_LinkCallbacks),
 	m_BufferEvent(bufferevent_socket_new(cNetworkSingleton::Get().GetEventBase(), a_Socket, BEV_OPT_CLOSE_ON_FREE | BEV_OPT_THREADSAFE)),
-	m_Server(a_Server)
+	m_Server(a_Server),
+	m_LocalPort(0),
+	m_RemotePort(0),
+	m_ShouldShutdown(false)
 {
 	LOGD("Created new cTCPLinkImpl at %p with BufferEvent at %p", this, m_BufferEvent);
 
@@ -111,7 +118,7 @@ void cTCPLinkImpl::Enable(cTCPLinkImplPtr a_Self)
 	m_Self = a_Self;
 
 	// Set the LibEvent callbacks and enable processing:
-	bufferevent_setcb(m_BufferEvent, ReadCallback, nullptr, EventCallback, this);
+	bufferevent_setcb(m_BufferEvent, ReadCallback, WriteCallback, EventCallback, this);
 	bufferevent_enable(m_BufferEvent, EV_READ | EV_WRITE);
 }
 
@@ -121,6 +128,11 @@ void cTCPLinkImpl::Enable(cTCPLinkImplPtr a_Self)
 
 bool cTCPLinkImpl::Send(const void * a_Data, size_t a_Length)
 {
+	if (m_ShouldShutdown)
+	{
+		LOGD("%s: Cannot send data, the link is already shut down.", __FUNCTION__);
+		return false;
+	}
 	return (bufferevent_write(m_BufferEvent, a_Data, a_Length) == 0);
 }
 
@@ -130,12 +142,15 @@ bool cTCPLinkImpl::Send(const void * a_Data, size_t a_Length)
 
 void cTCPLinkImpl::Shutdown(void)
 {
-	#ifdef _WIN32
-		shutdown(bufferevent_getfd(m_BufferEvent), SD_SEND);
-	#else
-		shutdown(bufferevent_getfd(m_BufferEvent), SHUT_WR);
-	#endif
-	bufferevent_disable(m_BufferEvent, EV_WRITE);
+	// If there's no outgoing data, shutdown the socket directly:
+	if (evbuffer_get_length(bufferevent_get_output(m_BufferEvent)) == 0)
+	{
+		DoActualShutdown();
+		return;
+	}
+
+	// There's still outgoing data in the LibEvent buffer, schedule a shutdown when it's written to OS's TCP stack:
+	m_ShouldShutdown = true;
 }
 
 
@@ -181,6 +196,24 @@ void cTCPLinkImpl::ReadCallback(bufferevent * a_BufferEvent, void * a_Self)
 
 
 
+void cTCPLinkImpl::WriteCallback(bufferevent * a_BufferEvent, void * a_Self)
+{
+	ASSERT(a_Self != nullptr);
+	auto Self = static_cast<cTCPLinkImpl *>(a_Self);
+	ASSERT(Self->m_Callbacks != nullptr);
+
+	// If there's no more data to write and the link has been scheduled for shutdown, do the shutdown:
+	auto OutLen = evbuffer_get_length(bufferevent_get_output(Self->m_BufferEvent));
+	if ((OutLen == 0) && (Self->m_ShouldShutdown))
+	{
+		Self->DoActualShutdown();
+	}
+}
+
+
+
+
+
 void cTCPLinkImpl::EventCallback(bufferevent * a_BufferEvent, short a_What, void * a_Self)
 {
 	LOGD("cTCPLink event callback for link %p, BEV %p; what = 0x%02x", a_Self, a_BufferEvent, a_What);
@@ -316,6 +349,20 @@ void cTCPLinkImpl::UpdateRemoteAddress(void)
 
 
 
+void cTCPLinkImpl::DoActualShutdown(void)
+{
+	#ifdef _WIN32
+		shutdown(bufferevent_getfd(m_BufferEvent), SD_SEND);
+	#else
+		shutdown(bufferevent_getfd(m_BufferEvent), SHUT_WR);
+	#endif
+	bufferevent_disable(m_BufferEvent, EV_WRITE);
+}
+
+
+
+
+
 ////////////////////////////////////////////////////////////////////////////////
 // cNetwork API:
 
-- 
cgit v1.2.3