Process metric callbacks from the main-loop thread

This avoids the callbacks from being processed on the worker thread
spawned by Civetweb. It fixes data race issues with lookups involving
global variables, amongst other threading issues.

auxil/prometheus-cpp

@@ -1 +1 @@
-Subproject commit 2fec7205d1a9cb4829b86c943d599696d53de85c
+Subproject commit 4649065e2a1dd21c81e41cd6007dce5486b77fc0

scripts/base/init-bare.zeek

@@ -5883,6 +5883,13 @@ export {
 
 	type MetricVector : vector of Metric;
 	type HistogramMetricVector : vector of HistogramMetric;
+
+	## Maximum amount of time for CivetWeb HTTP threads to
+	## wait for metric callbacks to complete on the IO loop.
+	const callback_timeout: interval = 5sec &redef;
+
+	## Number of CivetWeb threads to use.
+	const civetweb_threads: count = 2 &redef;
 }
 
 module GLOBAL;

src/telemetry/CMakeLists.txt

@@ -9,6 +9,7 @@ zeek_add_subdir_library(
     ProcessStats.cc
     Utils.cc
     BIFS
+    consts.bif
     telemetry.bif)
 
 # We don't need to include the civetweb headers across the whole project, only

src/telemetry/Manager.cc

@@ -6,6 +6,7 @@
 
 // CivetServer is from the civetweb submodule in prometheus-cpp
 #include <CivetServer.h>
+#include <prometheus/collectable.h>
 #include <prometheus/exposer.h>
 #include <prometheus/registry.h>
 #include <rapidjson/document.h>
@@ -16,19 +17,32 @@
 
 #include "zeek/3rdparty/doctest.h"
 #include "zeek/ID.h"
+#include "zeek/RunState.h"
 #include "zeek/ZeekString.h"
 #include "zeek/broker/Manager.h"
+#include "zeek/iosource/Manager.h"
 #include "zeek/telemetry/ProcessStats.h"
 #include "zeek/telemetry/Timer.h"
+#include "zeek/telemetry/consts.bif.h"
 #include "zeek/telemetry/telemetry.bif.h"
 #include "zeek/threading/formatters/detail/json.h"
 
 namespace zeek::telemetry {
 
-Manager::Manager() { prometheus_registry = std::make_shared<prometheus::Registry>(); }
+/**
+ * Prometheus Collectable interface used to insert Zeek callback processing
+ * before the Prometheus registry's collection of metric data.
+ */
+class ZeekCollectable : public prometheus::Collectable {
+public:
+    std::vector<prometheus::MetricFamily> Collect() const override {
+        telemetry_mgr->WaitForPrometheusCallbacks();
+        return {};
+    }
+};
+
+Manager::Manager() : IOSource(true) { prometheus_registry = std::make_shared<prometheus::Registry>(); }
 
-// This can't be defined as =default because of the use of unique_ptr with a forward-declared type
-// in Manager.h
 Manager::~Manager() {}
 
 void Manager::InitPostScript() {
@@ -75,7 +89,9 @@ void Manager::InitPostScript() {
 
         if ( ! getenv("ZEEKCTL_CHECK_CONFIG") ) {
             try {
-                prometheus_exposer = std::make_unique<prometheus::Exposer>(prometheus_url, 2, callbacks);
+                prometheus_exposer =
+                    std::make_unique<prometheus::Exposer>(prometheus_url, BifConst::Telemetry::civetweb_threads,
+                                                          callbacks);
 
                 // CivetWeb stores a copy of the callbacks, so we're safe to delete the pointer here
                 delete callbacks;
@@ -84,6 +100,13 @@ void Manager::InitPostScript() {
                                      prometheus_url.c_str());
             }
 
+            // This has to be inserted before the registry below. The exposer
+            // processes the collectors in order of insertion. We want to make
+            // sure that the callbacks get called and the values in the metrics
+            // are updated before prometheus-cpp scrapes them.
+            zeek_collectable = std::make_shared<ZeekCollectable>();
+            prometheus_exposer->RegisterCollectable(zeek_collectable);
+
             prometheus_exposer->RegisterCollectable(prometheus_registry);
         }
     }
@@ -130,6 +153,21 @@ void Manager::InitPostScript() {
                                   return metric;
                               });
 #endif
+
+    iosource_mgr->RegisterFd(collector_flare.FD(), this);
+}
+
+void Manager::Terminate() {
+    // Notify the collector condition so that it doesn't hang waiting for
+    // a collector request to complete.
+    collector_cv.notify_all();
+
+    // Shut down the exposer first of all so we stop getting requests for
+    // data. This keeps us from getting a request on another thread while
+    // we're shutting down.
+    prometheus_exposer.reset();
+
+    iosource_mgr->UnregisterFd(collector_flare.FD(), this);
 }
 
 // -- collect metric stuff -----------------------------------------------------
@@ -545,6 +583,39 @@ HistogramPtr Manager::HistogramInstance(std::string_view prefix, std::string_vie
     return HistogramInstance(prefix, name, lbls, bounds_span, helptext, unit);
 }
 
+void Manager::ProcessFd(int fd, int flags) {
+    std::unique_lock<std::mutex> lk(collector_cv_mtx);
+
+    collector_flare.Extinguish();
+
+    prometheus_registry->UpdateViaCallbacks();
+    collector_response_idx = collector_request_idx;
+
+    lk.unlock();
+    collector_cv.notify_all();
+}
+
+void Manager::WaitForPrometheusCallbacks() {
+    std::unique_lock<std::mutex> lk(collector_cv_mtx);
+
+    ++collector_request_idx;
+    uint64_t expected_idx = collector_request_idx;
+    collector_flare.Fire();
+
+    // It should *not* take 5 seconds to go through all of the callbacks, but
+    // set this to have a timeout anyways just to avoid a deadlock.
+    bool res = collector_cv.wait_for(lk,
+                                     std::chrono::microseconds(
+                                         static_cast<long>(BifConst::Telemetry::callback_timeout * 1000000)),
+                                     [expected_idx]() {
+                                         return telemetry_mgr->collector_response_idx >= expected_idx ||
+                                                zeek::run_state::terminating;
+                                     });
+
+    if ( ! res )
+        fprintf(stderr, "Timeout waiting for prometheus callbacks\n");
+}
+
 } // namespace zeek::telemetry
 
 // -- unit tests ---------------------------------------------------------------

src/telemetry/Manager.h

@@ -9,8 +9,10 @@
 #include <string_view>
 #include <vector>
 
+#include "zeek/Flare.h"
 #include "zeek/IntrusivePtr.h"
 #include "zeek/Span.h"
+#include "zeek/iosource/IOSource.h"
 #include "zeek/telemetry/Counter.h"
 #include "zeek/telemetry/Gauge.h"
 #include "zeek/telemetry/Histogram.h"
@@ -29,15 +31,16 @@ class Registry;
 
 namespace zeek::telemetry {
 
+class ZeekCollectable;
+
 /**
  * Manages a collection of metric families.
  */
-class Manager final {
+class Manager final : public iosource::IOSource {
 public:
     Manager();
 
     Manager(const Manager&) = delete;
-
     Manager& operator=(const Manager&) = delete;
 
     ~Manager();
@@ -50,6 +53,8 @@ public:
      */
     void InitPostScript();
 
+    void Terminate();
+
     /**
      * @return A VectorVal containing all counter and gauge metrics and their values matching prefix and name.
      * @param prefix The prefix pattern to use for filtering. Supports globbing.
@@ -88,8 +93,8 @@ public:
      * @param labels Values for all label dimensions of the metric.
      * @param helptext Short explanation of the metric.
      * @param unit Unit of measurement.
-     * @param callback Passing a callback method will enable asynchronous mode. The callback method will be called by
+     * @param callback Passing a callback method will enable asynchronous mode. The callback method will be called
-     * the metrics subsystem whenever data is requested.
+     * by the metrics subsystem whenever data is requested.
      */
     CounterPtr CounterInstance(std::string_view prefix, std::string_view name, Span<const LabelView> labels,
                                std::string_view helptext, std::string_view unit = "",
@@ -124,8 +129,8 @@ public:
      * @param labels Values for all label dimensions of the metric.
      * @param helptext Short explanation of the metric.
      * @param unit Unit of measurement.
-     * @param callback Passing a callback method will enable asynchronous mode. The callback method will be called by
+     * @param callback Passing a callback method will enable asynchronous mode. The callback method will be called
-     * the metrics subsystem whenever data is requested.
+     * by the metrics subsystem whenever data is requested.
      */
     GaugePtr GaugeInstance(std::string_view prefix, std::string_view name, Span<const LabelView> labels,
                            std::string_view helptext, std::string_view unit = "",
@@ -212,6 +217,12 @@ public:
      */
     std::shared_ptr<prometheus::Registry> GetRegistry() const { return prometheus_registry; }
 
+    // IOSource interface
+    double GetNextTimeout() override { return -1.0; }
+    void Process() override {}
+    const char* Tag() override { return "Telemetry::Manager"; }
+    void ProcessFd(int fd, int flags) override;
+
 protected:
     template<class F>
     static auto WithLabelNames(Span<const LabelView> xs, F continuation) {
@@ -231,6 +242,15 @@ protected:
         }
     }
 
+    friend class ZeekCollectable;
+
+    /**
+     * Fires the flare for prometheus-cpp callback handling and waits for it to complete.
+     * This can be called from other threads to ensure the callback handling happens on
+     * the main thread.
+     */
+    void WaitForPrometheusCallbacks();
+
 private:
     RecordValPtr GetMetricOptsRecord(const prometheus::MetricFamily& metric_family);
     void BuildClusterJson();
@@ -250,6 +270,14 @@ private:
     std::unique_ptr<prometheus::Exposer> prometheus_exposer;
 
     std::string cluster_json;
+
+    std::shared_ptr<ZeekCollectable> zeek_collectable;
+    zeek::detail::Flare collector_flare;
+    std::condition_variable collector_cv;
+    std::mutex collector_cv_mtx;
+    // Only modified under collector_cv_mtx!
+    uint64_t collector_request_idx = 0;
+    uint64_t collector_response_idx = 0;
 };
 
 } // namespace zeek::telemetry

src/telemetry/consts.bif

@@ -0,0 +1,2 @@
+const Telemetry::callback_timeout: interval;
+const Telemetry::civetweb_threads: count;

src/zeek-setup.cc

@@ -376,6 +376,7 @@ static void terminate_zeek() {
     input_mgr->Terminate();
     thread_mgr->Terminate();
     broker_mgr->Terminate();
+    telemetry_mgr->Terminate();
 
     event_mgr.Drain();
 
@@ -716,6 +717,7 @@ SetupResult setup(int argc, char** argv, Options* zopts) {
     // when that variable is defined.
     auto early_shutdown = [] {
         broker_mgr->Terminate();
+        telemetry_mgr->Terminate();
         delete iosource_mgr;
         delete telemetry_mgr;
     };

testing/btest/Baseline/coverage.bare-load-baseline/canonified_loaded_scripts.log

@@ -146,6 +146,7 @@ scripts/base/init-frameworks-and-bifs.zeek
     scripts/base/frameworks/files/magic/__load__.zeek
   scripts/base/frameworks/telemetry/options.zeek
   build/scripts/base/bif/__load__.zeek
+    build/scripts/base/bif/consts.bif.zeek
     build/scripts/base/bif/telemetry.bif.zeek
     build/scripts/base/bif/zeekygen.bif.zeek
     build/scripts/base/bif/pcap.bif.zeek

testing/btest/Baseline/coverage.default-load-baseline/canonified_loaded_scripts.log

@@ -146,6 +146,7 @@ scripts/base/init-frameworks-and-bifs.zeek
     scripts/base/frameworks/files/magic/__load__.zeek
   scripts/base/frameworks/telemetry/options.zeek
   build/scripts/base/bif/__load__.zeek
+    build/scripts/base/bif/consts.bif.zeek
     build/scripts/base/bif/telemetry.bif.zeek
     build/scripts/base/bif/zeekygen.bif.zeek
     build/scripts/base/bif/pcap.bif.zeek

testing/btest/Baseline/plugins.hooks/output

@@ -464,6 +464,7 @@
 0.000000   MetaHookPost  LoadFile(0, ./comm.bif.zeek, <...>/comm.bif.zeek) -> -1
 0.000000   MetaHookPost  LoadFile(0, ./communityid.bif.zeek, <...>/communityid.bif.zeek) -> -1
 0.000000   MetaHookPost  LoadFile(0, ./const.bif.zeek, <...>/const.bif.zeek) -> -1
+0.000000   MetaHookPost  LoadFile(0, ./consts.bif.zeek, <...>/consts.bif.zeek) -> -1
 0.000000   MetaHookPost  LoadFile(0, ./contents, <...>/contents.zeek) -> -1
 0.000000   MetaHookPost  LoadFile(0, ./control, <...>/control.zeek) -> -1
 0.000000   MetaHookPost  LoadFile(0, ./data.bif.zeek, <...>/data.bif.zeek) -> -1
@@ -758,6 +759,7 @@
 0.000000   MetaHookPost  LoadFileExtended(0, ./comm.bif.zeek, <...>/comm.bif.zeek) -> (-1, <no content>)
 0.000000   MetaHookPost  LoadFileExtended(0, ./communityid.bif.zeek, <...>/communityid.bif.zeek) -> (-1, <no content>)
 0.000000   MetaHookPost  LoadFileExtended(0, ./const.bif.zeek, <...>/const.bif.zeek) -> (-1, <no content>)
+0.000000   MetaHookPost  LoadFileExtended(0, ./consts.bif.zeek, <...>/consts.bif.zeek) -> (-1, <no content>)
 0.000000   MetaHookPost  LoadFileExtended(0, ./contents, <...>/contents.zeek) -> (-1, <no content>)
 0.000000   MetaHookPost  LoadFileExtended(0, ./control, <...>/control.zeek) -> (-1, <no content>)
 0.000000   MetaHookPost  LoadFileExtended(0, ./data.bif.zeek, <...>/data.bif.zeek) -> (-1, <no content>)
@@ -1384,6 +1386,7 @@
 0.000000   MetaHookPre   LoadFile(0, ./comm.bif.zeek, <...>/comm.bif.zeek)
 0.000000   MetaHookPre   LoadFile(0, ./communityid.bif.zeek, <...>/communityid.bif.zeek)
 0.000000   MetaHookPre   LoadFile(0, ./const.bif.zeek, <...>/const.bif.zeek)
+0.000000   MetaHookPre   LoadFile(0, ./consts.bif.zeek, <...>/consts.bif.zeek)
 0.000000   MetaHookPre   LoadFile(0, ./contents, <...>/contents.zeek)
 0.000000   MetaHookPre   LoadFile(0, ./control, <...>/control.zeek)
 0.000000   MetaHookPre   LoadFile(0, ./data.bif.zeek, <...>/data.bif.zeek)
@@ -1678,6 +1681,7 @@
 0.000000   MetaHookPre   LoadFileExtended(0, ./comm.bif.zeek, <...>/comm.bif.zeek)
 0.000000   MetaHookPre   LoadFileExtended(0, ./communityid.bif.zeek, <...>/communityid.bif.zeek)
 0.000000   MetaHookPre   LoadFileExtended(0, ./const.bif.zeek, <...>/const.bif.zeek)
+0.000000   MetaHookPre   LoadFileExtended(0, ./consts.bif.zeek, <...>/consts.bif.zeek)
 0.000000   MetaHookPre   LoadFileExtended(0, ./contents, <...>/contents.zeek)
 0.000000   MetaHookPre   LoadFileExtended(0, ./control, <...>/control.zeek)
 0.000000   MetaHookPre   LoadFileExtended(0, ./data.bif.zeek, <...>/data.bif.zeek)
@@ -2305,6 +2309,7 @@
 0.000000 | HookLoadFile  ./comm.bif.zeek <...>/comm.bif.zeek
 0.000000 | HookLoadFile  ./communityid.bif.zeek <...>/communityid.bif.zeek
 0.000000 | HookLoadFile  ./const.bif.zeek <...>/const.bif.zeek
+0.000000 | HookLoadFile  ./consts.bif.zeek <...>/consts.bif.zeek
 0.000000 | HookLoadFile  ./contents <...>/contents.zeek
 0.000000 | HookLoadFile  ./control <...>/control.zeek
 0.000000 | HookLoadFile  ./data.bif.zeek <...>/data.bif.zeek
@@ -2599,6 +2604,7 @@
 0.000000 | HookLoadFileExtended ./comm.bif.zeek <...>/comm.bif.zeek
 0.000000 | HookLoadFileExtended ./communityid.bif.zeek <...>/communityid.bif.zeek
 0.000000 | HookLoadFileExtended ./const.bif.zeek <...>/const.bif.zeek
+0.000000 | HookLoadFileExtended ./consts.bif.zeek <...>/consts.bif.zeek
 0.000000 | HookLoadFileExtended ./contents <...>/contents.zeek
 0.000000 | HookLoadFileExtended ./control <...>/control.zeek
 0.000000 | HookLoadFileExtended ./data.bif.zeek <...>/data.bif.zeek