Make ThreadLocalOverloadState virtual

Make ThreadLocalOverloadState an interface that the Overload Manager can
return an implementation of. This decouples the behavior specified by
the header from the details of the implementation.

Signed-off-by: Alex Konradi <akonradi@google.com>

include/envoy/server/overload_manager.h

@@ -31,27 +31,12 @@ using OverloadActionCb = std::function<void(OverloadActionState)>;
 /**
  * Thread-local copy of the state of each configured overload action.
  */
-class ThreadLocalOverloadState : public ThreadLocal::ThreadLocalObject {
+class ThreadLocalOverloadState {
 public:
-  const OverloadActionState& getState(const std::string& action) {
-    auto it = actions_.find(action);
-    if (it == actions_.end()) {
-      it = actions_.insert(std::make_pair(action, OverloadActionState::Inactive)).first;
-    }
-    return it->second;
-  }
-
-  void setState(const std::string& action, OverloadActionState state) {
-    auto it = actions_.find(action);
-    if (it == actions_.end()) {
-      actions_[action] = state;
-    } else {
-      it->second = state;
-    }
-  }
-
-private:
-  std::unordered_map<std::string, OverloadActionState> actions_;
+  virtual ~ThreadLocalOverloadState() = default;
+
+  // Get a thread-local reference to the value for the given action key.
+  virtual const OverloadActionState& getState(const std::string& action) PURE;
 };
 
 /**

source/server/admin/admin.cc

@@ -650,12 +650,6 @@ ConfigTracker& AdminImpl::getConfigTracker() { return config_tracker_; }
 AdminImpl::NullRouteConfigProvider::NullRouteConfigProvider(TimeSource& time_source)
     : config_(new Router::NullConfigImpl()), time_source_(time_source) {}
 
-OverloadTimerFactory AdminImpl::NullOverloadManager::NullThreadOverloadState::getTimerFactory() {
-  return [this](absl::string_view, Event::TimerCb callback) {
-    return dispatcher_.createTimer(callback);
-  };
-}
-
 void AdminImpl::startHttpListener(const std::string& access_log_path,
                                   const std::string& address_out_path,
                                   Network::Address::InstanceConstSharedPtr address,

source/server/admin/admin.h

@@ -254,8 +254,6 @@ class AdminImpl : public Admin,
     struct NullThreadOverloadState : public ThreadLocalOverloadState {
       NullThreadOverloadState(Event::Dispatcher& dispatcher) : dispatcher_(dispatcher) {}
       const OverloadActionState& getState(const std::string&) override { return inactive_; }
-      void setState(const std::string&, OverloadActionState) override {}
-      OverloadTimerFactory getTimerFactory() override;
 
       const OverloadActionState inactive_ = OverloadActionState::Inactive;
       Event::Dispatcher& dispatcher_;

source/server/overload_manager_impl.cc

@@ -35,6 +35,33 @@ class ThresholdTriggerImpl : public OverloadAction::Trigger {
   absl::optional<double> value_;
 };
 
+/**
+ * Thread-local copy of the state of each configured overload action.
+ */
+class ThreadLocalOverloadStateImpl : public ThreadLocalOverloadState,
+                                     public ThreadLocal::ThreadLocalObject {
+public:
+  const OverloadActionState& getState(const std::string& action) override {
+    auto it = actions_.find(action);
+    if (it == actions_.end()) {
+      it = actions_.insert(std::make_pair(action, OverloadActionState::Inactive)).first;
+    }
+    return it->second;
+  }
+
+  void setState(const std::string& action, OverloadActionState state) {
+    auto it = actions_.find(action);
+    if (it == actions_.end()) {
+      actions_[action] = state;
+    } else {
+      it->second = state;
+    }
+  }
+
+private:
+  std::unordered_map<std::string, OverloadActionState> actions_;
+};
+
 Stats::Counter& makeCounter(Stats::Scope& scope, absl::string_view a, absl::string_view b) {
   Stats::StatNameManagedStorage stat_name(absl::StrCat("overload.", a, ".", b),
                                           scope.symbolTable());
@@ -148,7 +175,7 @@ void OverloadManagerImpl::start() {
   started_ = true;
 
   tls_->set([](Event::Dispatcher&) -> ThreadLocal::ThreadLocalObjectSharedPtr {
-    return std::make_shared<ThreadLocalOverloadState>();
+    return std::make_shared<ThreadLocalOverloadStateImpl>();
   });
 
   if (resources_.empty()) {
@@ -191,7 +218,7 @@ bool OverloadManagerImpl::registerForAction(const std::string& action,
 }
 
 ThreadLocalOverloadState& OverloadManagerImpl::getThreadLocalOverloadState() {
-  return tls_->getTyped<ThreadLocalOverloadState>();
+  return tls_->getTyped<ThreadLocalOverloadStateImpl>();
 }
 
 void OverloadManagerImpl::updateResourcePressure(const std::string& resource, double pressure) {
@@ -208,7 +235,7 @@ void OverloadManagerImpl::updateResourcePressure(const std::string& resource, do
                     ENVOY_LOG(info, "Overload action {} became {}", action,
                               is_active ? "active" : "inactive");
                     tls_->runOnAllThreads([this, action, state] {
-                      tls_->getTyped<ThreadLocalOverloadState>().setState(action, state);
+                      tls_->getTyped<ThreadLocalOverloadStateImpl>().setState(action, state);
                     });
                     auto callback_range = action_to_callbacks_.equal_range(action);
                     std::for_each(callback_range.first, callback_range.second,

test/common/http/conn_manager_impl_test.cc

@@ -5600,11 +5600,12 @@ TEST(HttpConnectionManagerTracingStatsTest, verifyTracingStats) {
 }
 
 TEST_F(HttpConnectionManagerImplTest, NoNewStreamWhenOverloaded) {
-  setup(false, "");
+  Server::OverloadActionState stop_accepting_requests = Server::OverloadActionState::Active;
+  ON_CALL(overload_manager_.overload_state_,
+          getState(Server::OverloadActionNames::get().StopAcceptingRequests))
+      .WillByDefault(ReturnRef(stop_accepting_requests));
 
-  overload_manager_.overload_state_.setState(
-      Server::OverloadActionNames::get().StopAcceptingRequests,
-      Server::OverloadActionState::Active);
+  setup(false, "");
 
   EXPECT_CALL(*codec_, dispatch(_)).WillRepeatedly(Invoke([&](Buffer::Instance&) -> Http::Status {
     RequestDecoder* decoder = &conn_manager_->newStream(response_encoder_);
@@ -5630,10 +5631,12 @@ TEST_F(HttpConnectionManagerImplTest, NoNewStreamWhenOverloaded) {
 }
 
 TEST_F(HttpConnectionManagerImplTest, DisableKeepAliveWhenOverloaded) {
-  setup(false, "");
+  Server::OverloadActionState disable_http_keep_alive = Server::OverloadActionState::Active;
+  ON_CALL(overload_manager_.overload_state_,
+          getState(Server::OverloadActionNames::get().DisableHttpKeepAlive))
+      .WillByDefault(ReturnRef(disable_http_keep_alive));
 
-  overload_manager_.overload_state_.setState(
-      Server::OverloadActionNames::get().DisableHttpKeepAlive, Server::OverloadActionState::Active);
+  setup(false, "");
 
   std::shared_ptr<MockStreamDecoderFilter> filter(new NiceMock<MockStreamDecoderFilter>());
   EXPECT_CALL(filter_factory_, createFilterChain(_))

test/mocks/server/mocks.cc

@@ -76,6 +76,11 @@ MockHotRestart::MockHotRestart() : stats_allocator_(*symbol_table_) {
 }
 MockHotRestart::~MockHotRestart() = default;
 
+MockThreadLocalOverloadState::MockThreadLocalOverloadState()
+    : disabled_state_(OverloadActionState::Inactive) {
+  ON_CALL(*this, getState).WillByDefault(ReturnRef(disabled_state_));
+}
+
 MockOverloadManager::MockOverloadManager() {
   ON_CALL(*this, getThreadLocalOverloadState()).WillByDefault(ReturnRef(overload_state_));
 }

test/mocks/server/mocks.h

@@ -331,6 +331,15 @@ class MockWorker : public Worker {
   std::function<void()> remove_filter_chains_completion_;
 };
 
+class MockThreadLocalOverloadState : public ThreadLocalOverloadState {
+public:
+  MockThreadLocalOverloadState();
+  MOCK_METHOD(const OverloadActionState&, getState, (const std::string&), (override));
+
+private:
+  const OverloadActionState disabled_state_;
+};
+
 class MockOverloadManager : public OverloadManager {
 public:
   MockOverloadManager();
@@ -343,7 +352,7 @@ class MockOverloadManager : public OverloadManager {
                OverloadActionCb callback));
   MOCK_METHOD(ThreadLocalOverloadState&, getThreadLocalOverloadState, ());
 
-  ThreadLocalOverloadState overload_state_;
+  NiceMock<MockThreadLocalOverloadState> overload_state_;
 };
 
 class MockInstance : public Instance {