1 files changed, 36 insertions, 0 deletions

diff --git a/indra/llcorehttp/llwebsocketmgr.cpp b/indra/llcorehttp/llwebsocketmgr.cpp
index ea87d1e07e..be1ab54efd 100644
--- a/indra/llcorehttp/llwebsocketmgr.cpp
+++ b/indra/llcorehttp/llwebsocketmgr.cpp
@@ -65,6 +65,7 @@ void LLWebsocketMgr::cleanupSingleton()
 
 void LLWebsocketMgr::update()
 {
+    LL_PROFILE_ZONE_SCOPED_CATEGORY_WEBSOCKET;
     std::vector<WSServer::ptr_t> stops;
 
     for (auto &[name, server] : mServers)
@@ -255,6 +256,7 @@ struct Server_impl
             // Run controlled event loop with periodic stop flag checking
             while (!mOwner->mShouldStop && !mServer.stopped())
             {
+                LL_PROFILE_ZONE_NAMED_CATEGORY_WEBSOCKET("ws server run_for");
                 // Process events for up to 100ms, then check the stop flag
                 std::chrono::milliseconds timeout(100);
                 std::size_t handlers_run = mServer.get_io_service().run_for(timeout);
@@ -314,6 +316,7 @@ struct Server_impl
      */
     void onOpen(websocketpp::connection_hdl hdl) const
     {
+        LL_PROFILE_ZONE_SCOPED_CATEGORY_WEBSOCKET;
         LL_ERRS_IF(!mOwner, "WebSocket") << "mOwner should never be null. If it is, something is very wrong!" << LL_ENDL;
 
         mOwner->handleOpenConnection(hdl);
@@ -325,6 +328,7 @@ struct Server_impl
      */
     void onClose(websocketpp::connection_hdl hdl) const
     {
+        LL_PROFILE_ZONE_SCOPED_CATEGORY_WEBSOCKET;
         LL_ERRS_IF(!mOwner, "WebSocket") << "mOwner should never be null" << LL_ENDL;
         mOwner->handleCloseConnection(hdl);
     }
@@ -342,6 +346,7 @@ struct Server_impl
      */
     void onMessage(websocketpp::connection_hdl hdl, Server_t::message_ptr msg) const
     {
+        LL_PROFILE_ZONE_SCOPED_CATEGORY_WEBSOCKET;
         LL_ERRS_IF(!mOwner, "WebSocket") << "mOwner should never be null" << LL_ENDL;
         LLWebsocketMgr::WSConnection::ptr_t connection = mOwner->getConnection(hdl);
         if (!connection)
@@ -441,6 +446,9 @@ void LLWebsocketMgr::WSServer::stop()
 
         mShouldStop = true;
 
+        // Send close frames to all connected clients before stopping the ASIO loop
+        closeAllConnections(1001, "Server shutting down");
+
         // Stop the websocket server (this will cause the controlled run loop to exit)
         mImpl->stop();
     } // Release the lock here
@@ -463,6 +471,7 @@ bool LLWebsocketMgr::WSServer::isRunning() const
 
 void LLWebsocketMgr::WSServer::broadcastMessage(const std::string& message)
 {
+    LL_PROFILE_ZONE_SCOPED_CATEGORY_WEBSOCKET;
     LL_ERRS_IF(!mImpl, "WebSocket") << "WebSocket server " << mServerName << " implementation is null !" << LL_ENDL;
     LLMutexLock lock(&mConnectionMutex);
     for (const auto& [handle, conn] : mConnections)
@@ -473,6 +482,7 @@ void LLWebsocketMgr::WSServer::broadcastMessage(const std::string& message)
 
 bool LLWebsocketMgr::WSServer::sendMessageTo(const connection_h& handle, const std::string& message)
 {
+    LL_PROFILE_ZONE_SCOPED_CATEGORY_WEBSOCKET;
     LL_ERRS_IF(!mImpl, "WebSocket") << "WebSocket server " << mServerName << " implementation is null !" << LL_ENDL;
     websocketpp::lib::error_code ec;
     mImpl->mServer.send(handle, message, websocketpp::frame::opcode::text, ec);
@@ -546,6 +556,7 @@ LLWebsocketMgr::connection_state_t LLWebsocketMgr::WSServer::getConnectionState(
 
 void LLWebsocketMgr::WSServer::handleOpenConnection(const connection_h& handle)
 {
+    LL_PROFILE_ZONE_SCOPED_CATEGORY_WEBSOCKET;
     WSConnection::ptr_t connection;
     size_t              size(0);
     {
@@ -582,6 +593,7 @@ void LLWebsocketMgr::WSServer::handleOpenConnection(const connection_h& handle)
 
 void LLWebsocketMgr::WSServer::handleCloseConnection(const connection_h& handle)
 {
+    LL_PROFILE_ZONE_SCOPED_CATEGORY_WEBSOCKET;
     size_t              size(0);
     WSConnection::ptr_t connection;
     {
@@ -608,6 +620,7 @@ void LLWebsocketMgr::WSServer::handleCloseConnection(const connection_h& handle)
 
 void LLWebsocketMgr::WSServer::handleMessage(const connection_h& handle, const std::string& message)
 {
+    LL_PROFILE_ZONE_SCOPED_CATEGORY_WEBSOCKET;
     WSConnection::ptr_t connection = getConnection(handle);
     if (connection)
     {
@@ -672,3 +685,26 @@ bool LLWebsocketMgr::WSConnection::isConnected() const
     }
     return server->getConnectionState(mConnectionHandle) == connection_open;
 }
+
+LLWebsocketMgr::WSConnection::ptr_t LLWebsocketMgr::WSConnection::getSelfPtr()
+{
+    auto server = mOwningServer.lock();
+    if (!server) return nullptr;
+    return server->getConnection(mConnectionHandle);
+}
+
+void LLWebsocketMgr::WSServer::closeAllConnections(U16 code, const std::string& reason)
+{
+    std::vector<connection_h> handles;
+    {
+        LLMutexLock lock(&mConnectionMutex);
+        for (const auto& [handle, conn] : mConnections)
+        {
+            handles.push_back(handle);
+        }
+    }
+    for (const auto& handle : handles)
+    {
+        closeConnection(handle, code, reason);
+    }
+}