session_residency.h

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#pragma once
 
#include <cstddef>
#include <cstdint>
#include <string>
#include <unordered_map>
#include <vector>
 
namespace entropic {
 
constexpr int kNoSessionSlot = -1;
 
template <typename Tok>
class SessionResidency {
public:
    explicit SessionResidency(int slots = 1)
        : slots_(slots < 1 ? 1 : slots), resident_(
              static_cast<std::size_t>(slots < 1 ? 1 : slots)) {}
 
    int slot_for(const std::string& key) const {
        auto it = assigned_.find(key);
        return it == assigned_.end() ? kNoSessionSlot : it->second;
    }
 
    int acquire(const std::string& key, std::string* evicted) {
        if (evicted != nullptr) { evicted->clear(); }
        auto it = assigned_.find(key);
        if (it != assigned_.end()) {
            touch(key);
            return it->second;
        }
        int slot = free_slot();
        if (slot == kNoSessionSlot) {
            const std::string victim = lru_key();
            slot = assigned_[victim];
            forget(victim);
            if (evicted != nullptr) { *evicted = victim; }
        }
        assigned_[key] = slot;
        touch(key);
        return slot;
    }
 
    const std::vector<Tok>& resident(int slot) const {
        static const std::vector<Tok> kEmpty;
        return valid(slot) ? resident_[static_cast<std::size_t>(slot)]
                           : kEmpty;
    }
 
    void set_resident(int slot, std::vector<Tok> tokens) {
        if (valid(slot)) {
            resident_[static_cast<std::size_t>(slot)] = std::move(tokens);
        }
    }
 
    void invalidate(int slot) {
        if (valid(slot)) {
            resident_[static_cast<std::size_t>(slot)].clear();
        }
    }
 
    void invalidate_all() {
        for (auto& v : resident_) { v.clear(); }
    }
 
    int forget_session(const std::string& key) {
        auto it = assigned_.find(key);
        if (it == assigned_.end()) { return kNoSessionSlot; }
        const int slot = it->second;
        invalidate(slot);
        forget(key);
        return slot;
    }
 
    int slots() const { return slots_; }
 
    std::size_t assigned_count() const { return assigned_.size(); }
 
private:
    bool valid(int slot) const {
        return slot >= 0 && slot < slots_;
    }
 
    void touch(const std::string& key) { use_order_[key] = ++tick_; }
 
    void forget(const std::string& key) {
        assigned_.erase(key);
        use_order_.erase(key);
    }
 
    int free_slot() const {
        std::vector<bool> taken(static_cast<std::size_t>(slots_), false);
        for (const auto& [k, s] : assigned_) {
            if (s >= 0 && s < slots_) {
                taken[static_cast<std::size_t>(s)] = true;
            }
        }
        int found = kNoSessionSlot;
        for (int i = 0; i < slots_ && found == kNoSessionSlot; ++i) {
            if (!taken[static_cast<std::size_t>(i)]) { found = i; }
        }
        return found;
    }
 
    std::string lru_key() const {
        std::string oldest;
        uint64_t best = UINT64_MAX;
        for (const auto& [k, t] : use_order_) {
            if (t < best) { best = t; oldest = k; }
        }
        return oldest;
    }
 
    int slots_;
    std::vector<std::vector<Tok>> resident_;
    std::unordered_map<std::string, int> assigned_;
    std::unordered_map<std::string, uint64_t> use_order_;
    uint64_t tick_ = 0;
};
 
}  // namespace entropic