From 8b0e68f48e5ede435709751bc4975bc8797471c9 Mon Sep 17 00:00:00 2001
From: cyan317 <120398799+cindyyan317@users.noreply.github.com>
Date: Wed, 16 Oct 2024 17:06:27 +0100
Subject: [PATCH] chore: Add counter for total messages waiting to be sent
 (#1691)

---
 src/util/requests/impl/WsConnectionImpl.hpp |  2 +-
 src/web/Server.hpp                          | 12 ++--
 src/web/impl/WsBase.hpp                     | 20 +++++-
 tests/unit/web/ServerTests.cpp              | 67 +++++++++++++++++----
 4 files changed, 78 insertions(+), 23 deletions(-)

diff --git a/src/util/requests/impl/WsConnectionImpl.hpp b/src/util/requests/impl/WsConnectionImpl.hpp
index 04eb95e01..99ce0a3e0 100644
--- a/src/util/requests/impl/WsConnectionImpl.hpp
+++ b/src/util/requests/impl/WsConnectionImpl.hpp
@@ -134,7 +134,7 @@ class WsConnectionImpl : public WsConnection {
         // The timer below can be called with no error code even if the operation is completed before the timeout, so we
         // need an additional flag here
         timer.async_wait([&cancellationSignal, isCompleted](boost::system::error_code errorCode) {
-            if (!errorCode and not *isCompleted)
+            if (!errorCode and not*isCompleted)
                 cancellationSignal.emit(boost::asio::cancellation_type::terminal);
         });
         operation(cyield);
diff --git a/src/web/Server.hpp b/src/web/Server.hpp
index 373c57c5f..e1d5b5f2c 100644
--- a/src/web/Server.hpp
+++ b/src/web/Server.hpp
@@ -69,10 +69,8 @@ namespace web {
  * @tparam HandlerType The executor to handle the requests
  */
 template <
-    template <typename>
-    class PlainSessionType,
-    template <typename>
-    class SslSessionType,
+    template <typename> class PlainSessionType,
+    template <typename> class SslSessionType,
     SomeServerHandler HandlerType>
 class Detector : public std::enable_shared_from_this<Detector<PlainSessionType, SslSessionType, HandlerType>> {
     using std::enable_shared_from_this<Detector<PlainSessionType, SslSessionType, HandlerType>>::shared_from_this;
@@ -191,10 +189,8 @@ class Detector : public std::enable_shared_from_this<Detector<PlainSessionType,
  * @tparam HandlerType The handler to process the request and return response.
  */
 template <
-    template <typename>
-    class PlainSessionType,
-    template <typename>
-    class SslSessionType,
+    template <typename> class PlainSessionType,
+    template <typename> class SslSessionType,
     SomeServerHandler HandlerType>
 class Server : public std::enable_shared_from_this<Server<PlainSessionType, SslSessionType, HandlerType>> {
     using std::enable_shared_from_this<Server<PlainSessionType, SslSessionType, HandlerType>>::shared_from_this;
diff --git a/src/web/impl/WsBase.hpp b/src/web/impl/WsBase.hpp
index 889ed769e..d75335b6e 100644
--- a/src/web/impl/WsBase.hpp
+++ b/src/web/impl/WsBase.hpp
@@ -23,6 +23,9 @@
 #include "rpc/common/Types.hpp"
 #include "util/Taggable.hpp"
 #include "util/log/Logger.hpp"
+#include "util/prometheus/Gauge.hpp"
+#include "util/prometheus/Label.hpp"
+#include "util/prometheus/Prometheus.hpp"
 #include "web/DOSGuard.hpp"
 #include "web/interface/Concepts.hpp"
 #include "web/interface/ConnectionBase.hpp"
@@ -71,6 +74,8 @@ template <template <typename> typename Derived, SomeServerHandler HandlerType>
 class WsBase : public ConnectionBase, public std::enable_shared_from_this<WsBase<Derived, HandlerType>> {
     using std::enable_shared_from_this<WsBase<Derived, HandlerType>>::shared_from_this;
 
+    std::reference_wrapper<util::prometheus::GaugeInt> messagesLength_;
+
     boost::beast::flat_buffer buffer_;
     std::reference_wrapper<web::DOSGuard> dosGuard_;
     bool sending_ = false;
@@ -103,15 +108,26 @@ class WsBase : public ConnectionBase, public std::enable_shared_from_this<WsBase
         std::shared_ptr<HandlerType> const& handler,
         boost::beast::flat_buffer&& buffer
     )
-        : ConnectionBase(tagFactory, ip), buffer_(std::move(buffer)), dosGuard_(dosGuard), handler_(handler)
+        : ConnectionBase(tagFactory, ip)
+        , messagesLength_(PrometheusService::gaugeInt(
+              "ws_messages_length",
+              util::prometheus::Labels(),
+              "The total length of messages in the queue"
+          ))
+        , buffer_(std::move(buffer))
+        , dosGuard_(dosGuard)
+        , handler_(handler)
     {
         upgraded = true;  // NOLINT (cppcoreguidelines-pro-type-member-init)
+
         LOG(perfLog_.debug()) << tag() << "session created";
     }
 
     ~WsBase() override
     {
         LOG(perfLog_.debug()) << tag() << "session closed";
+        if (!messages_.empty())
+            messagesLength_.get() -= messages_.size();
         dosGuard_.get().decrement(clientIp);
     }
 
@@ -135,6 +151,7 @@ class WsBase : public ConnectionBase, public std::enable_shared_from_this<WsBase
     onWrite(boost::system::error_code ec, std::size_t)
     {
         messages_.pop();
+        --messagesLength_.get();
         sending_ = false;
         if (ec) {
             wsFail(ec, "Failed to write");
@@ -165,6 +182,7 @@ class WsBase : public ConnectionBase, public std::enable_shared_from_this<WsBase
             derived().ws().get_executor(),
             [this, self = derived().shared_from_this(), msg = std::move(msg)]() {
                 messages_.push(msg);
+                ++messagesLength_.get();
                 maybeSendNext();
             }
         );
diff --git a/tests/unit/web/ServerTests.cpp b/tests/unit/web/ServerTests.cpp
index 635d534c5..3dc320878 100644
--- a/tests/unit/web/ServerTests.cpp
+++ b/tests/unit/web/ServerTests.cpp
@@ -22,6 +22,7 @@
 #include "util/MockPrometheus.hpp"
 #include "util/TestHttpSyncClient.hpp"
 #include "util/config/Config.hpp"
+#include "util/prometheus/Gauge.hpp"
 #include "util/prometheus/Label.hpp"
 #include "util/prometheus/Prometheus.hpp"
 #include "web/DOSGuard.hpp"
@@ -42,6 +43,7 @@
 #include <boost/json/parse.hpp>
 #include <boost/system/system_error.hpp>
 #include <fmt/core.h>
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 #include <chrono>
@@ -202,6 +204,8 @@ class WebServerTest : public NoLoggerFixture {
     std::optional<std::thread> runner;
 };
 
+struct WebServerTestsWithMockPrometheus : WebServerTest, prometheus::WithMockPrometheus {};
+
 class EchoExecutor {
 public:
     void
@@ -263,7 +267,7 @@ makeServerSync(
 
 }  // namespace
 
-TEST_F(WebServerTest, Http)
+TEST_F(WebServerTestsWithMockPrometheus, Http)
 {
     auto e = std::make_shared<EchoExecutor>();
     auto const server = makeServerSync(cfg, ctx, std::nullopt, dosGuard, e);
@@ -271,8 +275,13 @@ TEST_F(WebServerTest, Http)
     EXPECT_EQ(res, R"({"Hello":1})");
 }
 
-TEST_F(WebServerTest, Ws)
+TEST_F(WebServerTestsWithMockPrometheus, Ws)
 {
+    ::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
+        makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
+    EXPECT_CALL(wsMessagesCounterMock, add(1));
+    EXPECT_CALL(wsMessagesCounterMock, add(-1));
+
     auto e = std::make_shared<EchoExecutor>();
     auto const server = makeServerSync(cfg, ctx, std::nullopt, dosGuard, e);
     WebSocketSyncClient wsClient;
@@ -282,7 +291,7 @@ TEST_F(WebServerTest, Ws)
     wsClient.disconnect();
 }
 
-TEST_F(WebServerTest, HttpInternalError)
+TEST_F(WebServerTestsWithMockPrometheus, HttpInternalError)
 {
     auto e = std::make_shared<ExceptionExecutor>();
     auto const server = makeServerSync(cfg, ctx, std::nullopt, dosGuard, e);
@@ -293,8 +302,13 @@ TEST_F(WebServerTest, HttpInternalError)
     );
 }
 
-TEST_F(WebServerTest, WsInternalError)
+TEST_F(WebServerTestsWithMockPrometheus, WsInternalError)
 {
+    ::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
+        makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
+    EXPECT_CALL(wsMessagesCounterMock, add(1));
+    EXPECT_CALL(wsMessagesCounterMock, add(-1));
+
     auto e = std::make_shared<ExceptionExecutor>();
     auto const server = makeServerSync(cfg, ctx, std::nullopt, dosGuard, e);
     WebSocketSyncClient wsClient;
@@ -307,8 +321,13 @@ TEST_F(WebServerTest, WsInternalError)
     );
 }
 
-TEST_F(WebServerTest, WsInternalErrorNotJson)
+TEST_F(WebServerTestsWithMockPrometheus, WsInternalErrorNotJson)
 {
+    ::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
+        makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
+    EXPECT_CALL(wsMessagesCounterMock, add(1));
+    EXPECT_CALL(wsMessagesCounterMock, add(-1));
+
     auto e = std::make_shared<ExceptionExecutor>();
     auto const server = makeServerSync(cfg, ctx, std::nullopt, dosGuard, e);
     WebSocketSyncClient wsClient;
@@ -321,7 +340,7 @@ TEST_F(WebServerTest, WsInternalErrorNotJson)
     );
 }
 
-TEST_F(WebServerTest, Https)
+TEST_F(WebServerTestsWithMockPrometheus, Https)
 {
     auto e = std::make_shared<EchoExecutor>();
     auto sslCtx = parseCertsForTest();
@@ -331,8 +350,13 @@ TEST_F(WebServerTest, Https)
     EXPECT_EQ(res, R"({"Hello":1})");
 }
 
-TEST_F(WebServerTest, Wss)
+TEST_F(WebServerTestsWithMockPrometheus, Wss)
 {
+    ::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
+        makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
+    EXPECT_CALL(wsMessagesCounterMock, add(1));
+    EXPECT_CALL(wsMessagesCounterMock, add(-1));
+
     auto e = std::make_shared<EchoExecutor>();
     auto sslCtx = parseCertsForTest();
     auto const ctxSslRef = sslCtx ? std::optional<std::reference_wrapper<ssl::context>>{sslCtx.value()} : std::nullopt;
@@ -345,7 +369,7 @@ TEST_F(WebServerTest, Wss)
     wsClient.disconnect();
 }
 
-TEST_F(WebServerTest, HttpRequestOverload)
+TEST_F(WebServerTestsWithMockPrometheus, HttpRequestOverload)
 {
     auto e = std::make_shared<EchoExecutor>();
     auto const server = makeServerSync(cfg, ctx, std::nullopt, dosGuardOverload, e);
@@ -358,8 +382,13 @@ TEST_F(WebServerTest, HttpRequestOverload)
     );
 }
 
-TEST_F(WebServerTest, WsRequestOverload)
+TEST_F(WebServerTestsWithMockPrometheus, WsRequestOverload)
 {
+    ::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
+        makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
+    EXPECT_CALL(wsMessagesCounterMock, add(1)).Times(2);
+    EXPECT_CALL(wsMessagesCounterMock, add(-1)).Times(2);
+
     auto e = std::make_shared<EchoExecutor>();
     auto const server = makeServerSync(cfg, ctx, std::nullopt, dosGuardOverload, e);
     WebSocketSyncClient wsClient;
@@ -377,7 +406,7 @@ TEST_F(WebServerTest, WsRequestOverload)
     );
 }
 
-TEST_F(WebServerTest, HttpPayloadOverload)
+TEST_F(WebServerTestsWithMockPrometheus, HttpPayloadOverload)
 {
     std::string const s100(100, 'a');
     auto e = std::make_shared<EchoExecutor>();
@@ -389,8 +418,13 @@ TEST_F(WebServerTest, HttpPayloadOverload)
     );
 }
 
-TEST_F(WebServerTest, WsPayloadOverload)
+TEST_F(WebServerTestsWithMockPrometheus, WsPayloadOverload)
 {
+    ::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
+        makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
+    EXPECT_CALL(wsMessagesCounterMock, add(1));
+    EXPECT_CALL(wsMessagesCounterMock, add(-1));
+
     std::string const s100(100, 'a');
     auto e = std::make_shared<EchoExecutor>();
     auto server = makeServerSync(cfg, ctx, std::nullopt, dosGuardOverload, e);
@@ -404,7 +438,7 @@ TEST_F(WebServerTest, WsPayloadOverload)
     );
 }
 
-TEST_F(WebServerTest, WsTooManyConnection)
+TEST_F(WebServerTestsWithMockPrometheus, WsTooManyConnection)
 {
     auto e = std::make_shared<EchoExecutor>();
     auto server = makeServerSync(cfg, ctx, std::nullopt, dosGuardOverload, e);
@@ -510,10 +544,17 @@ struct WebServerAdminTestParams {
     std::string expectedResponse;
 };
 
-class WebServerAdminTest : public WebServerTest, public ::testing::WithParamInterface<WebServerAdminTestParams> {};
+class WebServerAdminTest : public WebServerTest,
+                           public ::testing::WithParamInterface<WebServerAdminTestParams>,
+                           public prometheus::WithMockPrometheus {};
 
 TEST_P(WebServerAdminTest, WsAdminCheck)
 {
+    ::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
+        makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
+    EXPECT_CALL(wsMessagesCounterMock, add(1));
+    EXPECT_CALL(wsMessagesCounterMock, add(-1));
+
     auto e = std::make_shared<AdminCheckExecutor>();
     Config const serverConfig{parse(GetParam().config)};
     auto server = makeServerSync(serverConfig, ctx, std::nullopt, dosGuardOverload, e);