diff --git a/include/ghoul/io/socket/tcpsocket.h b/include/ghoul/io/socket/tcpsocket.h index 90474293..87ad065c 100644 --- a/include/ghoul/io/socket/tcpsocket.h +++ b/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 bool get(T* buffer, size_t nItems = 1); diff --git a/include/ghoul/io/socket/websocket.h b/include/ghoul/io/socket/websocket.h index d105b1af..500b76d7 100644 --- a/include/ghoul/io/socket/websocket.h +++ b/include/ghoul/io/socket/websocket.h @@ -115,6 +115,7 @@ class WebSocket : public Socket { std::condition_variable _inputNotifier; std::unique_ptr _tcpSocket; + bool _isMarkedForClosing = false; }; } // namespace ghoul::io diff --git a/src/io/socket/tcpsocket.cpp b/src/io/socket/tcpsocket.cpp index 2322fd33..0e53b4b6 100644 --- a/src/io/socket/tcpsocket.cpp +++ b/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)) { diff --git a/src/io/socket/websocket.cpp b/src/io/socket/websocket.cpp index 30d3db74..5ba8cc7c 100644 --- a/src/io/socket/websocket.cpp +++ b/src/io/socket/websocket.cpp @@ -58,7 +58,37 @@ WebSocket::WebSocket(std::unique_ptr 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(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(); } ); @@ -66,6 +96,7 @@ WebSocket::WebSocket(std::unique_ptr socket, 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(_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(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(_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