diff --git a/net/WebSocketHandler.hpp b/net/WebSocketHandler.hpp index 9bdf130a3a..6d2823d566 100644 --- a/net/WebSocketHandler.hpp +++ b/net/WebSocketHandler.hpp @@ -238,7 +238,8 @@ protected: public: void shutdown(const StatusCodes statusCode = StatusCodes::NORMAL_CLOSE, - const std::string& statusMessage = std::string()) + const std::string& statusMessage = std::string(), + bool hardShutdown = false) { std::shared_ptr socket = _socket.lock(); if (socket) @@ -250,6 +251,10 @@ public: socket->ignoreInput(); assert(socket->getInBuffer().empty() && "Socket buffer must be empty after ignoreInput"); + + // force close after writing this message + if (hardShutdown) + socket->shutdown(); } _wsPayload.clear(); @@ -259,6 +264,13 @@ public: _shuttingDown = false; } + /// Don't wait for the remote Websocket to handshake with us; go down fast. + void shutdownAfterWriting() + { + shutdown(WebSocketHandler::StatusCodes::NORMAL_CLOSE, std::string(), + true /* hard async shutdown & close */); + } + private: bool handleTCPStream(const std::shared_ptr& socket) {