From 26e9f23ada4f7dea37199bca89ce628e2a54b74e Mon Sep 17 00:00:00 2001
From: Jonas Devlieghere <jonas@devlieghere.com>
Date: Thu, 9 Feb 2023 13:00:45 -0800
Subject: [PATCH] [lldb] Add the ability to remove diagnostic callbacks

Add the ability to remove diagnostic callbacks. This is necessary for
diagnostics who's lifetime is tied to objects that can be destroyed.

Differential revision: https://reviews.llvm.org/D143548
---
 lldb/include/lldb/Utility/Diagnostics.h | 20 ++++++++++++++++++--
 lldb/source/Utility/Diagnostics.cpp     | 18 ++++++++++++++----
 2 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/lldb/include/lldb/Utility/Diagnostics.h b/lldb/include/lldb/Utility/Diagnostics.h
index 5d23fa8..c2e4b44 100644
--- a/lldb/include/lldb/Utility/Diagnostics.h
+++ b/lldb/include/lldb/Utility/Diagnostics.h
@@ -42,8 +42,10 @@ public:
   void Report(llvm::StringRef message);
 
   using Callback = std::function<llvm::Error(const FileSpec &)>;
+  using CallbackID = uint64_t;
 
-  void AddCallback(Callback callback);
+  CallbackID AddCallback(Callback callback);
+  void RemoveCallback(CallbackID id);
 
   static Diagnostics &Instance();
 
@@ -61,7 +63,21 @@ private:
 
   RotatingLogHandler m_log_handler;
 
-  llvm::SmallVector<Callback, 4> m_callbacks;
+  struct CallbackEntry {
+    CallbackEntry(CallbackID id, Callback callback)
+        : id(id), callback(std::move(callback)) {}
+    CallbackID id;
+    Callback callback;
+  };
+
+  /// Monotonically increasing callback identifier. Unique per Diagnostic
+  /// instance.
+  CallbackID m_callback_id;
+
+  /// List of callback entries.
+  llvm::SmallVector<CallbackEntry, 4> m_callbacks;
+
+  /// Mutex to protect callback list and callback identifier.
   std::mutex m_callbacks_mutex;
 };
 
diff --git a/lldb/source/Utility/Diagnostics.cpp b/lldb/source/Utility/Diagnostics.cpp
index a597501..1632ae0 100644
--- a/lldb/source/Utility/Diagnostics.cpp
+++ b/lldb/source/Utility/Diagnostics.cpp
@@ -43,9 +43,19 @@ Diagnostics::Diagnostics() : m_log_handler(g_num_log_messages) {}
 
 Diagnostics::~Diagnostics() {}
 
-void Diagnostics::AddCallback(Callback callback) {
+Diagnostics::CallbackID Diagnostics::AddCallback(Callback callback) {
   std::lock_guard<std::mutex> guard(m_callbacks_mutex);
-  m_callbacks.push_back(callback);
+  CallbackID id = m_callback_id++;
+  m_callbacks.emplace_back(id, callback);
+  return id;
+}
+
+void Diagnostics::RemoveCallback(CallbackID id) {
+  std::lock_guard<std::mutex> guard(m_callbacks_mutex);
+  m_callbacks.erase(
+      std::remove_if(m_callbacks.begin(), m_callbacks.end(),
+                     [id](const CallbackEntry &e) { return e.id == id; }),
+      m_callbacks.end());
 }
 
 bool Diagnostics::Dump(raw_ostream &stream) {
@@ -84,8 +94,8 @@ Error Diagnostics::Create(const FileSpec &dir) {
   if (Error err = DumpDiangosticsLog(dir))
     return err;
 
-  for (Callback c : m_callbacks) {
-    if (Error err = c(dir))
+  for (CallbackEntry e : m_callbacks) {
+    if (Error err = e.callback(dir))
       return err;
   }
 
-- 
2.7.4