Feature/networking issues

include/ghoul/io/socket/tcpsocket.h

@@ -73,6 +73,14 @@ class TcpSocket : public Socket {
     void interceptInput(InputInterceptor interceptor);
     void uninterceptInput();
 
+     /**
+      * Non-blocking counterpart to disconnect(), stops the I/O loops and closes the
+      * socket, but does not join the threads. Safe to call from within those threads
+      * themselves (unlike disconnect())
+      */
+    void closeConnection();
+    bool waitForOutputQueueDrained(std::chrono::milliseconds timeout);
+
     // Methods for binary communication
     template <typename T = char>
     bool get(T* buffer, size_t nItems = 1);

include/ghoul/io/socket/websocket.h

@@ -115,6 +115,7 @@ class WebSocket : public Socket {
     std::condition_variable _inputNotifier;
 
     std::unique_ptr<TcpSocket> _tcpSocket;
+    bool _isMarkedForClosing = false;
 };
 
 } // namespace ghoul::io

src/io/socket/tcpsocket.cpp

@@ -338,7 +338,7 @@ void TcpSocket::streamOutput() {
                 nBytesToSend,
                 0
             );
-            auto failed = [](ssize_t nBytes) { return nBytes == ssize_t(-1); };
+            auto failed = [](ssize_t nBytes) { return nBytes <= 0; };
 #endif // WIN32
 
             if (failed(nSentBytes)) {
@@ -348,8 +348,9 @@ void TcpSocket::streamOutput() {
                 _outputNotifier.notify_all();
                 return;
             }
-            _outputQueue.erase(_outputQueue.begin(), _outputQueue.begin() + nBytesToSend);
+            _outputQueue.erase(_outputQueue.begin(), _outputQueue.begin() + nSentBytes);
         }
+        _outputNotifier.notify_all(); // Let anyone waiting on drainage know
     }
 }
 
@@ -445,6 +446,25 @@ void TcpSocket::uninterceptInput() {
     _inputInterceptor = nullptr;
 }
 
+void TcpSocket::closeConnection() {
+    _shouldStopThreads = true;
+    closeSocket();
+    _inputNotifier.notify_all();
+    _outputNotifier.notify_all();
+}
+
+bool TcpSocket::waitForOutputQueueDrained(std::chrono::milliseconds timeout) {
+    std::unique_lock lock(_outputQueueMutex);
+    return _outputNotifier.wait_for(
+        lock,
+        timeout,
+        [this]() {
+            return _outputQueue.empty() || _shouldStopThreads ||
+                (!_isConnected && !_isConnecting);
+        }
+    );
+}
+
 bool TcpSocket::getBytes(char* buffer, size_t nItems) {
     waitForInput(nItems);
     if (_shouldStopThreads || (!_isConnected && !_isConnecting)) {

src/io/socket/websocket.cpp

@@ -58,14 +58,45 @@ WebSocket::WebSocket(std::unique_ptr<TcpSocket> socket,
 
      _tcpSocket->interceptInput(
         [this](const char* data, size_t nBytes) {
-            _socketConnection->read_some(data, nBytes);
+            std::string output;
+            {
+                // Guards _socketConnection/_outputStream against putMessage() on other
+                // threads. onOpen/onMessage/onClose run nested inside read_some(), so
+                // they must not lock _outputStreamMutex themselves
+                const std::unique_lock lock(_outputStreamMutex);
+                _socketConnection->read_some(data, nBytes);
+
+
+                // `read_some` can cause the websocketpp to generate protocol-level output
+                // e.g., a Close frame response, or a Pong reply to a Ping. Flush the
+                // responses so the client isn't left waiting
+                output = _outputStream.str();
+                if (!output.empty()) {
+                    _outputStream.str("");
+                    _outputStream.clear();
+                }
+            }
+
+            // Write outside the lock, since TcpSocket::put has its own internal locking
+            if (!output.empty()) {
+                _tcpSocket->put<char>(output.c_str(), output.size());
+            }
+            if (_isMarkedForClosing) {
+                const bool drained =
+                    _tcpSocket->waitForOutputQueueDrained(std::chrono::seconds(3));
+                if (!drained) {
+                    LWARNING("Timed out flushing final output before closing socket");
+                }
+                _tcpSocket->closeConnection();
+            }
             _inputNotifier.notify_one();
         }
     );
 }
 
 WebSocket::~WebSocket() {
     LDEBUG("Destroying socket connection");
+    _tcpSocket->uninterceptInput();
     _socketConnection->eof();
     _tcpSocket = nullptr;
 }
@@ -105,9 +136,22 @@ bool WebSocket::getMessage(std::string& message) {
 }
 
 bool WebSocket::putMessage(const std::string& message) {
-    _socketConnection->send(message);
-    _tcpSocket->put<char>(_outputStream.str().c_str(), _outputStream.str().size());
-    _outputStream.str("");
+    std::string output;
+
+    {
+        const std::unique_lock lock(_outputStreamMutex);
+        _socketConnection->send(message);
+        output = _outputStream.str();
+        if (!output.empty()) {
+            _outputStream.str("");
+            _outputStream.clear();
+        }
+
+    }
+
+    if (!output.empty()) {
+        _tcpSocket->put<char>(output.c_str(), output.size());
+    }
     return true;
 }
 
@@ -139,8 +183,6 @@ void WebSocket::onOpen(const websocketpp::connection_hdl& hdl) {
     ));
     const std::unique_lock lock(_connectionHandlesMutex);
     _connectionHandles.insert(hdl);
-    _tcpSocket->put<char>(_outputStream.str().c_str(), _outputStream.str().size());
-    _outputStream.str("");
 }
 
 void WebSocket::onClose(const websocketpp::connection_hdl& hdl) {
@@ -152,6 +194,7 @@ void WebSocket::onClose(const websocketpp::connection_hdl& hdl) {
     const std::unique_lock lock(_connectionHandlesMutex);
     _connectionHandles.erase(hdl);
     _inputNotifier.notify_one();
+    _isMarkedForClosing = true;
 }
 
 } // namespace ghoul::io