From e2aeaa0956fb6ff812aba5ce09c55816f75b28d1 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/web/impl/WsBase.hpp        | 21 +++++++++-
 tests/unit/web/ServerTests.cpp | 71 +++++++++++++++++++++++++++-------
 2 files changed, 76 insertions(+), 16 deletions(-)

diff --git a/src/web/impl/WsBase.hpp b/src/web/impl/WsBase.hpp
index 0ea0fae6..cdc606a1 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/DOSGuardInterface.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<dosguard::DOSGuardInterface> dosGuard_;
     bool sending_ = false;
@@ -103,15 +108,26 @@ public:
         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);
     }
 
@@ -134,7 +150,9 @@ public:
     void
     onWrite(boost::system::error_code ec, std::size_t)
     {
+        LOG(perfLog_.info()) << tag() << "xinmeng sent message";
         messages_.pop();
+        --messagesLength_.get();
         sending_ = false;
         if (ec) {
             wsFail(ec, "Failed to write");
@@ -165,6 +183,7 @@ public:
             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 a61b3074..b43469ee 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/Server.hpp"
@@ -42,6 +43,7 @@
 #include <boost/json/value.hpp>
 #include <boost/system/system_error.hpp>
 #include <fmt/core.h>
+#include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 #include <condition_variable>
@@ -144,6 +146,8 @@ private:
     std::optional<std::thread> runner;
 };
 
+struct WebServerTestsWithMockPrometheus : WebServerTest, prometheus::WithMockPrometheus {};
+
 class EchoExecutor {
 public:
     void
@@ -204,7 +208,7 @@ makeServerSync(
 
 }  // namespace
 
-TEST_F(WebServerTest, Http)
+TEST_F(WebServerTestsWithMockPrometheus, Http)
 {
     auto e = std::make_shared<EchoExecutor>();
     auto const server = makeServerSync(cfg, ctx, dosGuard, e);
@@ -212,8 +216,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, dosGuard, e);
     WebSocketSyncClient wsClient;
@@ -223,7 +232,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, dosGuard, e);
@@ -234,8 +243,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, dosGuard, e);
     WebSocketSyncClient wsClient;
@@ -248,8 +262,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, dosGuard, e);
     WebSocketSyncClient wsClient;
@@ -262,7 +281,7 @@ TEST_F(WebServerTest, WsInternalErrorNotJson)
     );
 }
 
-TEST_F(WebServerTest, IncompleteSslConfig)
+TEST_F(WebServerTestsWithMockPrometheus, IncompleteSslConfig)
 {
     auto e = std::make_shared<EchoExecutor>();
 
@@ -273,7 +292,7 @@ TEST_F(WebServerTest, IncompleteSslConfig)
     EXPECT_EQ(server, nullptr);
 }
 
-TEST_F(WebServerTest, WrongSslConfig)
+TEST_F(WebServerTestsWithMockPrometheus, WrongSslConfig)
 {
     auto e = std::make_shared<EchoExecutor>();
 
@@ -285,7 +304,7 @@ TEST_F(WebServerTest, WrongSslConfig)
     EXPECT_EQ(server, nullptr);
 }
 
-TEST_F(WebServerTest, Https)
+TEST_F(WebServerTestsWithMockPrometheus, Https)
 {
     auto e = std::make_shared<EchoExecutor>();
     cfg = Config{addSslConfig(generateJSONWithDynamicPort(port))};
@@ -294,8 +313,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>();
     cfg = Config{addSslConfig(generateJSONWithDynamicPort(port))};
     auto server = makeServerSync(cfg, ctx, dosGuard, e);
@@ -306,7 +330,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, dosGuardOverload, e);
@@ -319,8 +343,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, dosGuardOverload, e);
     WebSocketSyncClient wsClient;
@@ -338,7 +367,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>();
@@ -350,8 +379,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, dosGuardOverload, e);
@@ -365,7 +399,7 @@ TEST_F(WebServerTest, WsPayloadOverload)
     );
 }
 
-TEST_F(WebServerTest, WsTooManyConnection)
+TEST_F(WebServerTestsWithMockPrometheus, WsTooManyConnection)
 {
     auto e = std::make_shared<EchoExecutor>();
     auto server = makeServerSync(cfg, ctx, dosGuardOverload, e);
@@ -471,10 +505,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{boost::json::parse(GetParam().config)};
     auto server = makeServerSync(serverConfig, ctx, dosGuardOverload, e);