Skip to content

Commit

Permalink
chore: Add counter for total messages waiting to be sent (#1691)
Browse files Browse the repository at this point in the history
  • Loading branch information
cindyyan317 authored and kuznetsss committed Oct 18, 2024
1 parent 189098d commit 8b0e68f
Show file tree
Hide file tree
Showing 4 changed files with 78 additions and 23 deletions.
2 changes: 1 addition & 1 deletion src/util/requests/impl/WsConnectionImpl.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
12 changes: 4 additions & 8 deletions src/web/Server.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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;
Expand Down
20 changes: 19 additions & 1 deletion src/web/impl/WsBase.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}

Expand All @@ -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");
Expand Down Expand Up @@ -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();
}
);
Expand Down
67 changes: 54 additions & 13 deletions tests/unit/web/ServerTests.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -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>
Expand Down Expand Up @@ -202,6 +204,8 @@ class WebServerTest : public NoLoggerFixture {
std::optional<std::thread> runner;
};

struct WebServerTestsWithMockPrometheus : WebServerTest, prometheus::WithMockPrometheus {};

class EchoExecutor {
public:
void
Expand Down Expand Up @@ -263,16 +267,21 @@ 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);
auto const res = HttpSyncClient::syncPost("localhost", port, R"({"Hello":1})");
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;
Expand All @@ -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);
Expand All @@ -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;
Expand All @@ -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;
Expand All @@ -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();
Expand All @@ -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;
Expand All @@ -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);
Expand All @@ -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;
Expand All @@ -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>();
Expand All @@ -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);
Expand All @@ -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);
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit 8b0e68f

Please sign in to comment.