Make SlotImpl detachable from its owner, and add a new runOnAllThreads interface to Slot.

include/envoy/thread_local/thread_local.h

@@ -74,6 +74,17 @@ class Slot {
    */
   using InitializeCb = std::function<ThreadLocalObjectSharedPtr(Event::Dispatcher& dispatcher)>;
   virtual void set(InitializeCb cb) PURE;
+
+  /**
+   * UpdateCb takes the current stored data, and returns an updated/new version data.
+   * TLS will run the callback and replace the stored data with the returned value *in each thread*.
+   *
+   * NOTE: The update callback is not supposed to capture the Slot, or its owner. As the owner may
+   * be destructed in main thread before the update_cb gets called in a worker thread.
+   **/
+  using UpdateCb = std::function<ThreadLocalObjectSharedPtr(ThreadLocalObjectSharedPtr)>;
+  virtual void runOnAllThreads(const UpdateCb& update_cb) PURE;
+  virtual void runOnAllThreads(const UpdateCb& update_cb, Event::PostCb complete_cb) PURE;
 };
 
 using SlotPtr = std::unique_ptr<Slot>;

source/common/common/non_copyable.h

@@ -2,14 +2,19 @@
 
 namespace Envoy {
 /**
- * Mixin class that makes derived classes not copyable. Like boost::noncopyable without boost.
+ * Mixin class that makes derived classes not copyable and not moveable. Like boost::noncopyable
+ * without boost.
  */
 class NonCopyable {
 protected:
   NonCopyable() = default;
 
-private:
-  NonCopyable(const NonCopyable&);
-  NonCopyable& operator=(const NonCopyable&);
+  // Non-moveable.
+  NonCopyable(NonCopyable&&) noexcept = delete;
+  NonCopyable& operator=(NonCopyable&&) noexcept = delete;
+
+  // Non-copyable.
+  NonCopyable(const NonCopyable&) = delete;
+  NonCopyable& operator=(const NonCopyable&) = delete;
 };
 } // namespace Envoy

source/common/config/config_provider_impl.cc

@@ -23,6 +23,16 @@ ConfigSubscriptionCommonBase::~ConfigSubscriptionCommonBase() {
   init_target_.ready();
   config_provider_manager_.unbindSubscription(manager_identifier_);
 }
+
+void ConfigSubscriptionCommonBase::applyConfigUpdate(const ConfigUpdateCb& update_fn) {
+  tls_->runOnAllThreads([update_fn](ThreadLocal::ThreadLocalObjectSharedPtr previous)
+                            -> ThreadLocal::ThreadLocalObjectSharedPtr {
+    auto prev_thread_local_config = std::dynamic_pointer_cast<ThreadLocalConfig>(previous);
+    prev_thread_local_config->config_ = update_fn(prev_thread_local_config->config_);
+    return previous;
+  });
+}
+
 bool ConfigSubscriptionInstance::checkAndApplyConfigUpdate(const Protobuf::Message& config_proto,
                                                            const std::string& config_name,
                                                            const std::string& version_info) {

source/common/config/config_provider_impl.h

@@ -216,19 +216,8 @@ class ConfigSubscriptionCommonBase : protected Logger::Loggable<Logger::Id::conf
    *
    * @param update_fn the callback to run on each thread, it takes the previous version Config and
    * returns a updated/new version Config.
-   * @param complete_cb the callback to run when the update propagation is done.
    */
-  void applyConfigUpdate(
-      const ConfigUpdateCb& update_fn, const Event::PostCb& complete_cb = []() {}) {
-    tls_->runOnAllThreads(
-        [this, update_fn]() {
-          // NOTE: there is a known race condition between *this* subscription being teared down in
-          // main thread and the posted callback being executed before the destruction. See more
-          // details in https://github.com/envoyproxy/envoy/issues/7902
-          tls_->getTyped<ThreadLocalConfig>().config_ = update_fn(getConfig());
-        },
-        complete_cb);
-  }
+  void applyConfigUpdate(const ConfigUpdateCb& update_fn);
 
   void setLastUpdated() { last_updated_ = time_source_.systemTime(); }
 

source/common/router/rds_impl.cc

@@ -199,8 +199,12 @@ Router::ConfigConstSharedPtr RdsRouteConfigProviderImpl::config() {
 void RdsRouteConfigProviderImpl::onConfigUpdate() {
   ConfigConstSharedPtr new_config(
       new ConfigImpl(config_update_info_->routeConfiguration(), factory_context_, false));
-  tls_->runOnAllThreads(
-      [this, new_config]() -> void { tls_->getTyped<ThreadLocalConfig>().config_ = new_config; });
+  tls_->runOnAllThreads([new_config](ThreadLocal::ThreadLocalObjectSharedPtr previous)
+                            -> ThreadLocal::ThreadLocalObjectSharedPtr {
+    auto prev_config = std::dynamic_pointer_cast<ThreadLocalConfig>(previous);
+    prev_config->config_ = new_config;
+    return previous;
+  });
 }
 
 void RdsRouteConfigProviderImpl::validateConfig(

source/common/thread_local/thread_local_impl.cc

@@ -27,8 +27,9 @@ SlotPtr InstanceImpl::allocateSlot() {
 
   if (free_slot_indexes_.empty()) {
     std::unique_ptr<SlotImpl> slot(new SlotImpl(*this, slots_.size()));
-    slots_.push_back(slot.get());
-    return slot;
+    auto wrapper = std::make_unique<Bookkeeper>(*this, std::move(slot));
+    slots_.push_back(&wrapper->slot());
+    return wrapper;
   }
   const uint32_t idx = free_slot_indexes_.front();
   free_slot_indexes_.pop_front();
@@ -42,11 +43,59 @@ bool InstanceImpl::SlotImpl::currentThreadRegistered() {
   return thread_local_data_.data_.size() > index_;
 }
 
+void InstanceImpl::SlotImpl::runOnAllThreads(const UpdateCb& cb) {
+  parent_.runOnAllThreads([this, cb]() { setThreadLocal(index_, cb(get())); });
+}
+
+void InstanceImpl::SlotImpl::runOnAllThreads(const UpdateCb& cb, Event::PostCb complete_cb) {
+  parent_.runOnAllThreads([this, cb]() { setThreadLocal(index_, cb(get())); }, complete_cb);
+}
+
 ThreadLocalObjectSharedPtr InstanceImpl::SlotImpl::get() {
   ASSERT(currentThreadRegistered());
   return thread_local_data_.data_[index_];
 }
 
+InstanceImpl::Bookkeeper::Bookkeeper(InstanceImpl& parent, std::unique_ptr<SlotImpl>&& slot)
+    : parent_(parent), holder_(std::make_unique<SlotHolder>(std::move(slot))) {}
+
+ThreadLocalObjectSharedPtr InstanceImpl::Bookkeeper::get() { return slot().get(); }
+
+void InstanceImpl::Bookkeeper::runOnAllThreads(const UpdateCb& cb, Event::PostCb complete_cb) {
+  slot().runOnAllThreads(
+      [cb, ref_count = holder_->ref_count_](ThreadLocalObjectSharedPtr previous) {
+        return cb(std::move(previous));
+      },
+      complete_cb);
+}
+
+void InstanceImpl::Bookkeeper::runOnAllThreads(const UpdateCb& cb) {
+  slot().runOnAllThreads(
+      [cb, ref_count = holder_->ref_count_](ThreadLocalObjectSharedPtr previous) {
+        return cb(std::move(previous));
+      });
+}
+
+bool InstanceImpl::Bookkeeper::currentThreadRegistered() {
+  return slot().currentThreadRegistered();
+}
+
+void InstanceImpl::Bookkeeper::runOnAllThreads(Event::PostCb cb) {
+  // Use holder_.ref_count_ to bookkeep how many on-the-fly callback are out there.
+  slot().runOnAllThreads([cb, ref_count = holder_->ref_count_]() { cb(); });
+}
+
+void InstanceImpl::Bookkeeper::runOnAllThreads(Event::PostCb cb, Event::PostCb main_callback) {
+  // Use holder_.ref_count_ to bookkeep how many on-the-fly callback are out there.
+  slot().runOnAllThreads([cb, main_callback, ref_count = holder_->ref_count_]() { cb(); },
+                         main_callback);
+}
+
+void InstanceImpl::Bookkeeper::set(InitializeCb cb) {
+  slot().set([cb, ref_count = holder_->ref_count_](Event::Dispatcher& dispatcher)
+                 -> ThreadLocalObjectSharedPtr { return cb(dispatcher); });
+}
+
 void InstanceImpl::registerThread(Event::Dispatcher& dispatcher, bool main_thread) {
   ASSERT(std::this_thread::get_id() == main_thread_id_);
   ASSERT(!shutdown_);
@@ -61,6 +110,40 @@ void InstanceImpl::registerThread(Event::Dispatcher& dispatcher, bool main_threa
   }
 }
 
+// Deletes a Slot if it's recycleable(no on-the-fly callbacks points to it), otherwise puts it into
+// a deferred delete queue, and schedules a cleanup task to collect the Slot once it's recycleable.
+void InstanceImpl::recycle(std::unique_ptr<SlotHolder>&& holder) {
+  if (holder->isRecycleable()) {
+    holder.reset();
+    return;
+  }
+  deferred_deletes_.emplace_back(std::move(holder));
+  scheduleCleanup();
+}
+
+// Cleans up the deferred deletes queue.
+// Removes all the items with ref count 1, and reschedule another cleanup task if there are left
+// un-recycleable items.
+// This way ensures that, unless server shuts down, an enqueued item will be
+// deleted eventually.
+void InstanceImpl::scheduleCleanup() {
+  ASSERT(main_thread_dispatcher_ != nullptr);
+  if (shutdown_) {
+    return;
+  }
+  if (deferred_deletes_.empty()) {
+    return;
+  }
+  main_thread_dispatcher_->post([this]() {
+    deferred_deletes_.remove_if(
+        [](std::unique_ptr<SlotHolder>& holder) -> bool { return holder->isRecycleable(); });
+    if (!deferred_deletes_.empty()) {
+      // Reschedule another cleanup task if there are still non-recyclable slots.
+      scheduleCleanup();
+    }
+  });
+}
+
 void InstanceImpl::removeSlot(SlotImpl& slot) {
   ASSERT(std::this_thread::get_id() == main_thread_id_);
 

source/common/thread_local/thread_local_impl.h

@@ -8,14 +8,15 @@
 #include "envoy/thread_local/thread_local.h"
 
 #include "common/common/logger.h"
+#include "common/common/non_copyable.h"
 
 namespace Envoy {
 namespace ThreadLocal {
 
 /**
  * Implementation of ThreadLocal that relies on static thread_local objects.
  */
-class InstanceImpl : Logger::Loggable<Logger::Id::main>, public Instance {
+class InstanceImpl : Logger::Loggable<Logger::Id::main>, public NonCopyable, public Instance {
 public:
   InstanceImpl() : main_thread_id_(std::this_thread::get_id()) {}
   ~InstanceImpl() override;
@@ -35,6 +36,8 @@ class InstanceImpl : Logger::Loggable<Logger::Id::main>, public Instance {
     // ThreadLocal::Slot
     ThreadLocalObjectSharedPtr get() override;
     bool currentThreadRegistered() override;
+    void runOnAllThreads(const UpdateCb& cb) override;
+    void runOnAllThreads(const UpdateCb& cb, Event::PostCb complete_cb) override;
     void runOnAllThreads(Event::PostCb cb) override { parent_.runOnAllThreads(cb); }
     void runOnAllThreads(Event::PostCb cb, Event::PostCb main_callback) override {
       parent_.runOnAllThreads(cb, main_callback);
@@ -45,17 +48,55 @@ class InstanceImpl : Logger::Loggable<Logger::Id::main>, public Instance {
     const uint64_t index_;
   };
 
+  // A helper class for holding a SlotImpl and its bookkeeping shared_ptr which counts the number of
+  // update callbacks on-the-fly.
+  struct SlotHolder {
+    SlotHolder(std::unique_ptr<SlotImpl>&& slot) : slot_(std::move(slot)) {}
+    bool isRecycleable() { return ref_count_.use_count() == 1; }
+
+    const std::unique_ptr<SlotImpl> slot_;
+    const std::shared_ptr<int> ref_count_{new int(0)};
+  };
+
+  // A Wrapper of SlotImpl which on destruction returns the SlotImpl to the deferred delete queue
+  // (detaches it).
+  struct Bookkeeper : public Slot {
+    Bookkeeper(InstanceImpl& parent, std::unique_ptr<SlotImpl>&& slot);
+    ~Bookkeeper() override { parent_.recycle(std::move(holder_)); }
+    SlotImpl& slot() { return *(holder_->slot_); }
+
+    // ThreadLocal::Slot
+    ThreadLocalObjectSharedPtr get() override;
+    void runOnAllThreads(const UpdateCb& cb) override;
+    void runOnAllThreads(const UpdateCb& cb, Event::PostCb complete_cb) override;
+    bool currentThreadRegistered() override;
+    void runOnAllThreads(Event::PostCb cb) override;
+    void runOnAllThreads(Event::PostCb cb, Event::PostCb main_callback) override;
+    void set(InitializeCb cb) override;
+
+    InstanceImpl& parent_;
+    std::unique_ptr<SlotHolder> holder_;
+  };
+
   struct ThreadLocalData {
     Event::Dispatcher* dispatcher_{};
     std::vector<ThreadLocalObjectSharedPtr> data_;
   };
 
+  void recycle(std::unique_ptr<SlotHolder>&& holder);
+  // Cleanup the deferred deletes queue.
+  void scheduleCleanup();
   void removeSlot(SlotImpl& slot);
   void runOnAllThreads(Event::PostCb cb);
   void runOnAllThreads(Event::PostCb cb, Event::PostCb main_callback);
   static void setThreadLocal(uint32_t index, ThreadLocalObjectSharedPtr object);
 
   static thread_local ThreadLocalData thread_local_data_;
+
+  // A queue for Slots that has to be deferred to delete due to out-going callbacks
+  // pointing to the Slot.
+  std::list<std::unique_ptr<SlotHolder>> deferred_deletes_;
+
   std::vector<SlotImpl*> slots_;
   // A list of index of freed slots.
   std::list<uint32_t> free_slot_indexes_;

test/common/thread_local/thread_local_impl_test.cc

@@ -85,6 +85,40 @@ TEST_F(ThreadLocalInstanceImplTest, All) {
   tls_.shutdownThread();
 }
 
+// Test that the config passed into the update callback is the previous version stored in the slot.
+TEST_F(ThreadLocalInstanceImplTest, UpdateCallback) {
+  InSequence s;
+
+  SlotPtr slot = tls_.allocateSlot();
+
+  auto newer_version = std::make_shared<TestThreadLocalObject>();
+  bool update_called = false;
+
+  TestThreadLocalObject& object_ref = setObject(*slot);
+  auto update_cb = [&object_ref, &update_called,
+                    newer_version](ThreadLocalObjectSharedPtr obj) -> ThreadLocalObjectSharedPtr {
+    // The unit test setup have two dispatchers registered, but only one thread, this lambda will be
+    // called twice in the same thread.
+    if (!update_called) {
+      EXPECT_EQ(obj.get(), &object_ref);
+      update_called = true;
+    } else {
+      EXPECT_EQ(obj.get(), newer_version.get());
+    }
+
+    return newer_version;
+  };
+  EXPECT_CALL(thread_dispatcher_, post(_));
+  EXPECT_CALL(object_ref, onDestroy());
+  EXPECT_CALL(*newer_version, onDestroy());
+  slot->runOnAllThreads(update_cb);
+
+  EXPECT_EQ(newer_version.get(), &slot->getTyped<TestThreadLocalObject>());
+
+  tls_.shutdownGlobalThreading();
+  tls_.shutdownThread();
+}
+
 // TODO(ramaraochavali): Run this test with real threads. The current issue in the unit
 // testing environment is, the post to main_dispatcher is not working as expected.
 

test/mocks/thread_local/mocks.h

@@ -63,6 +63,14 @@ class MockInstance : public Instance {
     void runOnAllThreads(Event::PostCb cb, Event::PostCb main_callback) override {
       parent_.runOnAllThreads(cb, main_callback);
     }
+    void runOnAllThreads(const UpdateCb& cb) override {
+      parent_.runOnAllThreads([cb, this]() { parent_.data_[index_] = cb(parent_.data_[index_]); });
+    }
+    void runOnAllThreads(const UpdateCb& cb, Event::PostCb main_callback) override {
+      parent_.runOnAllThreads([cb, this]() { parent_.data_[index_] = cb(parent_.data_[index_]); },
+                              main_callback);
+    }
+
     void set(InitializeCb cb) override { parent_.data_[index_] = cb(parent_.dispatcher_); }
 
     MockInstance& parent_;