Skip to content

Commit 9b0b4f5

Browse files
cindyyan317cindyyan317
authored
chore: Add counter for total messages waiting to be sent (#1691)
1 parent a210117 commit 9b0b4f5

File tree

2 files changed

+76
-16
lines changed

2 files changed

+76
-16
lines changed

‎src/web/impl/WsBase.hpp‎

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,9 @@
2323
#include "rpc/common/Types.hpp"
2424
#include "util/Taggable.hpp"
2525
#include "util/log/Logger.hpp"
26+
#include "util/prometheus/Gauge.hpp"
27+
#include "util/prometheus/Label.hpp"
28+
#include "util/prometheus/Prometheus.hpp"
2629
#include "web/dosguard/DOSGuardInterface.hpp"
2730
#include "web/interface/Concepts.hpp"
2831
#include "web/interface/ConnectionBase.hpp"
@@ -71,6 +74,8 @@ template <template <typename> typename Derived, SomeServerHandler HandlerType>
7174
class WsBase : public ConnectionBase, public std::enable_shared_from_this<WsBase<Derived, HandlerType>> {
7275
using std::enable_shared_from_this<WsBase<Derived, HandlerType>>::shared_from_this;
7376

77+
std::reference_wrapper<util::prometheus::GaugeInt> messagesLength_;
78+
7479
boost::beast::flat_buffer buffer_;
7580
std::reference_wrapper<dosguard::DOSGuardInterface> dosGuard_;
7681
bool sending_ = false;
@@ -103,15 +108,26 @@ class WsBase : public ConnectionBase, public std::enable_shared_from_this<WsBase
103108
std::shared_ptr<HandlerType> const& handler,
104109
boost::beast::flat_buffer&& buffer
105110
)
106-
: ConnectionBase(tagFactory, ip), buffer_(std::move(buffer)), dosGuard_(dosGuard), handler_(handler)
111+
: ConnectionBase(tagFactory, ip)
112+
, messagesLength_(PrometheusService::gaugeInt(
113+
"ws_messages_length",
114+
util::prometheus::Labels(),
115+
"The total length of messages in the queue"
116+
))
117+
, buffer_(std::move(buffer))
118+
, dosGuard_(dosGuard)
119+
, handler_(handler)
107120
{
108121
upgraded = true; // NOLINT (cppcoreguidelines-pro-type-member-init)
122+
109123
LOG(perfLog_.debug()) << tag() << "session created";
110124
}
111125

112126
~WsBase() override
113127
{
114128
LOG(perfLog_.debug()) << tag() << "session closed";
129+
if (!messages_.empty())
130+
messagesLength_.get() -= messages_.size();
115131
dosGuard_.get().decrement(clientIp);
116132
}
117133

@@ -134,7 +150,9 @@ class WsBase : public ConnectionBase, public std::enable_shared_from_this<WsBase
134150
void
135151
onWrite(boost::system::error_code ec, std::size_t)
136152
{
153+
LOG(perfLog_.info()) << tag() << "xinmeng sent message";
137154
messages_.pop();
155+
--messagesLength_.get();
138156
sending_ = false;
139157
if (ec) {
140158
wsFail(ec, "Failed to write");
@@ -165,6 +183,7 @@ class WsBase : public ConnectionBase, public std::enable_shared_from_this<WsBase
165183
derived().ws().get_executor(),
166184
[this, self = derived().shared_from_this(), msg = std::move(msg)]() {
167185
messages_.push(msg);
186+
++messagesLength_.get();
168187
maybeSendNext();
169188
}
170189
);

‎tests/unit/web/ServerTests.cpp‎

Lines changed: 56 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
#include "util/MockPrometheus.hpp"
2323
#include "util/TestHttpSyncClient.hpp"
2424
#include "util/config/Config.hpp"
25+
#include "util/prometheus/Gauge.hpp"
2526
#include "util/prometheus/Label.hpp"
2627
#include "util/prometheus/Prometheus.hpp"
2728
#include "web/Server.hpp"
@@ -42,6 +43,7 @@
4243
#include <boost/json/value.hpp>
4344
#include <boost/system/system_error.hpp>
4445
#include <fmt/core.h>
46+
#include <gmock/gmock.h>
4547
#include <gtest/gtest.h>
4648

4749
#include <condition_variable>
@@ -144,6 +146,8 @@ struct WebServerTest : NoLoggerFixture {
144146
std::optional<std::thread> runner;
145147
};
146148

149+
struct WebServerTestsWithMockPrometheus : WebServerTest, prometheus::WithMockPrometheus {};
150+
147151
class EchoExecutor {
148152
public:
149153
void
@@ -204,16 +208,21 @@ makeServerSync(
204208

205209
} // namespace
206210

207-
TEST_F(WebServerTest, Http)
211+
TEST_F(WebServerTestsWithMockPrometheus, Http)
208212
{
209213
auto e = std::make_shared<EchoExecutor>();
210214
auto const server = makeServerSync(cfg, ctx, dosGuard, e);
211215
auto const res = HttpSyncClient::syncPost("localhost", port, R"({"Hello":1})");
212216
EXPECT_EQ(res, R"({"Hello":1})");
213217
}
214218

215-
TEST_F(WebServerTest, Ws)
219+
TEST_F(WebServerTestsWithMockPrometheus, Ws)
216220
{
221+
::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
222+
makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
223+
EXPECT_CALL(wsMessagesCounterMock, add(1));
224+
EXPECT_CALL(wsMessagesCounterMock, add(-1));
225+
217226
auto e = std::make_shared<EchoExecutor>();
218227
auto const server = makeServerSync(cfg, ctx, dosGuard, e);
219228
WebSocketSyncClient wsClient;
@@ -223,7 +232,7 @@ TEST_F(WebServerTest, Ws)
223232
wsClient.disconnect();
224233
}
225234

226-
TEST_F(WebServerTest, HttpInternalError)
235+
TEST_F(WebServerTestsWithMockPrometheus, HttpInternalError)
227236
{
228237
auto e = std::make_shared<ExceptionExecutor>();
229238
auto const server = makeServerSync(cfg, ctx, dosGuard, e);
@@ -234,8 +243,13 @@ TEST_F(WebServerTest, HttpInternalError)
234243
);
235244
}
236245

237-
TEST_F(WebServerTest, WsInternalError)
246+
TEST_F(WebServerTestsWithMockPrometheus, WsInternalError)
238247
{
248+
::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
249+
makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
250+
EXPECT_CALL(wsMessagesCounterMock, add(1));
251+
EXPECT_CALL(wsMessagesCounterMock, add(-1));
252+
239253
auto e = std::make_shared<ExceptionExecutor>();
240254
auto const server = makeServerSync(cfg, ctx, dosGuard, e);
241255
WebSocketSyncClient wsClient;
@@ -248,8 +262,13 @@ TEST_F(WebServerTest, WsInternalError)
248262
);
249263
}
250264

251-
TEST_F(WebServerTest, WsInternalErrorNotJson)
265+
TEST_F(WebServerTestsWithMockPrometheus, WsInternalErrorNotJson)
252266
{
267+
::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
268+
makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
269+
EXPECT_CALL(wsMessagesCounterMock, add(1));
270+
EXPECT_CALL(wsMessagesCounterMock, add(-1));
271+
253272
auto e = std::make_shared<ExceptionExecutor>();
254273
auto const server = makeServerSync(cfg, ctx, dosGuard, e);
255274
WebSocketSyncClient wsClient;
@@ -262,7 +281,7 @@ TEST_F(WebServerTest, WsInternalErrorNotJson)
262281
);
263282
}
264283

265-
TEST_F(WebServerTest, IncompleteSslConfig)
284+
TEST_F(WebServerTestsWithMockPrometheus, IncompleteSslConfig)
266285
{
267286
auto e = std::make_shared<EchoExecutor>();
268287

@@ -273,7 +292,7 @@ TEST_F(WebServerTest, IncompleteSslConfig)
273292
EXPECT_EQ(server, nullptr);
274293
}
275294

276-
TEST_F(WebServerTest, WrongSslConfig)
295+
TEST_F(WebServerTestsWithMockPrometheus, WrongSslConfig)
277296
{
278297
auto e = std::make_shared<EchoExecutor>();
279298

@@ -285,7 +304,7 @@ TEST_F(WebServerTest, WrongSslConfig)
285304
EXPECT_EQ(server, nullptr);
286305
}
287306

288-
TEST_F(WebServerTest, Https)
307+
TEST_F(WebServerTestsWithMockPrometheus, Https)
289308
{
290309
auto e = std::make_shared<EchoExecutor>();
291310
cfg = Config{addSslConfig(generateJSONWithDynamicPort(port))};
@@ -294,8 +313,13 @@ TEST_F(WebServerTest, Https)
294313
EXPECT_EQ(res, R"({"Hello":1})");
295314
}
296315

297-
TEST_F(WebServerTest, Wss)
316+
TEST_F(WebServerTestsWithMockPrometheus, Wss)
298317
{
318+
::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
319+
makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
320+
EXPECT_CALL(wsMessagesCounterMock, add(1));
321+
EXPECT_CALL(wsMessagesCounterMock, add(-1));
322+
299323
auto e = std::make_shared<EchoExecutor>();
300324
cfg = Config{addSslConfig(generateJSONWithDynamicPort(port))};
301325
auto server = makeServerSync(cfg, ctx, dosGuard, e);
@@ -306,7 +330,7 @@ TEST_F(WebServerTest, Wss)
306330
wsClient.disconnect();
307331
}
308332

309-
TEST_F(WebServerTest, HttpRequestOverload)
333+
TEST_F(WebServerTestsWithMockPrometheus, HttpRequestOverload)
310334
{
311335
auto e = std::make_shared<EchoExecutor>();
312336
auto const server = makeServerSync(cfg, ctx, dosGuardOverload, e);
@@ -319,8 +343,13 @@ TEST_F(WebServerTest, HttpRequestOverload)
319343
);
320344
}
321345

322-
TEST_F(WebServerTest, WsRequestOverload)
346+
TEST_F(WebServerTestsWithMockPrometheus, WsRequestOverload)
323347
{
348+
::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
349+
makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
350+
EXPECT_CALL(wsMessagesCounterMock, add(1)).Times(2);
351+
EXPECT_CALL(wsMessagesCounterMock, add(-1)).Times(2);
352+
324353
auto e = std::make_shared<EchoExecutor>();
325354
auto const server = makeServerSync(cfg, ctx, dosGuardOverload, e);
326355
WebSocketSyncClient wsClient;
@@ -338,7 +367,7 @@ TEST_F(WebServerTest, WsRequestOverload)
338367
);
339368
}
340369

341-
TEST_F(WebServerTest, HttpPayloadOverload)
370+
TEST_F(WebServerTestsWithMockPrometheus, HttpPayloadOverload)
342371
{
343372
std::string const s100(100, 'a');
344373
auto e = std::make_shared<EchoExecutor>();
@@ -350,8 +379,13 @@ TEST_F(WebServerTest, HttpPayloadOverload)
350379
);
351380
}
352381

353-
TEST_F(WebServerTest, WsPayloadOverload)
382+
TEST_F(WebServerTestsWithMockPrometheus, WsPayloadOverload)
354383
{
384+
::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
385+
makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
386+
EXPECT_CALL(wsMessagesCounterMock, add(1));
387+
EXPECT_CALL(wsMessagesCounterMock, add(-1));
388+
355389
std::string const s100(100, 'a');
356390
auto e = std::make_shared<EchoExecutor>();
357391
auto server = makeServerSync(cfg, ctx, dosGuardOverload, e);
@@ -365,7 +399,7 @@ TEST_F(WebServerTest, WsPayloadOverload)
365399
);
366400
}
367401

368-
TEST_F(WebServerTest, WsTooManyConnection)
402+
TEST_F(WebServerTestsWithMockPrometheus, WsTooManyConnection)
369403
{
370404
auto e = std::make_shared<EchoExecutor>();
371405
auto server = makeServerSync(cfg, ctx, dosGuardOverload, e);
@@ -471,10 +505,17 @@ struct WebServerAdminTestParams {
471505
std::string expectedResponse;
472506
};
473507

474-
class WebServerAdminTest : public WebServerTest, public ::testing::WithParamInterface<WebServerAdminTestParams> {};
508+
class WebServerAdminTest : public WebServerTest,
509+
public ::testing::WithParamInterface<WebServerAdminTestParams>,
510+
public prometheus::WithMockPrometheus {};
475511

476512
TEST_P(WebServerAdminTest, WsAdminCheck)
477513
{
514+
::testing::StrictMock<util::prometheus::MockCounterImplInt>& wsMessagesCounterMock =
515+
makeMock<util::prometheus::GaugeInt>("ws_messages_length", "");
516+
EXPECT_CALL(wsMessagesCounterMock, add(1));
517+
EXPECT_CALL(wsMessagesCounterMock, add(-1));
518+
478519
auto e = std::make_shared<AdminCheckExecutor>();
479520
Config const serverConfig{boost::json::parse(GetParam().config)};
480521
auto server = makeServerSync(serverConfig, ctx, dosGuardOverload, e);

0 commit comments

Comments
 (0)