Fix missing disconnect

XRPLF · Nov 20, 2024 · 27e9ed1 · 27e9ed1
1 parent 908b7e3
commit 27e9ed1
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 1 deletion.
diff --git a/src/rpc/common/Types.hpp b/src/rpc/common/Types.hpp
@@ -33,7 +33,6 @@
 
 #include <cstdint>
 #include <expected>
-#include <memory>
 #include <string>
 #include <utility>
 #include <variant>

diff --git a/src/web/ng/impl/ConnectionHandler.cpp b/src/web/ng/impl/ConnectionHandler.cpp
@@ -167,6 +167,10 @@ ConnectionHandler::processConnection(ConnectionPtr connectionPtr, boost::asio::y
             shouldCloseGracefully = parallelRequestResponseLoop(connectionRef, subscriptionContextInterfacePtr, yield);
             break;
     }
+
+    if (subscriptionContext != nullptr)
+        subscriptionContext->disconnect(yield);
+
     if (shouldCloseGracefully)
         connectionRef.close(yield);
 

diff --git a/tests/unit/web/ng/impl/ConnectionHandlerTests.cpp b/tests/unit/web/ng/impl/ConnectionHandlerTests.cpp
@@ -251,6 +251,36 @@ TEST_F(ConnectionHandlerSequentialProcessingTest, SendSubscriptionMessage)
     });
 }
 
+TEST_F(ConnectionHandlerSequentialProcessingTest, SubscriptionContextIsDisconnectedAfterProcessingFinished)
+{
+    testing::StrictMock<testing::MockFunction<
+        Response(Request const&, ConnectionMetadata const&, web::SubscriptionContextPtr, boost::asio::yield_context)>>
+        wsHandlerMock;
+    connectionHandler_.onWs(wsHandlerMock.AsStdFunction());
+
+    testing::StrictMock<testing::MockFunction<void(web::SubscriptionContextInterface*)>> onDisconnectHook;
+
+    EXPECT_CALL(*mockWsConnection_, wasUpgraded).WillOnce(Return(true));
+    testing::Expectation const expectationReceiveCalled = EXPECT_CALL(*mockWsConnection_, receive)
+                                                              .WillOnce(Return(makeRequest("", Request::HttpHeaders{})))
+                                                              .WillOnce(Return(makeError(websocket::error::closed)));
+
+    EXPECT_CALL(wsHandlerMock, Call)
+        .WillOnce([&](Request const& request, auto&&, web::SubscriptionContextPtr subscriptionContext, auto&&) {
+            EXPECT_NE(subscriptionContext, nullptr);
+            subscriptionContext->onDisconnect(onDisconnectHook.AsStdFunction());
+            return Response(http::status::ok, "", request);
+        });
+
+    EXPECT_CALL(*mockWsConnection_, send).WillOnce(Return(std::nullopt));
+
+    EXPECT_CALL(onDisconnectHook, Call).After(expectationReceiveCalled);
+
+    runSpawn([this](boost::asio::yield_context yield) {
+        connectionHandler_.processConnection(std::move(mockWsConnection_), yield);
+    });
+}
+
 TEST_F(ConnectionHandlerSequentialProcessingTest, SubscriptionContextIsNullForHttpConnection)
 {
     std::string const target = "/some/target";