include/tsutil/UdiTable.h

/** @file

  Fixed-size table for tracking frequently observed keys with bounded memory.

  @section license License

  Licensed to the Apache Software Foundation (ASF) under one or more contributor license
  agreements.  See the NOTICE file distributed with this work for additional information regarding
  copyright ownership.  The ASF licenses this file to you under the Apache License, Version 2.0
  (the "License"); you may not use this file except in compliance with the License.  You may obtain
  a copy of the License at

      http://www.apache.org/licenses/LICENSE-2.0

  Unless required by applicable law or agreed to in writing, software distributed under the License
  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
  or implied. See the License for the specific language governing permissions and limitations under
  the License.
*/

#pragma once

#include <algorithm>
#include <atomic>
#include <chrono>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>

namespace ts
{

/** Default eviction policy for UdiTable. */
template <typename Data> struct UdiTableAlwaysEvictable {
  bool
  operator()(Data const &) const noexcept
  {
    return true;
  }
};

/** A fixed-size hash table using the Udi "King of the Hill" algorithm.
 *
 * Instantiations of this table track the keys/entities (IPs, URLs, etc.) with
 * the highest rates of events (e.g. number of requests, number of errors, etc.).
 *
 * Key properties:
 * - Fixed memory: N slots = bounded memory, no unbounded growth
 * - Self-cleaning: No cleanup thread needed, table manages itself
 * - Hot tracking: High-score entries naturally stay in the table
 * - Simple locking: Single mutex for all operations
 * - Safe references: Returns shared_ptr so data survives eviction
 *
 * @tparam Key Default-constructible entity type (e.g., IP address, URL)
 * @tparam Data Default-constructible user data type stored with each key.
 * @tparam Hash Hash function for keys (defaults to std::hash)
 * @tparam CanEvict Predicate callable with @c Data @c const& before replacing
 *   an occupied slot. The default permits every occupied slot to participate.
 *
 * The table owns the key and score for each entry. Users provide only their custom
 * Data type which is stored in a shared_ptr for safe access.
 *
 * Thread Safety:
 * All operations are protected by a single mutex and are serialized.
 * Returned shared_ptr<Data> remains valid even after the slot is evicted.
 * The eviction predicate is invoked while the mutex is held and must not
 * reenter the same table.
 *
 * Example usage:
 * @code
 * struct MyData {
 *   std::atomic<uint32_t> error_count{0};
 *   std::atomic<uint32_t> success_count{0};
 * };
 *
 * ts::UdiTable<std::string, MyData> table(10000);
 *
 * auto data = table.process_event("some_key", 1);
 * if (data) {
 *   data->error_count.fetch_add(1);
 * }
 * @endcode
 */
template <typename Key, typename Data, typename Hash = std::hash<Key>, typename CanEvict = UdiTableAlwaysEvictable<Data>>
class UdiTable
{
public:
  using key_type  = Key;
  using data_type = Data;
  using data_ptr  = std::shared_ptr<Data>;

  enum class ProcessStatus {
    TRACKED,      ///< Existing entry or successful contest.
    CONTEST_LOST, ///< Ordinary score-based contest loss.
    NO_CANDIDATE, ///< Bounded eviction-predicate scan found no candidate.
  };

  using const_data_ptr = std::shared_ptr<Data const>;
  using data_format_fn = std::function<std::string(Key const &, uint32_t, const_data_ptr const &)>;

  // =========================================================================
  // Public API - Declarations
  // =========================================================================

  /**
   * @param[in] num_slots Total number of slots to allocate.
   * @param[in] can_evict Predicate that decides whether occupied slots may be replaced.
   * Called under the table mutex: any predicate locks must follow that lock.
   * For abuse_shield the order is table mutex, then RuleBuckets mutex.
   */
  explicit UdiTable(size_t num_slots, CanEvict can_evict = {});

  // No copying or moving
  UdiTable(UdiTable const &)            = delete;
  UdiTable &operator=(UdiTable const &) = delete;
  UdiTable(UdiTable &&)                 = delete;
  UdiTable &operator=(UdiTable &&)      = delete;

  /** Retrieve the configured data for a given key.
   *
   * @param[in] key The key to look up.
   * @return The data associated with @a key if found, nullptr otherwise.
   *
   * Thread-safe: Uses mutex lock.
   * The returned shared_ptr remains valid even if the slot is later evicted.
   */
  data_ptr       find(Key const &key);
  const_data_ptr find(Key const &key) const;

  /** Process an event for a key, creating a slot if needed via @a contest.
   *
   * If the key is already tracked, increments its score and returns the Data.
   * If not, a call to @a contest is made to see whether the @a key should evict
   * an entry and take its place.
   *
   * @param[in] key The key for which an event is being processed.
   * @param[in] score_delta Score to add (typically 1 for events).
   * @param[out] status Optional result that distinguishes an ordinary contest
   * loss from exhausting the bounded protected-slot scan.
   * @return The data for @a key if tracked or contest won, nullptr otherwise.
   *
   * Thread-safe: Uses mutex lock.
   * The returned shared_ptr remains valid even if the slot is later evicted.
   */
  data_ptr process_event(Key const &key, uint32_t score_delta = 1, ProcessStatus *status = nullptr);

  /** Remove a key from the table.
   *
   * @param[in] key The key to remove.
   * @return Whether @a key was found and removed.
   *
   * Note: Existing shared_ptr references to the removed Data remain valid.
   */
  bool remove(Key const &key);

  /// Snapshot occupied data for maintenance without holding the table lock.
  std::vector<data_ptr>
  data_snapshot() const
  {
    std::lock_guard       lock(mutex_);
    std::vector<data_ptr> result;
    result.reserve(lookup_.size());
    for (const auto &[key, index] : lookup_) {
      result.push_back(slots_[index].data);
    }
    return result;
  }

  /// Clear entries; outstanding shared pointers remain valid.
  void
  clear()
  {
    std::lock_guard lock(mutex_);
    lookup_.clear();
    for (auto &slot : slots_) {
      slot.clear();
    }
    contest_ptr_ = 0;
  }

  // Statistics
  size_t   num_slots() const;
  size_t   slots_used() const;
  uint64_t contests() const;
  uint64_t contests_won() const;
  uint64_t evictions() const;

  /** Reset table-level metrics to zero.
   *
   * Does NOT modify any entries in the table.
   */
  void reset_metrics();

  /** Get the timestamp of the last reset (or table creation).
   */
  std::chrono::system_clock::time_point last_reset_time() const;

  /** Get the number of seconds since last reset (or table creation).
   */
  uint64_t seconds_since_reset() const;

  /** Dump all entries to a string.
   *
   * This may be useful to evaluate the behavior of the table for debugging or
   * diagnostic purposes.
   *
   * @param[in] format_data Optional function to format each entry (key, score, data).
   * @return A string representation of all table entries.
   */
  std::string dump(data_format_fn format_data = nullptr) const;

private:
  /** The data used to track the key and score used to determine eviction.
   */
  struct Slot {
    Key      key{};
    uint32_t score{0};
    data_ptr data;

    /** Check whether there is an entity associated with this slot.
     *
     * Note that a score of zero does not necessarily mean that the slot is
     * empty. Upon table initialization, all slots are empty.
     *
     * @return Whether the slot has no entity assigned.
     */
    bool
    is_empty() const
    {
      return !data;
    }

    /** Remove the entity associated with this slot.
     */
    void
    clear()
    {
      key   = Key{};
      score = 0;
      data.reset();
    }
  };

  /** Determine whether the given key should evict a slot.
   *
   * Assumes @a mutex_ already held exclusively.
   *
   * @param[in] key The key trying to enter.
   * @param[in] incoming_score The score of the incoming key.
   * @return The data for @a key if contest won, nullptr if contest lost.
   */
  data_ptr contest(Key const &key, uint32_t incoming_score, ProcessStatus *status);

  static constexpr size_t MAX_CONTEST_PROBES = 1024;

  /// Single global mutex for all operations
  mutable std::mutex mutex_;

  /// Lookup map: key -> slot index
  std::unordered_map<Key, size_t, Hash> lookup_;

  /// The slots representing the table.
  std::vector<Slot> slots_;

  /// The contest pointer - rotates through all slots as @a contest() is called.
  size_t contest_ptr_{0};

  /// Policy for determining whether an occupied slot may participate in a contest.
  CanEvict can_evict_;

  /// Metrics.
  std::atomic<uint64_t> metric_contests_{0};
  std::atomic<uint64_t> metric_contests_won_{0};
  std::atomic<uint64_t> metric_evictions_{0};

  /// Timestamp of last reset (or construction)
  std::chrono::system_clock::time_point last_reset_time_;
  /// Monotonic timestamp for elapsed-time metrics.
  std::chrono::steady_clock::time_point last_reset_steady_{std::chrono::steady_clock::now()};
};

// ===========================================================================
// Implementation
// ===========================================================================

template <typename Key, typename Data, typename Hash, typename CanEvict>
UdiTable<Key, Data, Hash, CanEvict>::UdiTable(size_t num_slots, CanEvict can_evict)
  : slots_(num_slots), can_evict_(std::move(can_evict)), last_reset_time_(std::chrono::system_clock::now())
{
  lookup_.reserve(num_slots);
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
std::shared_ptr<Data>
UdiTable<Key, Data, Hash, CanEvict>::find(Key const &key)
{
  std::lock_guard<std::mutex> lock(mutex_);

  auto it = lookup_.find(key);
  if (it != lookup_.end()) {
    return slots_[it->second].data;
  }
  return nullptr;
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
std::shared_ptr<Data const>
UdiTable<Key, Data, Hash, CanEvict>::find(Key const &key) const
{
  return const_cast<UdiTable *>(this)->find(key);
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
std::shared_ptr<Data>
UdiTable<Key, Data, Hash, CanEvict>::process_event(Key const &key, uint32_t score_delta, ProcessStatus *status)
{
  std::lock_guard<std::mutex> lock(mutex_);

  if (status) {
    *status = ProcessStatus::TRACKED;
  }

  // Check if already tracked
  auto it = lookup_.find(key);
  if (it != lookup_.end()) {
    Slot &slot = slots_[it->second];
    slot.score = score_delta > UINT32_MAX - slot.score ? UINT32_MAX : slot.score + score_delta;
    return slot.data;
  }

  // Not tracked - contest for a slot
  return contest(key, score_delta, status);
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
bool
UdiTable<Key, Data, Hash, CanEvict>::remove(Key const &key)
{
  std::lock_guard<std::mutex> lock(mutex_);

  auto it = lookup_.find(key);
  if (it == lookup_.end()) {
    return false;
  }

  slots_[it->second].clear();
  lookup_.erase(it);
  return true;
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
size_t
UdiTable<Key, Data, Hash, CanEvict>::num_slots() const
{
  return slots_.size();
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
size_t
UdiTable<Key, Data, Hash, CanEvict>::slots_used() const
{
  std::lock_guard<std::mutex> lock(mutex_);
  return lookup_.size();
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
uint64_t
UdiTable<Key, Data, Hash, CanEvict>::contests() const
{
  return metric_contests_.load(std::memory_order_relaxed);
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
uint64_t
UdiTable<Key, Data, Hash, CanEvict>::contests_won() const
{
  return metric_contests_won_.load(std::memory_order_relaxed);
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
uint64_t
UdiTable<Key, Data, Hash, CanEvict>::evictions() const
{
  return metric_evictions_.load(std::memory_order_relaxed);
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
void
UdiTable<Key, Data, Hash, CanEvict>::reset_metrics()
{
  std::lock_guard<std::mutex> lock(mutex_);
  metric_contests_.store(0, std::memory_order_relaxed);
  metric_contests_won_.store(0, std::memory_order_relaxed);
  metric_evictions_.store(0, std::memory_order_relaxed);
  last_reset_time_   = std::chrono::system_clock::now();
  last_reset_steady_ = std::chrono::steady_clock::now();
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
std::chrono::system_clock::time_point
UdiTable<Key, Data, Hash, CanEvict>::last_reset_time() const
{
  std::lock_guard<std::mutex> lock(mutex_);
  return last_reset_time_;
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
uint64_t
UdiTable<Key, Data, Hash, CanEvict>::seconds_since_reset() const
{
  std::lock_guard<std::mutex> lock(mutex_);
  auto elapsed = std::chrono::duration_cast<std::chrono::seconds>(std::chrono::steady_clock::now() - last_reset_steady_);
  return static_cast<uint64_t>(elapsed.count());
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
std::string
UdiTable<Key, Data, Hash, CanEvict>::dump(data_format_fn format_data) const
{
  struct Snapshot {
    Key            key;
    uint32_t       score;
    const_data_ptr data;
    size_t         slot_idx;
  };

  std::vector<Snapshot> snapshot;
  {
    std::lock_guard<std::mutex> lock(mutex_);
    snapshot.reserve(lookup_.size());
    for (auto const &[key, slot_idx] : lookup_) {
      Slot const &slot = slots_[slot_idx];
      snapshot.push_back({slot.key, slot.score, slot.data, slot_idx});
    }
  }

  std::string result;
  result.reserve(snapshot.size() * 64);

  for (auto const &entry : snapshot) {
    if (format_data) {
      result += format_data(entry.key, entry.score, entry.data);
    } else {
      result += "slot[" + std::to_string(entry.slot_idx) + "] score=" + std::to_string(entry.score) + "\n";
    }
  }
  return result;
}

template <typename Key, typename Data, typename Hash, typename CanEvict>
std::shared_ptr<Data>
UdiTable<Key, Data, Hash, CanEvict>::contest(Key const &key, uint32_t incoming_score, ProcessStatus *status)
{
  // Called with mutex_ already held exclusively

  metric_contests_.fetch_add(1, std::memory_order_relaxed);

  // This should never happen, but just in case
  if (slots_.empty()) {
    if (status) {
      *status = ProcessStatus::NO_CANDIDATE;
    }
    return nullptr;
  }

  // Bound the scan so a predicate that rejects every occupied slot cannot
  // serialize all callers behind a full-table walk.
  size_t slot_idx    = slots_.size();
  size_t probe_limit = std::min(slots_.size(), MAX_CONTEST_PROBES);
  for (size_t count = 0; count < probe_limit; ++count) {
    size_t candidate = contest_ptr_;
    contest_ptr_     = (contest_ptr_ + 1) % slots_.size();
    Slot const &slot = slots_[candidate];
    if (slot.is_empty() || can_evict_(*slot.data)) {
      slot_idx = candidate;
      break;
    }
  }
  if (slot_idx == slots_.size()) {
    if (status) {
      *status = ProcessStatus::NO_CANDIDATE;
    }
    return nullptr;
  }

  Slot    &slot          = slots_[slot_idx];
  uint32_t current_score = slot.score;

  if (slot.is_empty() || incoming_score > current_score) {
    // New key wins - takes the slot
    if (!slot.is_empty()) {
      // Evict the old key from the lookup
      lookup_.erase(slot.key);
      metric_evictions_.fetch_add(1, std::memory_order_relaxed);
    }

    // Initialize the slot with new key and fresh Data
    slot.key     = key;
    slot.score   = incoming_score;
    slot.data    = std::make_shared<Data>();
    lookup_[key] = slot_idx;

    metric_contests_won_.fetch_add(1, std::memory_order_relaxed);
    return slot.data;
  } else {
    // New key loses - existing slot survives but is weakened
    if (current_score > 0) {
      --slot.score;
    }
    if (status) {
      *status = ProcessStatus::CONTEST_LOST;
    }
    return nullptr;
  }
}

} // namespace ts