metrics-collector fixes

Author: jcdkiki

src/main.cpp

@@ -1,28 +1,90 @@
-#include <iostream>
+#include "metrics-collector.hpp"
 #include "file.hpp"
 
+#include <iostream>
+#include <cstring>
+#include <thread>
+#include <signal.h>
+
 void print_usage(const std::string& program_name) {
-    std::cerr << "Usage: " << program_name 
+    std::cerr << "Usage: " << program_name << " hash" 
               << " PATH ALGORITHM" << std::endl
               << "Supported algorithms are: "
               << "md2, md5, sha, " 
               << "sha1, sha224, sha256, sha384, sha512, "
               << "mdc2 and ripemd160" << std::endl;
 }
 
-int main(int argc, char* argv[]) {
-    if (argc < 3) {
+int hash_main(int argc, char* argv[]) {
+    if (argc < 4) {
         print_usage(argv[0]);
         return 1;
     }
 
     try {
-        std::string hash = File::calculate_hash(argv[1], argv[2]);
+        std::string hash = File::calculate_hash(argv[2], argv[3]);
         std::cout << "Digest is: " << hash << std::endl;
     } catch (const std::exception& e) {
         std::cerr << "Error: " << e.what() << std::endl;
         return 1;
     }
 
     return 0;
-}
+}
+
+
+bool is_running = true;
+
+void signal_handler(int signo) {
+    std::cerr << "[INFO] Received " << strsignal(signo) << " signal. Stopping" << std::endl;
+    is_running = false;
+}
+
+int metrics_main(int argc, char *argv[]) {
+    signal(SIGTERM, signal_handler);
+    signal(SIGINT, signal_handler);
+
+    if (argc != 5) {
+        std::cerr << "usage: worker" << argv[0] << " metrics" 
+                  << " ADDRESS PORT WORKER_NAME" << std::endl;
+        return 1;
+    }
+    
+    MetricsCollector metrics_collector(argv[2], argv[3], argv[4]);
+   
+    srand(time(NULL));
+    while (is_running) {
+        std::this_thread::sleep_for(std::chrono::seconds(rand() % 5 + 1));
+        std::cerr << "[INFO] Doing task" << std::endl;
+        metrics_collector.StartTask();
+    
+        // some "work" that requires memory
+        void *data = malloc(rand() % (1024 * 1024 * 50)); // 0 - 50 megabytes
+        for (volatile int i = 0; i < 1000000000 + rand() % 10000000000; i++) {}
+        free(data);
+    
+        metrics_collector.StopTask();
+        std::cerr << "[INFO] Task done" << std::endl;
+    }
+
+    return 0;
+}
+
+int main(int argc, char *argv[])
+{
+    static const char usage[] = "Specify hash/metrics";
+
+    if (argc < 2) {
+        std::cerr << usage << std::endl;
+        return 1; 
+    }
+
+    if (strcmp(argv[1], "hash") == 0) {
+        return hash_main(argc, argv);
+    } else if (strcmp(argv[1], "metrics") == 0) {
+        return metrics_main(argc, argv);
+    }
+
+    std::cerr << usage << std::endl;
+    return 1;
+}

src/metrics-collector.cpp

@@ -4,6 +4,22 @@
 #include <chrono>
 #include <unistd.h>
 #include <fstream>
+#include <iostream>
+
+namespace {
+    double GetMemoryUsed()
+    {
+        std::ifstream file("/proc/self/statm");
+        if (!file.is_open()) {
+            return 0;
+        }
+
+        long mem_pages = 0;
+        file >> mem_pages;
+        file.close();
+        return mem_pages * (double)getpagesize();
+    }
+}
 
 MetricsCollector::MetricsCollector(const char *gateway_address, const char *gateway_port, const char *worker_name)
     : gateway(gateway_address, gateway_port, worker_name),
@@ -31,8 +47,8 @@ MetricsCollector::MetricsCollector(const char *gateway_address, const char *gate
     while (true) {
         CPUInfo cpu;
 
-        file >> cpu_name >> cpu.last_total_user >> cpu.last_total_user_low
-             >> cpu.last_total_sys >> cpu.last_total_idle
+        file >> cpu_name >> cpu.time.user >> cpu.time.user_low
+             >> cpu.time.sys >> cpu.time.idle
              >> ign >> ign >> ign >> ign >> ign >> ign;
 
         if (cpu_name.find("cpu") != 0)
@@ -48,30 +64,30 @@ MetricsCollector::MetricsCollector(const char *gateway_address, const char *gate
     task_processing_time_gauge = &task_processing_time_family.Add({});
 
     gateway.RegisterCollectable(registry);
-
-    is_running = true;
-    thread = std::thread(&MetricsCollector::mainLoop, this);
-
+    thread = std::thread(&MetricsCollector::MainLoop, this);
     is_task_running = false;
 }
 
-void MetricsCollector::mainLoop()
+void MetricsCollector::MainLoop()
 {
     while (is_running) {
         std::this_thread::sleep_for(std::chrono::seconds(1));
 
-        getMemoryUsed();
-        getCPUUsage();
+        memory_used_gauge->Set(::GetMemoryUsed());
+        GetCPUUsage();
 
         if (is_task_running) {
-            struct timespec time;
-            clock_gettime(CLOCK_MONOTONIC, &time);
-            task_processing_time_gauge->Set((time.tv_sec - task_start.tv_sec) + (time.tv_nsec - task_start.tv_nsec) * 1e-9);
+            auto cur_time = std::chrono::high_resolution_clock::now();
+            task_processing_time_gauge->Set(std::chrono::duration<double>(cur_time - task_start).count());
         }
-        else
+        else {
             task_processing_time_gauge->Set(0);
+        }
 
-        gateway.PushAdd();
+        int status = gateway.PushAdd();
+        if (status != 200) {
+            std::cerr << "[ERROR] Failed to push metrics. Status " << status << std::endl;
+        }
     }
 }
 
@@ -81,73 +97,53 @@ MetricsCollector::~MetricsCollector()
     thread.join();
 }
 
-void MetricsCollector::getMemoryUsed()
-{
-    std::ifstream file("/proc/self/statm");
-    if (!file.is_open()) {
-        memory_used_gauge->Set(0);
-        return;
-    }
-
-    long mem_pages = 0;
-    file >> mem_pages;
-    file.close();
-
-    memory_used_gauge->Set(mem_pages * (double)getpagesize());
-}
-
-void MetricsCollector::getCPUUsage()
+void MetricsCollector::GetCPUUsage()
 {
     std::ifstream file("/proc/stat");
-    uint64_t total_user, total_user_low, total_sys, total_idle, total;
+    CPUInfo::Time cur_time;
     double percent;
 
     std::string cpu_name;
     int ign;
 
     while (true) {
-        file >> cpu_name >> total_user >> total_user_low >> total_sys >> total_idle
+        file >> cpu_name >> cur_time.user >> cur_time.user_low >> cur_time.sys >> cur_time.idle
              >> ign >> ign >> ign >> ign >> ign >> ign;
 
         if (cpu_name.find("cpu") != 0)
             break;
 
         CPUInfo &cpu = cpu_usage[cpu_name];
-        if (total_user < cpu.last_total_user || total_user_low < cpu.last_total_user_low ||
-            total_sys < cpu.last_total_sys || total_idle < cpu.last_total_idle) {
+        if (cur_time.user < cpu.time.user || cur_time.user_low < cpu.time.user_low ||
+            cur_time.sys < cpu.time.sys || cur_time.idle < cpu.time.idle) {
             // overflow detection
             percent = -1.0;
         }
         else {
-            total = (total_user - cpu.last_total_user) + (total_user_low - cpu.last_total_user_low) +
-                (total_sys - cpu.last_total_sys);
+            uint64_t total = (cur_time.user - cpu.time.user) + (cur_time.user_low - cpu.time.user_low) +
+                (cur_time.sys - cpu.time.sys);
 
             percent = total;
-            total += (total_idle - cpu.last_total_idle);
-            percent /= total;
-            percent *= 100;
+            total += (cur_time.idle - cpu.time.idle);
+            percent = (total == 0) ? -1.0 : (percent / total) * 100.0;
         }
 
-        cpu.last_total_user = total_user;
-        cpu.last_total_user_low = total_user_low;
-        cpu.last_total_sys = total_sys;
-        cpu.last_total_idle = total_idle;
-
+        cpu.time = cur_time;
         cpu.gauge->Set(percent);
     }
 
     file.close();
 }
 
-void MetricsCollector::startTask()
+void MetricsCollector::StartTask()
 {
     is_task_running = true;
-    clock_gettime(CLOCK_MONOTONIC, &task_start);
+    task_start = std::chrono::high_resolution_clock::now();
 }
 
-void MetricsCollector::stopTask()
+void MetricsCollector::StopTask()
 {
     is_task_running = false;
     task_processing_time_gauge->Set(0);
     gateway.PushAdd();
-}
+}

src/metrics-collector.hpp

@@ -1,36 +1,43 @@
 #include <prometheus/gauge.h>
 #include <prometheus/registry.h>
 #include <prometheus/gateway.h>
+#include <chrono>
 
 class MetricsCollector {
     prometheus::Gateway gateway;
     std::shared_ptr<prometheus::Registry> registry;
 
     struct CPUInfo {
         prometheus::Gauge *gauge;
-        uint64_t last_total_user;
-        uint64_t last_total_user_low;
-        uint64_t last_total_sys;
-        uint64_t last_total_idle;
+        
+        struct Time {
+            uint64_t user;
+            uint64_t user_low;
+            uint64_t sys;
+            uint64_t idle;
+        };
+
+        Time time;
     };
 
-    prometheus::Gauge *memory_used_gauge;
     std::unordered_map<std::string, CPUInfo> cpu_usage;
+    
+    prometheus::Gauge *memory_used_gauge;
     prometheus::Gauge *task_processing_time_gauge;
-    bool is_running;
+    
+    std::atomic<bool> is_running {true};
     std::thread thread;
 
     bool is_task_running;
-    struct timespec task_start;
+    std::chrono::time_point<std::chrono::high_resolution_clock> task_start;
 
-    void getMemoryUsed();
-    void getCPUUsage();
-    void mainLoop();
+    void GetCPUUsage();
+    void MainLoop();
 
 public:
     MetricsCollector(const char *gateway_address, const char *gateway_port, const char *worker_name);
     ~MetricsCollector();
 
-    void startTask();
-    void stopTask();
-};
+    void StartTask();
+    void StopTask();
+};