lsm_tree.h

#ifndef STORAGE_LSM_TREE_H
#define STORAGE_LSM_TREE_H
 
#include <vkdb/concepts.h>
#include <vkdb/sstable.h>
#include <vkdb/mem_table.h>
#include <vkdb/write_ahead_log.h>
#include <vkdb/lru_cache.h>
#include <vkdb/wal_lsm.h>
 
#include <algorithm>
#include <ranges>
#include <deque>
#include <set>
#include <functional>
 
namespace vkdb {
using TimeSeriesKeyFilter = std::function<bool(const TimeSeriesKey&)>;
 
static const TimeSeriesKeyFilter TRUE_TIME_SERIES_KEY_FILTER =
  [](const TimeSeriesKey&) { return true; };
 
template <ArithmeticNoCVRefQuals TValue>
class LSMTree {
public:
  using key_type = TimeSeriesKey;
  using mapped_type = std::optional<TValue>;
  using value_type = std::pair<const key_type, mapped_type>;
  using size_type = uint64_t;
  using table_type = typename MemTable<TValue>::table_type;
 
  static constexpr size_type LAYER_COUNT{8};
 
  static constexpr size_type CACHE_CAPACITY{10'000};
 
  explicit LSMTree(FilePath path) noexcept
    : wal_{path}
    , path_{std::move(path)}
    , sstable_id_{0}
    , ck_layers_{LAYER_COUNT}
    , cache_{CACHE_CAPACITY} {
      std::filesystem::create_directories(path_);
      load_sstables();
    }
 
  LSMTree(LSMTree&&) noexcept = default;
  
  LSMTree& operator=(LSMTree&&) noexcept = default;  
 
  LSMTree(const LSMTree&) = delete;
  
  LSMTree& operator=(const LSMTree&) = delete;
 
  ~LSMTree() noexcept = default;
 
  void put(const key_type& key, const TValue& value, bool log = true) {
    if (log) {
      wal_.append({WALRecordType::PUT, {key, value}});
    }
    mem_table_.put(key, value);
    dirty_table_[key] = true;
    if (mem_table_.size() == MemTable<TValue>::C0_LAYER_SSTABLE_MAX_ENTRIES) {
      flush();
    }
  }
 
  void remove(const key_type& key, bool log = true) {
    if (log) {
      wal_.append({WALRecordType::REMOVE, {key, std::nullopt}});
    }
    mem_table_.put(key, std::nullopt);
    dirty_table_[key] = true;
    if (mem_table_.size() == MemTable<TValue>::C0_LAYER_SSTABLE_MAX_ENTRIES) {
      flush();
    }
  }
 
  [[nodiscard]] mapped_type get(const key_type& key) const noexcept {
    if (!dirty_table_[key] && cache_.contains(key)) {
      return cache_.get(key);
    }
 
    try {
      return search_memtable(key);
    } catch (const std::invalid_argument& e) {
      const auto c0_value{iterative_layer_search(0, key)};
      if (c0_value.has_value()) {
        return c0_value;
      }
 
      for (size_type k{1}; k < LAYER_COUNT; ++k) {
        const auto ck_value{binary_layer_search(k, key)};
        if (ck_value.has_value()) {
          return ck_value;
        }
      }
 
      return std::nullopt;
    }
  }
 
  [[nodiscard]] std::vector<value_type> getRange(
    const key_type& start,
    const key_type& end,
    TimeSeriesKeyFilter&& filter
  ) const {
    table_type entry_table;
    for (size_type k{LAYER_COUNT - 1}; k > 0; --k) {
      binary_layer_search_range(k, start, end, filter, entry_table);
    }
    for (const auto& sstable : ck_layers_[0]) {
      update_entries_with_sstable_range(
        sstable, start, end, filter, entry_table
      );
    }
    for (const auto& [key, value] : mem_table_.getRange(start, end)) {
      update_entries_with_filtered_key(key, value, filter, entry_table);
    }
    return {entry_table.begin(), entry_table.end()};
  }
 
  void replayWAL() {
    wal_.replay(*this);
  }
 
  void clear() noexcept {
    for (const auto& ck_layer : ck_layers_) {
      for (const auto& sstable : ck_layer) {
        std::filesystem::remove(sstable.path());
        std::filesystem::remove(sstable.metadataPath());
      }
    }
    std::filesystem::remove(wal_.path());
    mem_table_.clear();
    ck_layers_ = CkLayers{LAYER_COUNT};
    wal_.clear();
    sstable_id_ = std::array<size_type, LAYER_COUNT>{0};
    cache_.clear();
    dirty_table_.clear();
  }
 
  [[nodiscard]] std::string str() const noexcept {
    std::stringstream ss;
    ss << mem_table_.str();
    for (const auto& ck_layer : ck_layers_) {
      for (const auto& sstable : ck_layer) {
        ss << sstable.str();
      }
    }
    return ss.str();
  }
 
  [[nodiscard]] size_type sstableCount() const noexcept {
    size_type count{0};
    for (const auto& ck_layer : ck_layers_) {
      count += ck_layer.size();
    }
    return count;
  }
 
  [[nodiscard]] size_type sstableCount(size_type k) const {
    if (k < 0 || k >= LAYER_COUNT) {
      throw std::out_of_range{
        "LSMTree::sstableCount(): Layer index out of range."
      };
    }
    return ck_layers_[k].size();
  }
 
  [[nodiscard]] bool empty() const noexcept {
    auto empty{mem_table_.empty()};
    for (const auto& ck_layer : ck_layers_) {
      empty = empty && ck_layer.empty();
      if (!empty) {
        break;
      }
    }
    return empty;
  }
 
private:
  using SSTables = std::deque<SSTable<TValue>>;
 
  using CkLayer = SSTables;
 
  using CkLayers = std::vector<CkLayer>;
 
  using TimeWindow = TimeRange;
 
  using Cache = LRUCache<key_type, TValue>;
 
  using DirtyTable = std::unordered_map<key_type, bool>;
 
  using TimeWindowToEntriesMap = std::map<TimeWindow, table_type, std::greater<>>;
 
  static constexpr std::array<size_type, LAYER_COUNT> CK_LAYER_WINDOW_SIZE{
    0,
    86'400,
    604'800,
    2'592'000,
    7'776'000,
    15'552'000,
    31'536'000,
    94'608'000
  };
 
  static constexpr std::array<size_type, LAYER_COUNT> CK_LAYER_TABLE_COUNT{
    32,
    2'048,
    1'024,
    512,
    256,
    128,
    64,
    32
  };
 
  [[nodiscard]] FilePath get_next_file_path(size_type k) noexcept {
    return path_ / ("sstable_l" + std::to_string(k) + "_"
      + std::to_string(sstable_id_[k]++) + ".sst");
  }
 
  void load_sstables() {
    std::array<std::set<FilePath>, LAYER_COUNT> sstable_files;
    for (const auto& file : std::filesystem::directory_iterator(path_)) {
      if (!file.is_regular_file() || file.path().extension() != ".sst") {
        continue;
      }
      const auto l_pos{file.path().filename().string().find("_l")};
      const auto layer_idx{
        std::stoull(file.path().filename().string().substr(l_pos + 2))
      };
      sstable_files[layer_idx].insert(file.path());
      const auto id{
        std::stoull(file.path().filename().string().substr(l_pos + 4))
      };
      sstable_id_[layer_idx] = std::max(
        sstable_id_[layer_idx],
        static_cast<size_type>(id + 1)
      );
    }
    for (size_type k{0}; k < LAYER_COUNT; ++k) {
      for (const auto& sstable_file : sstable_files[k]) {
        ck_layers_[k].emplace_back(std::move(sstable_file));
      }
    }
    compact();
  }
 
  void validate_layer_index(size_type k) const {
    if (k < 0 || k >= LAYER_COUNT) {
      throw std::out_of_range{
        "LSMTree::validate_layer_index(): Layer index out of range."
      };
    }
  }
 
  void reset_layer(size_type k) noexcept {
    validate_layer_index(k);
    ck_layers_[k].clear();
    sstable_id_[k] = 0;
  }
 
  void flush() {
    FilePath sstable_file_path{get_next_file_path(0)};
    ck_layers_[0].emplace_back(
      sstable_file_path,
      std::move(mem_table_),
      mem_table_.size()
    );
    compact();
    mem_table_.clear();
    wal_.clear();
  }
 
  void update_entries_with_filtered_key(
    const key_type& key,
    const mapped_type& value,
    const TimeSeriesKeyFilter& filter,
    table_type& entry_table
  ) const noexcept {
    if (!filter(key)) {
      return;
    }
    if (!value.has_value()) {
      entry_table.erase(key);
      return;
    }
    entry_table[key] = value;
  }
 
  void update_entries_with_sstable_range(
    const SSTable<TValue>& sstable,
    const key_type& start,
    const key_type& end,
    const TimeSeriesKeyFilter& filter,
    table_type& entry_table
  ) const noexcept {
    for (const auto& [key, value] : sstable.getRange(start, end)) {
      update_entries_with_filtered_key(key, value, filter, entry_table);
    }
  }
 
  mapped_type search_memtable(const key_type& key) const {
    const auto mem_table_value{mem_table_.get(key)};
    cache_.put(key, mem_table_value);
    dirty_table_[key] = false;
    return mem_table_value;
  }
 
  mapped_type iterative_layer_search(
    size_type k,
    const key_type& key
  ) const {
    validate_layer_index(k);
    for (const auto& sstable : ck_layers_[k] | std::views::reverse) {
      const auto sstable_value{sstable.get(key)};
      if (!sstable_value.has_value()) {
        continue;
      }
      cache_.put(key, sstable_value);
      dirty_table_[key] = false;
      return sstable_value;
    }
    return std::nullopt;
  }
 
  mapped_type binary_layer_search(
    size_type k,
    const key_type& key
  ) const {
    const auto timestamp{key.timestamp()};
    const auto& ck_layer{ck_layers_[k]};
    auto sstable_it{std::lower_bound(
      ck_layer.begin(),
      ck_layer.end(),
      timestamp,
      [](const auto& sstable, const auto target_time) {
        return target_time < sstable.timeRange().lower();
      }
    )};
 
    if (
      sstable_it != ck_layer.end() && 
      sstable_it->timeRange().lower() <= timestamp && 
      timestamp <= sstable_it->timeRange().upper()
    ) {
      const auto sstable_value{sstable_it->get(key)};
      if (sstable_value.has_value()) {
        cache_.put(key, sstable_value);
        dirty_table_[key] = false;
        return sstable_value;
      }
      return std::nullopt;
    }
    return std::nullopt;
  }
 
 
  void binary_layer_search_range(
    size_type k,
    const key_type& start,
    const key_type& end,
    const TimeSeriesKeyFilter& filter,
    table_type& entry_table
  ) const {
    const auto& ck_layer{ck_layers_[k]};
    auto start_it{std::lower_bound(
      ck_layer.begin(),
      ck_layer.end(),
      start.timestamp(),
      [](const auto& sstable, const auto target_time) {
        return sstable.timeRange().upper() < target_time;
      }
    )};
    
    auto end_it{std::upper_bound(
      start_it,
      ck_layer.end(),
      end.timestamp(),
      [](const auto target_time, const auto& sstable) {
        return target_time < sstable.timeRange().lower();
      }
    )};
 
    for (
      const auto& sstable :
      std::ranges::subrange(start_it, end_it) | std::views::reverse
    ) {
      update_entries_with_sstable_range(
        sstable, start, end, filter, entry_table
      );
    }
  }
 
  void compact() {
    compact_layer_if_needed(0);
  }
 
  void compact_layer_if_needed(size_type k) {
    validate_layer_index(k);
    if (k == LAYER_COUNT - 1) {
      return;
    }
    if (ck_layers_[k].size() <= CK_LAYER_TABLE_COUNT[k]) {
      return;
    }
    compact_layer(k);
    compact_layer_if_needed(k + 1);
  }
 
  void compact_layer(size_type k) {
    if (k == 0) {
      merge_all_in_layer(k);
      return;
    }
    while (ck_layers_[k].size() > CK_LAYER_TABLE_COUNT[k]) {
      merge_oldest_in_layer(k);
    }
  }
 
  void update_entries_with_sstable(
    SSTable<TValue>&& sstable,
    size_type window_size,
    TimeWindowToEntriesMap& time_window_to_entries,
    std::vector<FilePath>& files_to_remove
  ) const noexcept {
    for (const auto& [key, value] : sstable.entries()) {
      const auto start{(key.timestamp() / window_size) * window_size};
      const auto time_window{TimeWindow{start, start + window_size}};
      time_window_to_entries[time_window][key] = value;
    }
    files_to_remove.push_back(sstable.path());
    files_to_remove.push_back(sstable.metadataPath());
  }
 
  void merge_all_in_layer(size_type k) {
    auto& curr_layer{ck_layers_[k]};
    auto& next_layer{ck_layers_[k + 1]};
    const auto window_size{CK_LAYER_WINDOW_SIZE[k + 1]};
 
    TimeWindowToEntriesMap time_window_to_entries;
    std::vector<FilePath> files_to_remove;
 
    for (auto&& sstable : next_layer) {
      update_entries_with_sstable(
        std::move(sstable), window_size,
        time_window_to_entries,
        files_to_remove
      );
    }
    reset_layer(k + 1);
 
    for (auto&& sstable : curr_layer) {
      update_entries_with_sstable(
        std::move(sstable), window_size,
        time_window_to_entries,
        files_to_remove
      );
    }
    reset_layer(k);
 
    for (const auto& file : files_to_remove) {
      if (!std::filesystem::remove(file)) {
        throw std::runtime_error{
          "LSMTree::merge_all_in_layer(): Failed to remove file: '"
          + std::string{file} + "' after compaction."
        };
      }
    }
 
    for (auto& [time_window, entries] : time_window_to_entries) {
      auto sstable{merge_entries(std::move(entries), k)};
      next_layer.push_back(std::move(sstable));
    }
  }
 
  void merge_oldest_in_layer(size_type k) {
    auto& curr_layer{ck_layers_[k]};
    auto& next_layer{ck_layers_[k + 1]};
    const auto window_size{CK_LAYER_WINDOW_SIZE[k + 1]};
 
    TimeWindowToEntriesMap time_window_to_entries;
    std::vector<FilePath> files_to_remove;
 
    for (auto&& sstable : next_layer) {
      update_entries_with_sstable(
        std::move(sstable), window_size,
        time_window_to_entries,
        files_to_remove
      );
    }
    reset_layer(k + 1);
 
    while (curr_layer.size() > CK_LAYER_TABLE_COUNT[k]) {
      auto& oldest_sstable{curr_layer.front()};
      update_entries_with_sstable(
        std::move(oldest_sstable), window_size,
        time_window_to_entries,
        files_to_remove
      );
      curr_layer.pop_front();
    }
 
    for (const auto& file : files_to_remove) {
      if (!std::filesystem::remove(file)) {
        throw std::runtime_error{
          "LSMTree::merge_oldest_in_layer(): Failed to remove file: '"
          + std::string{file} + "' after compaction."
        };
      }
    }
 
    for (auto& [time_window, entries] : time_window_to_entries) {
      auto sstable{merge_entries(std::move(entries), k)};
      next_layer.push_back(std::move(sstable));
    }
  }
 
  SSTable<TValue> merge_entries(table_type&& entries, size_type k) {
    auto memtable{MemTable<TValue>{std::move(entries)}};
    const auto memtable_size{memtable.size()};
    auto sstable{SSTable<TValue>{
      get_next_file_path(k + 1),
      std::move(memtable),
      memtable_size
    }};
    return sstable;
  }
 
  MemTable<TValue> mem_table_;
 
  CkLayers ck_layers_;
 
  WriteAheadLog<TValue> wal_;
 
  FilePath path_;
 
  std::array<size_type, LAYER_COUNT> sstable_id_;
 
  mutable Cache cache_;
 
  mutable DirtyTable dirty_table_;
};
}  // namespace vkdb
 
#endif // STORAGE_LSM_TREE_H