tls: add the missing wrapper around SlotImpl allocation. (envoyproxy#…

…8290)

Signed-off-by: Xin Zhuang <[email protected]>

source/common/thread_local/thread_local_impl.cc

@@ -36,7 +36,7 @@ SlotPtr InstanceImpl::allocateSlot() {
   ASSERT(idx < slots_.size());
   std::unique_ptr<SlotImpl> slot(new SlotImpl(*this, idx));
   slots_[idx] = slot.get();
-  return slot;
+  return std::make_unique<Bookkeeper>(*this, std::move(slot));
 }
 
 bool InstanceImpl::SlotImpl::currentThreadRegistered() {

source/common/thread_local/thread_local_impl.h

@@ -99,6 +99,9 @@ class InstanceImpl : Logger::Loggable<Logger::Id::main>, public NonCopyable, pub
   std::thread::id main_thread_id_;
   Event::Dispatcher* main_thread_dispatcher_{};
   std::atomic<bool> shutdown_{};
+
+  // Test only.
+  friend class ThreadLocalInstanceImplTest;
 };
 
 } // namespace ThreadLocal

test/common/thread_local/thread_local_impl_test.cc

@@ -14,7 +14,6 @@ using testing::ReturnPointee;
 
 namespace Envoy {
 namespace ThreadLocal {
-namespace {
 
 class TestThreadLocalObject : public ThreadLocalObject {
 public:
@@ -46,8 +45,10 @@ class ThreadLocalInstanceImplTest : public testing::Test {
     object.reset();
     return object_ref;
   }
-
+  int deferredDeletesMapSize() { return tls_.deferred_deletes_.size(); }
+  int freeSlotIndexesListSize() { return tls_.free_slot_indexes_.size(); }
   InstanceImpl tls_;
+
   Event::MockDispatcher main_dispatcher_;
   Event::MockDispatcher thread_dispatcher_;
 };
@@ -59,15 +60,20 @@ TEST_F(ThreadLocalInstanceImplTest, All) {
   EXPECT_CALL(thread_dispatcher_, post(_));
   SlotPtr slot1 = tls_.allocateSlot();
   slot1.reset();
+  EXPECT_EQ(deferredDeletesMapSize(), 0);
+  EXPECT_EQ(freeSlotIndexesListSize(), 1);
 
   // Create a new slot which should take the place of the old slot. ReturnPointee() is used to
   // avoid "leaks" when using InSequence and shared_ptr.
   SlotPtr slot2 = tls_.allocateSlot();
   TestThreadLocalObject& object_ref2 = setObject(*slot2);
+  EXPECT_EQ(freeSlotIndexesListSize(), 0);
 
   EXPECT_CALL(thread_dispatcher_, post(_));
   EXPECT_CALL(object_ref2, onDestroy());
+  EXPECT_EQ(freeSlotIndexesListSize(), 0);
   slot2.reset();
+  EXPECT_EQ(freeSlotIndexesListSize(), 1);
 
   // Make two new slots, shutdown global threading, and delete them. We should not see any
   // cross-thread posts at this point. We should also see destruction in reverse order.
@@ -79,12 +85,53 @@ TEST_F(ThreadLocalInstanceImplTest, All) {
   tls_.shutdownGlobalThreading();
   slot3.reset();
   slot4.reset();
+  EXPECT_EQ(freeSlotIndexesListSize(), 0);
+  EXPECT_EQ(deferredDeletesMapSize(), 2);
 
   EXPECT_CALL(object_ref4, onDestroy());
   EXPECT_CALL(object_ref3, onDestroy());
   tls_.shutdownThread();
 }
 
+TEST_F(ThreadLocalInstanceImplTest, DeferredRecycle) {
+  InSequence s;
+
+  // Free a slot without ever calling set.
+  EXPECT_CALL(thread_dispatcher_, post(_));
+  SlotPtr slot1 = tls_.allocateSlot();
+  slot1.reset();
+  // Slot destructed directly, as there is no out-going callbacks.
+  EXPECT_EQ(deferredDeletesMapSize(), 0);
+  EXPECT_EQ(freeSlotIndexesListSize(), 1);
+
+  // Allocate a slot and set value, hold the posted callback and the slot will only be returned
+  // after the held callback is destructed.
+  {
+    SlotPtr slot2 = tls_.allocateSlot();
+    EXPECT_EQ(freeSlotIndexesListSize(), 0);
+    {
+      Event::PostCb holder;
+      EXPECT_CALL(thread_dispatcher_, post(_)).WillOnce(Invoke([&](Event::PostCb cb) {
+        // Holds the posted callback.
+        holder = cb;
+      }));
+      slot2->set(
+          [](Event::Dispatcher&) -> ThreadLocal::ThreadLocalObjectSharedPtr { return nullptr; });
+      slot2.reset();
+      // Not released yet, as holder has a copy of the ref_count_.
+      EXPECT_EQ(freeSlotIndexesListSize(), 0);
+      EXPECT_EQ(deferredDeletesMapSize(), 1);
+      // This post is called when the holder dies.
+      EXPECT_CALL(thread_dispatcher_, post(_));
+    }
+    // Slot is deleted now that there holder destructs.
+    EXPECT_EQ(deferredDeletesMapSize(), 0);
+    EXPECT_EQ(freeSlotIndexesListSize(), 1);
+  }
+
+  tls_.shutdownGlobalThreading();
+}
+
 // Test that the config passed into the update callback is the previous version stored in the slot.
 TEST_F(ThreadLocalInstanceImplTest, UpdateCallback) {
   InSequence s;
@@ -179,6 +226,5 @@ TEST(ThreadLocalInstanceImplDispatcherTest, Dispatcher) {
   tls.shutdownThread();
 }
 
-} // namespace
 } // namespace ThreadLocal
 } // namespace Envoy