#include "source/server/overload_manager_impl.h"

#include <chrono>

#include "envoy/common/exception.h"

#include "envoy/config/overload/v3/overload.pb.h"

#include "envoy/config/overload/v3/overload.pb.validate.h"

#include "envoy/stats/scope.h"

#include "source/common/common/fmt.h"

#include "source/common/config/utility.h"

#include "source/common/event/scaled_range_timer_manager_impl.h"

#include "source/common/protobuf/utility.h"

#include "source/common/stats/symbol_table.h"

#include "source/server/resource_monitor_config_impl.h"

#include "absl/container/node_hash_map.h"

#include "absl/strings/str_cat.h"

namespace Envoy {

namespace Server {

namespace {

using TriggerPtr = std::unique_ptr<Trigger>;

class ThresholdTriggerImpl final : public Trigger {

public:

  ThresholdTriggerImpl(const envoy::config::overload::v3::ThresholdTrigger& config)

      : threshold_(config.value()), state_(OverloadActionState::inactive()) {}

  bool updateValue(double value) override {

    const OverloadActionState state = actionState();

    state_ =

        value >= threshold_ ? OverloadActionState::saturated() : OverloadActionState::inactive();

    // This is a floating point comparison, though state_ is always either

    // saturated or inactive so there's no risk due to floating point precision.

    return state.value() != actionState().value();

  }

  OverloadActionState actionState() const override { return state_; }

private:

  const double threshold_;

  OverloadActionState state_;

};

class ScaledTriggerImpl final : public Trigger {

public:

  static absl::StatusOr<std::unique_ptr<ScaledTriggerImpl>>

  create(const envoy::config::overload::v3::ScaledTrigger& config) {

    if (config.scaling_threshold() >= config.saturation_threshold()) {

      return absl::InvalidArgumentError("scaling_threshold must be less than saturation_threshold");

    }

    return std::unique_ptr<ScaledTriggerImpl>(new ScaledTriggerImpl(config));

  }

  bool updateValue(double value) override {

    const OverloadActionState old_state = actionState();

    if (value <= scaling_threshold_) {

      state_ = OverloadActionState::inactive();

    } else if (value >= saturated_threshold_) {

      state_ = OverloadActionState::saturated();

    } else {

      state_ = OverloadActionState(

          UnitFloat((value - scaling_threshold_) / (saturated_threshold_ - scaling_threshold_)));

    }

    // All values of state_ are produced via this same code path. Even if

    // old_state and state_ should be approximately equal, there's no harm in

    // signaling for a small change if they're not float::operator== equal.

    return state_.value() != old_state.value();

  }

  OverloadActionState actionState() const override { return state_; }

private:

  ScaledTriggerImpl(const envoy::config::overload::v3::ScaledTrigger& config)

      : scaling_threshold_(config.scaling_threshold()),

        saturated_threshold_(config.saturation_threshold()),

        state_(OverloadActionState::inactive()) {}

  const double scaling_threshold_;

  const double saturated_threshold_;

  OverloadActionState state_;

};

absl::StatusOr<TriggerPtr>

createTriggerFromConfig(const envoy::config::overload::v3::Trigger& trigger_config) {

  TriggerPtr trigger;

  switch (trigger_config.trigger_oneof_case()) {

  case envoy::config::overload::v3::Trigger::TriggerOneofCase::kThreshold:

    trigger = std::make_unique<ThresholdTriggerImpl>(trigger_config.threshold());

    break;

  case envoy::config::overload::v3::Trigger::TriggerOneofCase::kScaled: {

    auto trigger_or_error = ScaledTriggerImpl::create(trigger_config.scaled());

    RETURN_IF_NOT_OK(trigger_or_error.status());

    trigger = std::move(trigger_or_error.value());

    break;

  }

  case envoy::config::overload::v3::Trigger::TriggerOneofCase::TRIGGER_ONEOF_NOT_SET:

    return absl::InvalidArgumentError(

        absl::StrCat("action not set for trigger ", trigger_config.name()));

  }

  return trigger;

}

Stats::Counter& makeCounter(Stats::Scope& scope, absl::string_view name_of_stat) {

  Stats::StatNameManagedStorage stat_name(name_of_stat, scope.symbolTable());

  return scope.counterFromStatName(stat_name.statName());

}

Stats::Counter& makeCounter(Stats::Scope& scope, absl::string_view a, absl::string_view b) {

  return makeCounter(scope, absl::StrCat("overload.", a, ".", b));

}

Stats::Gauge& makeGauge(Stats::Scope& scope, absl::string_view a, absl::string_view b,

                        Stats::Gauge::ImportMode import_mode) {

  Stats::StatNameManagedStorage stat_name(absl::StrCat("overload.", a, ".", b),

                                          scope.symbolTable());

  return scope.gaugeFromStatName(stat_name.statName(), import_mode);

}

Stats::Histogram& makeHistogram(Stats::Scope& scope, absl::string_view name,

                                Stats::Histogram::Unit unit) {

  Stats::StatNameManagedStorage stat_name(absl::StrCat("overload.", name), scope.symbolTable());

  return scope.histogramFromStatName(stat_name.statName(), unit);

}

absl::StatusOr<Event::ScaledTimerType> parseTimerType(

    envoy::config::overload::v3::ScaleTimersOverloadActionConfig::TimerType config_timer_type) {

  using Config = envoy::config::overload::v3::ScaleTimersOverloadActionConfig;

  switch (config_timer_type) {

  case Config::HTTP_DOWNSTREAM_CONNECTION_IDLE:

    return Event::ScaledTimerType::HttpDownstreamIdleConnectionTimeout;

  case Config::HTTP_DOWNSTREAM_STREAM_IDLE:

    return Event::ScaledTimerType::HttpDownstreamIdleStreamTimeout;

  case Config::TRANSPORT_SOCKET_CONNECT:

    return Event::ScaledTimerType::TransportSocketConnectTimeout;

  default:

    return absl::InvalidArgumentError(

        fmt::format("Unknown timer type {}", static_cast<int>(config_timer_type)));

  }

}

absl::StatusOr<Event::ScaledTimerTypeMap>

parseTimerMinimums(const ProtobufWkt::Any& typed_config,

                   ProtobufMessage::ValidationVisitor& validation_visitor) {

  using Config = envoy::config::overload::v3::ScaleTimersOverloadActionConfig;

  const Config action_config =

      MessageUtil::anyConvertAndValidate<Config>(typed_config, validation_visitor);

  Event::ScaledTimerTypeMap timer_map;

  for (const auto& scale_timer : action_config.timer_scale_factors()) {

    auto timer_or_error = parseTimerType(scale_timer.timer());

    RETURN_IF_NOT_OK(timer_or_error.status());

    const Event::ScaledTimerType& timer_type = *timer_or_error;

    const Event::ScaledTimerMinimum minimum =

        scale_timer.has_min_timeout()

            ? Event::ScaledTimerMinimum(Event::AbsoluteMinimum(std::chrono::milliseconds(

                  DurationUtil::durationToMilliseconds(scale_timer.min_timeout()))))

            : Event::ScaledTimerMinimum(

                  Event::ScaledMinimum(UnitFloat(scale_timer.min_scale().value() / 100.0)));

    auto [_, inserted] = timer_map.insert(std::make_pair(timer_type, minimum));

    UNREFERENCED_PARAMETER(_);

    if (!inserted) {

      return absl::InvalidArgumentError(fmt::format("Found duplicate entry for timer type {}",

                                                    Config::TimerType_Name(scale_timer.timer())));

    }

  }

  return timer_map;

}

} // namespace

/**

 * Thread-local copy of the state of each configured overload action.

 */

class ThreadLocalOverloadStateImpl : public ThreadLocalOverloadState {

public:

  explicit ThreadLocalOverloadStateImpl(

      const NamedOverloadActionSymbolTable& action_symbol_table,

      std::shared_ptr<absl::node_hash_map<OverloadProactiveResourceName, ProactiveResource>>&

          proactive_resources)

      : action_symbol_table_(action_symbol_table),

        actions_(action_symbol_table.size(), OverloadActionState(UnitFloat::min())),

        proactive_resources_(proactive_resources) {}

  const OverloadActionState& getState(const std::string& action) override {

    if (const auto symbol = action_symbol_table_.lookup(action); symbol != absl::nullopt) {

      return actions_[symbol->index()];

    }

    return always_inactive_;

  }

  void setState(NamedOverloadActionSymbolTable::Symbol action, OverloadActionState state) {

    actions_[action.index()] = state;

  }

  bool tryAllocateResource(OverloadProactiveResourceName resource_name,

                           int64_t increment) override {

    const auto proactive_resource = proactive_resources_->find(resource_name);

    if (proactive_resource == proactive_resources_->end()) {

      ENVOY_LOG_MISC(warn, "Failed to allocate resource usage, resource monitor is not configured");

      return false;

    }

    return proactive_resource->second.tryAllocateResource(increment);

  }

  bool tryDeallocateResource(OverloadProactiveResourceName resource_name,

                             int64_t decrement) override {

    const auto proactive_resource = proactive_resources_->find(resource_name);

    if (proactive_resource == proactive_resources_->end()) {

      ENVOY_LOG_MISC(warn,

                     "Failed to deallocate resource usage, resource monitor is not configured");

      return false;

    }

    return proactive_resource->second.tryDeallocateResource(decrement);

  }

  bool isResourceMonitorEnabled(OverloadProactiveResourceName resource_name) override {

    const auto proactive_resource = proactive_resources_->find(resource_name);

    return proactive_resource != proactive_resources_->end();

  }

  ProactiveResourceMonitorOptRef

  getProactiveResourceMonitorForTest(OverloadProactiveResourceName resource_name) override {

    const auto proactive_resource = proactive_resources_->find(resource_name);

    if (proactive_resource == proactive_resources_->end()) {

      ENVOY_LOG_MISC(warn, "Failed to get resource usage, resource monitor is not configured");

      return makeOptRefFromPtr<ProactiveResourceMonitor>(nullptr);

    }

    return proactive_resource->second.getProactiveResourceMonitorForTest();

  }

private:

  static const OverloadActionState always_inactive_;

  const NamedOverloadActionSymbolTable& action_symbol_table_;

  std::vector<OverloadActionState> actions_;

  std::shared_ptr<absl::node_hash_map<OverloadProactiveResourceName, ProactiveResource>>

      proactive_resources_;

};

const OverloadActionState ThreadLocalOverloadStateImpl::always_inactive_{UnitFloat::min()};

NamedOverloadActionSymbolTable::Symbol

NamedOverloadActionSymbolTable::get(absl::string_view string) {

  if (auto it = table_.find(string); it != table_.end()) {

    return Symbol(it->second);

  }

  size_t index = table_.size();

  names_.emplace_back(string);

  table_.emplace(std::make_pair(string, index));

  return Symbol(index);

}

absl::optional<NamedOverloadActionSymbolTable::Symbol>

NamedOverloadActionSymbolTable::lookup(absl::string_view string) const {

  if (auto it = table_.find(string); it != table_.end()) {

    return Symbol(it->second);

  }

  return absl::nullopt;

}

const absl::string_view NamedOverloadActionSymbolTable::name(Symbol symbol) const {

  return names_.at(symbol.index());

}

absl::StatusOr<std::unique_ptr<OverloadAction>>

OverloadAction::create(const envoy::config::overload::v3::OverloadAction& config,

                       Stats::Scope& stats_scope) {

  absl::Status creation_status = absl::OkStatus();

  auto ret =

      std::unique_ptr<OverloadAction>(new OverloadAction(config, stats_scope, creation_status));

  RETURN_IF_NOT_OK(creation_status);

  return ret;

}

OverloadAction::OverloadAction(const envoy::config::overload::v3::OverloadAction& config,

                               Stats::Scope& stats_scope, absl::Status& creation_status)

    : state_(OverloadActionState::inactive()),

      active_gauge_(

          makeGauge(stats_scope, config.name(), "active", Stats::Gauge::ImportMode::NeverImport)),

      scale_percent_gauge_(makeGauge(stats_scope, config.name(), "scale_percent",

                                     Stats::Gauge::ImportMode::NeverImport)) {

  for (const auto& trigger_config : config.triggers()) {

    absl::StatusOr<TriggerPtr> trigger_or_error = createTriggerFromConfig(trigger_config);

    SET_AND_RETURN_IF_NOT_OK(trigger_or_error.status(), creation_status);

    if (!triggers_.try_emplace(trigger_config.name(), std::move(*trigger_or_error)).second) {

      creation_status = absl::InvalidArgumentError(

          absl::StrCat("Duplicate trigger resource for overload action ", config.name()));

      return;

    }

  }

  active_gauge_.set(0);

  scale_percent_gauge_.set(0);

}

bool OverloadAction::updateResourcePressure(const std::string& name, double pressure) {

  const OverloadActionState old_state = getState();

  auto it = triggers_.find(name);

  ASSERT(it != triggers_.end());

  if (!it->second->updateValue(pressure)) {

    return false;

  }

  const auto trigger_new_state = it->second->actionState();

  active_gauge_.set(trigger_new_state.isSaturated() ? 1 : 0);

  scale_percent_gauge_.set(trigger_new_state.value().value() * 100);

  {

    // Compute the new state as the maximum over all trigger states.

    OverloadActionState new_state = OverloadActionState::inactive();

    for (auto& trigger : triggers_) {

      const auto trigger_state = trigger.second->actionState();

      if (trigger_state.value() > new_state.value()) {

        new_state = trigger_state;

      }

    }

    state_ = new_state;

  }

  return state_.value() != old_state.value();

}

OverloadActionState OverloadAction::getState() const { return state_; }

absl::StatusOr<std::unique_ptr<LoadShedPointImpl>>

LoadShedPointImpl::create(const envoy::config::overload::v3::LoadShedPoint& config,

                          Stats::Scope& stats_scope, Random::RandomGenerator& random_generator) {

  absl::Status creation_status = absl::OkStatus();

  auto ret = std::unique_ptr<LoadShedPointImpl>(

      new LoadShedPointImpl(config, stats_scope, random_generator, creation_status));

  RETURN_IF_NOT_OK(creation_status);

  return ret;

}

LoadShedPointImpl::LoadShedPointImpl(const envoy::config::overload::v3::LoadShedPoint& config,

                                     Stats::Scope& stats_scope,

                                     Random::RandomGenerator& random_generator,

                                     absl::Status& creation_status)

    : scale_percent_(makeGauge(stats_scope, config.name(), "scale_percent",

                               Stats::Gauge::ImportMode::NeverImport)),

      shed_load_counter_(makeCounter(stats_scope, config.name(), "shed_load_count")),

      random_generator_(random_generator) {

  for (const auto& trigger_config : config.triggers()) {

    auto trigger_or_error = createTriggerFromConfig(trigger_config);

    SET_AND_RETURN_IF_NOT_OK(trigger_or_error.status(), creation_status);

    if (!triggers_.try_emplace(trigger_config.name(), std::move(*trigger_or_error)).second) {

      creation_status = absl::InvalidArgumentError(

          absl::StrCat("Duplicate trigger resource for LoadShedPoint ", config.name()));

      return;

    }

  }

};

void LoadShedPointImpl::updateResource(absl::string_view resource_name,

                                       double resource_utilization) {

  auto it = triggers_.find(resource_name);

  if (it == triggers_.end()) {

    return;

  }

  it->second->updateValue(resource_utilization);

  updateProbabilityShedLoad();

}

void LoadShedPointImpl::updateProbabilityShedLoad() {

  float max_unit_float = 0.0f;

  for (const auto& trigger : triggers_) {

    max_unit_float = std::max(trigger.second->actionState().value().value(), max_unit_float);

  }

  probability_shed_load_.store(max_unit_float);

  // Update stats.

  scale_percent_.set(100 * max_unit_float);

}

bool LoadShedPointImpl::shouldShedLoad() {

  float unit_float_probability_shed_load = probability_shed_load_.load();

  // This should be ok as we're using unit float which saturates at 1.0f.

  if (unit_float_probability_shed_load == 1.0f) {

    shed_load_counter_.inc();

    return true;

  }

  if (random_generator_.bernoulli(UnitFloat(unit_float_probability_shed_load))) {

    shed_load_counter_.inc();

    return true;

  }

  return false;

}

absl::StatusOr<std::unique_ptr<OverloadManagerImpl>>

OverloadManagerImpl::create(Event::Dispatcher& dispatcher, Stats::Scope& stats_scope,

                            ThreadLocal::SlotAllocator& slot_allocator,

                            const envoy::config::overload::v3::OverloadManager& config,

                            ProtobufMessage::ValidationVisitor& validation_visitor, Api::Api& api,

                            const Server::Options& options) {

  absl::Status creation_status = absl::OkStatus();

  auto ret = std::unique_ptr<OverloadManagerImpl>(

      new OverloadManagerImpl(dispatcher, stats_scope, slot_allocator, config, validation_visitor,

                              api, options, creation_status));

  RETURN_IF_NOT_OK(creation_status);

  return ret;

}

OverloadManagerImpl::OverloadManagerImpl(Event::Dispatcher& dispatcher, Stats::Scope& stats_scope,

                                         ThreadLocal::SlotAllocator& slot_allocator,

                                         const envoy::config::overload::v3::OverloadManager& config,

                                         ProtobufMessage::ValidationVisitor& validation_visitor,

                                         Api::Api& api, const Server::Options& options,

                                         absl::Status& creation_status)

    : dispatcher_(dispatcher), time_source_(api.timeSource()), tls_(slot_allocator),

      refresh_interval_(

          std::chrono::milliseconds(PROTOBUF_GET_MS_OR_DEFAULT(config, refresh_interval, 1000))),

      refresh_interval_delays_(makeHistogram(stats_scope, "refresh_interval_delay",

                                             Stats::Histogram::Unit::Milliseconds)),

      proactive_resources_(

          std::make_unique<

              absl::node_hash_map<OverloadProactiveResourceName, ProactiveResource>>()) {

  Configuration::ResourceMonitorFactoryContextImpl context(dispatcher, options, api,

                                                           validation_visitor);

  // We should hide impl details from users, for them there should be no distinction between

  // proactive and regular resource monitors in configuration API. But internally we will maintain

  // two distinct collections of proactive and regular resources. Proactive resources are not

  // subject to periodic flushes and can be recalculated/updated on demand by invoking

  // `tryAllocateResource/tryDeallocateResource` via thread local overload state.

  for (const auto& resource : config.resource_monitors()) {

    const auto& name = resource.name();

    // Check if it is a proactive resource.

    auto proactive_resource_it =

        OverloadProactiveResources::get().proactive_action_name_to_resource_.find(name);

    ENVOY_LOG(debug, "Evaluating resource {}", name);

    bool result = false;

    if (proactive_resource_it !=

        OverloadProactiveResources::get().proactive_action_name_to_resource_.end()) {

      ENVOY_LOG(debug, "Adding proactive resource monitor for {}", name);

      auto& factory =

          Config::Utility::getAndCheckFactory<Configuration::ProactiveResourceMonitorFactory>(

              resource);

      auto config =

          Config::Utility::translateToFactoryConfig(resource, validation_visitor, factory);

      auto monitor = factory.createProactiveResourceMonitor(*config, context);

      result =

          proactive_resources_

              ->try_emplace(proactive_resource_it->second, name, std::move(monitor), stats_scope)

              .second;

    } else {

      ENVOY_LOG(debug, "Adding resource monitor for {}", name);

      auto& factory =

          Config::Utility::getAndCheckFactory<Configuration::ResourceMonitorFactory>(resource);

      auto config =

          Config::Utility::translateToFactoryConfig(resource, validation_visitor, factory);

      auto monitor = factory.createResourceMonitor(*config, context);

      result = resources_.try_emplace(name, name, std::move(monitor), *this, stats_scope).second;

    }

    if (!result) {

      creation_status =

          absl::InvalidArgumentError(absl::StrCat("Duplicate resource monitor ", name));

      return;

    }

  }

  for (const auto& action : config.actions()) {

    const auto& name = action.name();

    const auto symbol = action_symbol_table_.get(name);

    ENVOY_LOG(debug, "Adding overload action {}", name);

    // Validate that this is a well known overload action.

    const auto& well_known_actions = OverloadActionNames::get().WellKnownActions;

    if (std::find(well_known_actions.begin(), well_known_actions.end(), name) ==

        well_known_actions.end()) {

      creation_status =

          absl::InvalidArgumentError(absl::StrCat("Unknown Overload Manager Action ", name));

      return;

    }

    // TODO: use in place construction once https://github.com/abseil/abseil-cpp/issues/388 is

    // addressed

    // We cannot currently use in place construction as the OverloadAction constructor may fail,

    // causing an inconsistent internal state of the actions_ map, which on destruction results in

    // an invalid free.

    auto action_or_error = OverloadAction::create(action, stats_scope);

    SET_AND_RETURN_IF_NOT_OK(action_or_error.status(), creation_status);

    auto result = actions_.try_emplace(symbol, std::move(*action_or_error));

    if (!result.second) {

      creation_status =

          absl::InvalidArgumentError(absl::StrCat("Duplicate overload action ", name));

      return;

    }

    if (name == OverloadActionNames::get().ReduceTimeouts) {

      auto timer_or_error = parseTimerMinimums(action.typed_config(), validation_visitor);

      SET_AND_RETURN_IF_NOT_OK(timer_or_error.status(), creation_status);

      timer_minimums_ =

          std::make_shared<const Event::ScaledTimerTypeMap>(std::move(*timer_or_error));

    } else if (name == OverloadActionNames::get().ResetStreams) {

      if (!config.has_buffer_factory_config()) {

        creation_status = absl::InvalidArgumentError(

            fmt::format("Overload action \"{}\" requires buffer_factory_config.", name));

        return;

      }

      makeCounter(api.rootScope(), OverloadActionStatsNames::get().ResetStreamsCount);

    } else if (action.has_typed_config()) {

      creation_status = absl::InvalidArgumentError(fmt::format(

          "Overload action \"{}\" has an unexpected value for the typed_config field", name));

      return;

    }

    for (const auto& trigger : action.triggers()) {

      const std::string& resource = trigger.name();

      auto proactive_resource_it =

          OverloadProactiveResources::get().proactive_action_name_to_resource_.find(resource);

      if (resources_.find(resource) == resources_.end() &&

          proactive_resource_it ==

              OverloadProactiveResources::get().proactive_action_name_to_resource_.end()) {

        creation_status = absl::InvalidArgumentError(

            fmt::format("Unknown trigger resource {} for overload action {}", resource, name));

        return;

      }

      resource_to_actions_.insert(std::make_pair(resource, symbol));

    }

  }

  // Validate the trigger resources for Load shedPoints.

  for (const auto& point : config.loadshed_points()) {

    for (const auto& trigger : point.triggers()) {

      if (!resources_.contains(trigger.name())) {

        creation_status = absl::InvalidArgumentError(fmt::format(

            "Unknown trigger resource {} for loadshed point {}", trigger.name(), point.name()));

        return;

      }

    }

    auto load_shed_or_error =

        LoadShedPointImpl::create(point, api.rootScope(), api.randomGenerator());

    SET_AND_RETURN_IF_NOT_OK(load_shed_or_error.status(), creation_status);

    const auto result = loadshed_points_.try_emplace(point.name(), *std::move(load_shed_or_error));

    if (!result.second) {

      creation_status =

          absl::InvalidArgumentError(absl::StrCat("Duplicate loadshed point ", point.name()));

      return;

    }

  }

}

void OverloadManagerImpl::start() {

  ASSERT(!started_);

  started_ = true;

  tls_.set([this](Event::Dispatcher&) {

    return std::make_shared<ThreadLocalOverloadStateImpl>(action_symbol_table_,

                                                          proactive_resources_);

  });

  if (resources_.empty()) {

    return;

  }

  timer_ = dispatcher_.createTimer([this]() -> void {

    // Guarantee that all resource updates get flushed after no more than one refresh_interval_.

    flushResourceUpdates();

    // Start a new flush epoch. If all resource updates complete before this callback runs, the last

    // resource update will call flushResourceUpdates to flush the whole batch early.

    ++flush_epoch_;

    flush_awaiting_updates_ = resources_.size();

    for (auto& resource : resources_) {

      resource.second.update(flush_epoch_);

    }

    // Record delay.

    auto now = time_source_.monotonicTime();

    std::chrono::milliseconds delay =

        std::chrono::duration_cast<std::chrono::milliseconds>(now - time_resources_last_measured_);

    refresh_interval_delays_.recordValue(delay.count());

    time_resources_last_measured_ = now;

    timer_->enableTimer(refresh_interval_);

  });

  time_resources_last_measured_ = time_source_.monotonicTime();

  timer_->enableTimer(refresh_interval_);

}

void OverloadManagerImpl::stop() {

  // Disable any pending timeouts.

  if (timer_) {

    timer_->disableTimer();

  }

  // Clear the resource map to block on any pending updates.

  resources_.clear();

  // TODO(nezdolik): wrap proactive monitors into atomic? and clear it here

}

bool OverloadManagerImpl::registerForAction(const std::string& action,

                                            Event::Dispatcher& dispatcher,

                                            OverloadActionCb callback) {

  ASSERT(!started_);

  const auto symbol = action_symbol_table_.get(action);

  if (actions_.find(symbol) == actions_.end()) {

    ENVOY_LOG(debug, "No overload action is configured for {}.", action);

    return false;

  }

  action_to_callbacks_.emplace(std::piecewise_construct, std::forward_as_tuple(symbol),

                               std::forward_as_tuple(dispatcher, callback));

  return true;

}

ThreadLocalOverloadState& OverloadManagerImpl::getThreadLocalOverloadState() { return *tls_; }

Event::ScaledRangeTimerManagerFactory OverloadManagerImpl::scaledTimerFactory() {

  return [this](Event::Dispatcher& dispatcher) {

    auto manager = createScaledRangeTimerManager(dispatcher, timer_minimums_);

    registerForAction(OverloadActionNames::get().ReduceTimeouts, dispatcher,

                      [manager = manager.get()](OverloadActionState scale_state) {

                        manager->setScaleFactor(

                            // The action state is 0 for no overload up to 1 for maximal overload,

                            // but the scale factor for timers is 1 for no scaling and 0 for maximal

                            // scaling, so invert the value to pass in (1-value).

                            scale_state.value().invert());

                      });

    return manager;

  };

}

LoadShedPoint* OverloadManagerImpl::getLoadShedPoint(absl::string_view point_name) {

  if (auto it = loadshed_points_.find(point_name); it != loadshed_points_.end()) {

    return it->second.get();

  }

  return nullptr;

}

Event::ScaledRangeTimerManagerPtr OverloadManagerImpl::createScaledRangeTimerManager(

    Event::Dispatcher& dispatcher,

    const Event::ScaledTimerTypeMapConstSharedPtr& timer_minimums) const {

  return std::make_unique<Event::ScaledRangeTimerManagerImpl>(dispatcher, timer_minimums);

}

void OverloadManagerImpl::updateResourcePressure(const std::string& resource, double pressure,

                                                 FlushEpochId flush_epoch) {

  auto [start, end] = resource_to_actions_.equal_range(resource);

  std::for_each(start, end, [&](ResourceToActionMap::value_type& entry) {

    const NamedOverloadActionSymbolTable::Symbol action = entry.second;

    auto action_it = actions_.find(action);

    ASSERT(action_it != actions_.end());

    const OverloadActionState old_state = action_it->second->getState();

    if (action_it->second->updateResourcePressure(resource, pressure)) {

      const auto state = action_it->second->getState();

      if (old_state.isSaturated() != state.isSaturated()) {

        ENVOY_LOG(debug, "Overload action {} became {}", action_symbol_table_.name(action),

                  (state.isSaturated() ? "saturated" : "scaling"));

      }

      // Record the updated value to be sent to workers on the next thread-local-state flush, along

      // with any update callbacks. This might overwrite a previous action state change caused by a

      // pressure update for a different resource that hasn't been flushed yet. That's okay because

      // the state recorded here includes the information from all previous resource updates. So

      // even if resource 1 causes an action to have value A, and a later update to resource 2

      // causes the action to have value B, B would have been the result for whichever order the

      // updates to resources 1 and 2 came in.

      state_updates_to_flush_.insert_or_assign(action, state);

      auto [callbacks_start, callbacks_end] = action_to_callbacks_.equal_range(action);

      std::for_each(callbacks_start, callbacks_end, [&](ActionToCallbackMap::value_type& cb_entry) {

        callbacks_to_flush_.insert_or_assign(&cb_entry.second, state);

      });

    }

  });

  for (auto& loadshed_point : loadshed_points_) {

    loadshed_point.second->updateResource(resource, pressure);

  }

  // Eagerly flush updates if this is the last call to updateResourcePressure expected for the

  // current epoch. This assert is always valid because flush_awaiting_updates_ is initialized

  // before each batch of updates, and even if a resource monitor performs a double update, or a

  // previous update callback is late, the logic in OverloadManager::Resource::update() will prevent

  // unexpected calls to this function.

  ASSERT(flush_awaiting_updates_ > 0);

  --flush_awaiting_updates_;

  if (flush_epoch == flush_epoch_ && flush_awaiting_updates_ == 0) {

    flushResourceUpdates();

  }

}

void OverloadManagerImpl::flushResourceUpdates() {

  if (!state_updates_to_flush_.empty()) {

    auto shared_updates = std::make_shared<

        absl::flat_hash_map<NamedOverloadActionSymbolTable::Symbol, OverloadActionState>>();

    std::swap(*shared_updates, state_updates_to_flush_);

    tls_.runOnAllThreads(

        [updates = std::move(shared_updates)](OptRef<ThreadLocalOverloadStateImpl> overload_state) {

          for (const auto& [action, state] : *updates) {

            overload_state->setState(action, state);

          }

        });

  }

  for (const auto& [cb, state] : callbacks_to_flush_) {

    cb->dispatcher_.post([cb = cb, state = state]() { cb->callback_(state); });

  }

  callbacks_to_flush_.clear();

}

OverloadManagerImpl::Resource::Resource(const std::string& name, ResourceMonitorPtr monitor,

                                        OverloadManagerImpl& manager, Stats::Scope& stats_scope)

    : name_(name), monitor_(std::move(monitor)), manager_(manager),

      pressure_gauge_(

          makeGauge(stats_scope, name, "pressure", Stats::Gauge::ImportMode::NeverImport)),

      failed_updates_counter_(makeCounter(stats_scope, name, "failed_updates")),

      skipped_updates_counter_(makeCounter(stats_scope, name, "skipped_updates")) {}

void OverloadManagerImpl::Resource::update(FlushEpochId flush_epoch) {

  if (!pending_update_) {

    pending_update_ = true;

    flush_epoch_ = flush_epoch;

    monitor_->updateResourceUsage(*this);

    return;

  }

  ENVOY_LOG(debug, "Skipping update for resource {} which has pending update", name_);

  skipped_updates_counter_.inc();

}

void OverloadManagerImpl::Resource::onSuccess(const ResourceUsage& usage) {

  pending_update_ = false;

  manager_.updateResourcePressure(name_, usage.resource_pressure_, flush_epoch_);

  pressure_gauge_.set(usage.resource_pressure_ * 100); // convert to percent

}

void OverloadManagerImpl::Resource::onFailure(const EnvoyException& error) {

  pending_update_ = false;

  ENVOY_LOG(info, "Failed to update resource {}: {}", name_, error.what());

  failed_updates_counter_.inc();

}

} // namespace Server

} // namespace Envoy