Add packet counting to the prometheus exporter

CubicMC · Nov 9, 2023 · 09eb8a8 · 09eb8a8
1 parent 6831719
commit 09eb8a8
Show file tree

Hide file tree

Showing 5 changed files with 44 additions and 2 deletions.
diff --git a/cubic-server/Client.cpp b/cubic-server/Client.cpp
@@ -148,6 +148,7 @@ void Client::handleParsedClientPacket(std::unique_ptr<protocol::BaseServerPacket
 {
     using namespace protocol;
 
+    PEXP(incrementPacketRxCounter);
     switch (_status) {
     case ClientStatus::Initial:
         switch (packetID) {

diff --git a/cubic-server/PrometheusExporter.cpp b/cubic-server/PrometheusExporter.cpp
@@ -7,7 +7,8 @@
 #include "options.hpp"
 
 PrometheusExporter::PrometheusExporter(const std::string &bind):
-    _exposer(bind)
+    _exposer(bind),
+    _ready(false)
 {
 }
 
@@ -19,11 +20,16 @@ void PrometheusExporter::registerMetrics()
     // To add anything else to the metrics follow the readme of prometheus-cpp:
     // https://github.com/jupp0r/prometheus-cpp/blob/master/README.md
     // clang-format off
-    UNUSED auto &packet_counter = prometheus::BuildCounter()
+    auto &packet_counter = prometheus::BuildCounter()
         .Name("observed_packets_total")
         .Help("Number of observed packets")
         .Register(*_registry);
     // clang-format on
 
+    _packet_rx_counter = &packet_counter.Add({{"direction", "rx"}});
+    _packet_tx_counter = &packet_counter.Add({{"direction", "tx"}});
+
     _exposer.RegisterCollectable(_registry);
+
+    _ready = true;
 }
diff --git a/cubic-server/PrometheusExporter.hpp b/cubic-server/PrometheusExporter.hpp
@@ -4,17 +4,33 @@
 #include <memory>
 #include <string>
 
+#include "prometheus/counter.h"
 #include "prometheus/exposer.h"
 #include "prometheus/registry.h"
 
 class PrometheusExporter {
 public:
     PrometheusExporter(const std::string &bind);
     void registerMetrics();
+    void incrementPacketRxCounter()
+    {
+        if (_ready)
+            _packet_rx_counter->Increment();
+    }
+    void incrementPacketTxCounter()
+    {
+        if (_ready)
+            _packet_tx_counter->Increment();
+    }
 
 private:
     prometheus::Exposer _exposer;
     std::shared_ptr<prometheus::Registry> _registry;
+
+    prometheus::Counter *_packet_tx_counter;
+    prometheus::Counter *_packet_rx_counter;
+
+    bool _ready;
 };
 
 #endif /* D5AEBC4C_6028_4506_86F2_DA6261F54832 */
diff --git a/cubic-server/Server.cpp b/cubic-server/Server.cpp
@@ -156,10 +156,13 @@ void Server::launch(const configuration::ConfigHandler &config)
 
     _writeThread = std::thread(&Server::_writeLoop, this);
 
+#if PROMETHEUS_SUPPORT == 1
     if (CONFIG["monitoring-prometheus-enable"].as<bool>()) {
         _prometheusExporter = std::make_unique<PrometheusExporter>("0.0.0.0:4242");
         _prometheusExporter->registerMetrics();
+        _prometheusExporterOn = true;
     }
+#endif
 
     _doAccept();
 
@@ -201,6 +204,7 @@ void Server::_writeLoop()
             }
             boost::system::error_code ec;
             boost::asio::write(client->getSocket(), boost::asio::buffer(data.data->data(), data.data->size()), ec);
+            PEXP(incrementPacketTxCounter);
             // TODO(huntears): Handle errors properly xd
             if (ec) {
                 client->disconnect("Network error");

diff --git a/cubic-server/Server.hpp b/cubic-server/Server.hpp
@@ -57,6 +57,19 @@ constexpr uint16_t MS_PER_TICK = 50;
 #define SOUND_EVENT_CONVERTER Server::getInstance()->getSoundEventConverter()
 #define CONFIG Server::getInstance()->getConfig()
 
+#if PROMETHEUS_SUPPORT == 1
+#define PROMETHEUS Server::getInstance()->getPrometheusExporter()
+#define PEXP(method)                                   \
+    if (Server::getInstance()->prometheusExporterOn()) \
+        Server::getInstance()->getPrometheusExporter().method();
+#define PEXPP(method, ...)                             \
+    if (Server::getInstance()->prometheusExporterOn()) \
+        Server::getInstance()->getPrometheusExporter().method(__VA_ARGS__);
+#else
+#define PEXP(method)
+#define PEXPP(method, ...)
+#endif
+
 class Client;
 class WorldGroup;
 
@@ -161,10 +174,12 @@ class Server {
 
 #if PROMETHEUS_SUPPORT == 1
     std::unique_ptr<PrometheusExporter> _prometheusExporter;
+    bool _prometheusExporterOn;
 
 public:
     NODISCARD inline const PrometheusExporter &getPrometheusExporter() const { return *_prometheusExporter; }
     NODISCARD inline PrometheusExporter &getPrometheusExporter() { return *_prometheusExporter; }
+    bool prometheusExporterOn() const noexcept { return _prometheusExporterOn; }
 #endif
 
 public: