pull: Add RemoveCollectable

Closes: #415
jupp0r · Dec 17, 2020 · bd887da · bd887da
1 parent 4061c10
commit bd887da
Show file tree

Hide file tree

Showing 7 changed files with 42 additions and 0 deletions.
diff --git a/pull/include/prometheus/exposer.h b/pull/include/prometheus/exposer.h
@@ -33,6 +33,9 @@ class PROMETHEUS_CPP_PULL_EXPORT Exposer {
       const std::string& realm = "Prometheus-cpp Exporter",
       const std::string& uri = std::string("/metrics"));
 
+  void RemoveCollectable(const std::weak_ptr<Collectable>& collectable,
+                         const std::string& uri = std::string("/metrics"));
+
   std::vector<int> GetListeningPorts() const;
 
  private:

diff --git a/pull/src/endpoint.cc b/pull/src/endpoint.cc
@@ -41,6 +41,11 @@ void Endpoint::RegisterAuth(
   auth_handler_ = std::move(new_handler);
 }
 
+void Endpoint::RemoveCollectable(
+    const std::weak_ptr<Collectable>& collectable) {
+  metrics_handler_->RemoveCollectable(collectable);
+}
+
 const std::string& Endpoint::GetURI() const { return uri_; }
 
 }  // namespace detail

diff --git a/pull/src/endpoint.h b/pull/src/endpoint.h
@@ -24,6 +24,7 @@ class Endpoint {
   void RegisterAuth(
       std::function<bool(const std::string&, const std::string&)> authCB,
       const std::string& realm);
+  void RemoveCollectable(const std::weak_ptr<Collectable>& collectable);
 
   const std::string& GetURI() const;
 

diff --git a/pull/src/exposer.cc b/pull/src/exposer.cc
@@ -35,6 +35,12 @@ void Exposer::RegisterAuth(
   endpoint.RegisterAuth(std::move(authCB), realm);
 }
 
+void Exposer::RemoveCollectable(const std::weak_ptr<Collectable>& collectable,
+                                const std::string& uri) {
+  auto& endpoint = GetEndpointForUri(uri);
+  endpoint.RemoveCollectable(collectable);
+}
+
 std::vector<int> Exposer::GetListeningPorts() const {
   return server_->getListeningPorts();
 }

diff --git a/pull/src/handler.cc b/pull/src/handler.cc
@@ -121,6 +121,20 @@ void MetricsHandler::RegisterCollectable(
   collectables_.push_back(collectable);
 }
 
+void MetricsHandler::RemoveCollectable(
+    const std::weak_ptr<Collectable>& collectable) {
+  std::lock_guard<std::mutex> lock{collectables_mutex_};
+
+  auto locked = collectable.lock();
+  auto same_pointer = [&locked](const std::weak_ptr<Collectable>& candidate) {
+    return locked == candidate.lock();
+  };
+
+  collectables_.erase(std::remove_if(std::begin(collectables_),
+                                     std::end(collectables_), same_pointer),
+                      std::end(collectables_));
+}
+
 bool MetricsHandler::handleGet(CivetServer*, struct mg_connection* conn) {
   auto start_time_of_request = std::chrono::steady_clock::now();
 

diff --git a/pull/src/handler.h b/pull/src/handler.h
@@ -16,6 +16,7 @@ class MetricsHandler : public CivetHandler {
   explicit MetricsHandler(Registry& registry);
 
   void RegisterCollectable(const std::weak_ptr<Collectable>& collectable);
+  void RemoveCollectable(const std::weak_ptr<Collectable>& collectable);
 
   bool handleGet(CivetServer* server, struct mg_connection* conn) override;
 

diff --git a/pull/tests/integration/integration_test.cc b/pull/tests/integration/integration_test.cc
@@ -122,5 +122,17 @@ TEST_F(IntegrationTest, exposesCountersOnDifferentUrls) {
   EXPECT_THAT(second_metrics.body, Not(HasSubstr(first_counter_name)));
 }
 
+TEST_F(IntegrationTest, unexposeRegistry) {
+  const std::string counter_name = "some_counter_total";
+  const auto registry =
+      RegisterSomeCounter(counter_name, default_metrics_path_);
+
+  exposer_->RemoveCollectable(registry, default_metrics_path_);
+
+  const auto metrics = FetchMetrics(default_metrics_path_);
+  ASSERT_EQ(metrics.code, 200);
+  EXPECT_THAT(metrics.body, Not(HasSubstr(counter_name)));
+}
+
 }  // namespace
 }  // namespace prometheus