From 12a2a95d750aa5ae309ecd8374ba9f6348087495 Mon Sep 17 00:00:00 2001
From: Xin Zhuang <stevenzzz@google.com>
Date: Tue, 20 Aug 2019 17:15:57 -0400
Subject: [PATCH] maintain a free slot index set in TLS InstanceImpl to
 allocate in O(1) time

Signed-off-by: Xin Zhuang <stevenzzz@google.com>
---
 .../common/thread_local/thread_local_impl.cc  | 20 ++++++++++---------
 .../common/thread_local/thread_local_impl.h   |  5 +++++
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/source/common/thread_local/thread_local_impl.cc b/source/common/thread_local/thread_local_impl.cc
index 4e8c32fed776..d9b8469bd06d 100644
--- a/source/common/thread_local/thread_local_impl.cc
+++ b/source/common/thread_local/thread_local_impl.cc
@@ -24,16 +24,16 @@ SlotPtr InstanceImpl::allocateSlot() {
   ASSERT(std::this_thread::get_id() == main_thread_id_);
   ASSERT(!shutdown_);
 
-  for (uint64_t i = 0; i < slots_.size(); i++) {
-    if (slots_[i] == nullptr) {
-      std::unique_ptr<SlotImpl> slot(new SlotImpl(*this, i));
-      slots_[i] = slot.get();
-      return slot;
-    }
+  if (free_slot_indexes_.empty()) {
+    std::unique_ptr<SlotImpl> slot(new SlotImpl(*this, slots_.size()));
+    slots_.push_back(slot.get());
+    return slot;
   }
-
-  std::unique_ptr<SlotImpl> slot(new SlotImpl(*this, slots_.size()));
-  slots_.push_back(slot.get());
+  uint32_t idx = *free_slot_indexes_.begin();
+  free_slot_indexes_.erase(idx);
+  ASSERT(idx < slots_.size());
+  std::unique_ptr<SlotImpl> slot(new SlotImpl(*this, idx));
+  slots_[idx] = slot.get();
   return slot;
 }
 
@@ -73,6 +73,8 @@ void InstanceImpl::removeSlot(SlotImpl& slot) {
 
   const uint64_t index = slot.index_;
   slots_[index] = nullptr;
+  auto ret = free_slot_indexes_.emplace(index);
+  ASSERT(ret.second, fmt::format("slot index {} already in free slot set!", index));
   runOnAllThreads([index]() -> void {
     // This runs on each thread and clears the slot, making it available for a new allocations.
     // This is safe even if a new allocation comes in, because everything happens with post() and
diff --git a/source/common/thread_local/thread_local_impl.h b/source/common/thread_local/thread_local_impl.h
index 3e8e39c8fa89..73b5433134d5 100644
--- a/source/common/thread_local/thread_local_impl.h
+++ b/source/common/thread_local/thread_local_impl.h
@@ -9,6 +9,8 @@
 
 #include "common/common/logger.h"
 
+#include "absl/container/flat_hash_set.h"
+
 namespace Envoy {
 namespace ThreadLocal {
 
@@ -57,6 +59,9 @@ class InstanceImpl : Logger::Loggable<Logger::Id::main>, public Instance {
 
   static thread_local ThreadLocalData thread_local_data_;
   std::vector<SlotImpl*> slots_;
+  // A set of index of freed slots.
+  absl::flat_hash_set<uint32_t> free_slot_indexes_;
+
   std::list<std::reference_wrapper<Event::Dispatcher>> registered_threads_;
   std::thread::id main_thread_id_;
   Event::Dispatcher* main_thread_dispatcher_{};