FastCheckpointer.hpp

// <FastCheckpointer> -*- C++ -*-
 
#pragma once
 
#include <iostream>
#include <sstream>
#include <stack>
#include <queue>
 
#include "sparta/simulation/TreeNode.hpp"
#include "sparta/functional/ArchData.hpp"
#include "sparta/utils/SpartaException.hpp"
#include "sparta/utils/SpartaAssert.hpp"
 
#include "sparta/serialization/checkpoint/DeltaCheckpoint.hpp"
 
#ifndef DEFAULT_SNAPSHOT_THRESH
#define DEFAULT_SNAPSHOT_THRESH 20
#endif
 
namespace sparta::serialization::checkpoint
{
    class FastCheckpointer : public Checkpointer
    {
    public:
 
        typedef DeltaCheckpoint<storage::VectorStorage> checkpoint_type;
        //typedef DeltaCheckpoint<storage::StringStreamStorage> checkpoint_type;
 
 
        FastCheckpointer(TreeNode& root, Scheduler* sched=nullptr) :
            Checkpointer(root, sched),
            snap_thresh_(DEFAULT_SNAPSHOT_THRESH),
            next_chkpt_id_(checkpoint_type::MIN_CHECKPOINT),
            num_alive_checkpoints_(0),
            num_alive_snapshots_(0),
            num_dead_checkpoints_(0)
        { }
 
        FastCheckpointer(const std::vector<sparta::TreeNode*>& roots, Scheduler* sched=nullptr) :
            Checkpointer(roots, sched),
            snap_thresh_(DEFAULT_SNAPSHOT_THRESH),
            next_chkpt_id_(checkpoint_type::MIN_CHECKPOINT),
            num_alive_checkpoints_(0),
            num_alive_snapshots_(0),
            num_dead_checkpoints_(0)
        { }
 
        ~FastCheckpointer() {
            // Reverse iterate and flag all as free
            for(auto itr = chkpts_.rbegin(); itr != chkpts_.rend(); ++itr){
                checkpoint_type* d = static_cast<checkpoint_type*>(itr->second.get());
                if(!d->isFlaggedDeleted()){
                    d->flagDeleted();
                }
            }
        }
 
 
 
        uint32_t getSnapshotThreshold() const noexcept { return snap_thresh_; }
 
        void setSnapshotThreshold(uint32_t thresh) noexcept {
            snap_thresh_ = thresh;
        }
 
        uint64_t getTotalMemoryUse() const noexcept override {
            uint64_t mem = 0;
            for(auto& cp : chkpts_){
                mem += cp.second->getTotalMemoryUse();
            }
            return mem;
        }
 
        uint64_t getContentMemoryUse() const noexcept override {
            uint64_t mem = 0;
            for(auto& cp : chkpts_){
                mem += cp.second->getContentMemoryUse();
            }
            return mem;
        }
 
 
 
        void deleteCheckpoint(chkpt_id_t id) override {
 
            // Flag checkpoint as deleted
            checkpoint_type* d = findCheckpoint_(id);
            if(!d){
                throw CheckpointError("Could not delete checkpoint ID=")
                    << id << " because no checkpoint by this ID was found";
            }
 
            // Allow deletion and change ID to UNIDENTIFIED_CHECKPOINT.
            // This is still part of a chain though until there are no
            // dependencies on it.
            if(!d->isFlaggedDeleted()){
                num_dead_checkpoints_++;
                if(d->isSnapshot()){
                    num_alive_snapshots_--;
                }
                num_alive_checkpoints_--;
                d->flagDeleted();
            }
 
            // Delete this and all contiguous previous checkpoint which were
            // flagged deleted if possible. Stop if current_ is encountered
            cleanupChain_(d);
        }
 
        void loadCheckpoint(chkpt_id_t id) override {
            checkpoint_type* d = findCheckpoint_(id);
            if(!d){
                throw CheckpointError("Could not load checkpoint ID=")
                    << id << " because no checkpoint by this ID was found";
            }
 
            d->load(getArchDatas());
 
            // Move current to another checkpoint. Anything between head and the
            // old current_ is fair game for removal if allowed
            checkpoint_type* rmv = static_cast<checkpoint_type*>(getCurrent_());
            setCurrent_(d);
 
            // Restore scheduler tick number
            if(sched_){
                sched_->restartAt(getCurrentTick());
            }
 
            // Remove all checkpoints which can be. Stop if the new current_ is
            // encountered again.
            // Note that is is OK if current_ was moved to a later position in
            // the chain. No important checkpoints will be removed. The
            // important thing is never to remove current_.
            cleanupChain_(rmv);
        }
 
        std::vector<chkpt_id_t> getCheckpointsAt(tick_t t) override {
            std::vector<chkpt_id_t> results;
            for(auto& p : chkpts_){
                const Checkpoint* cp = p.second.get();
                const checkpoint_type* dcp = static_cast<const checkpoint_type*>(cp);
                if(cp->getTick() == t && !dcp->isFlaggedDeleted()){
                    results.push_back(cp->getID());
                }
            }
            return results;
        }
 
        std::vector<chkpt_id_t> getCheckpoints() override {
            std::vector<chkpt_id_t> results;
            for(auto& p : chkpts_){
                const Checkpoint* cp = p.second.get();
                const checkpoint_type* dcp = static_cast<const checkpoint_type*>(cp);
                if(!dcp->isFlaggedDeleted()){
                    results.push_back(cp->getID());
                }
            }
            return results;
        }
 
        uint32_t getNumCheckpoints() const noexcept override {
            return num_alive_checkpoints_;
        }
 
        uint32_t getNumSnapshots() const noexcept {
            return num_alive_snapshots_;
        }
 
        uint32_t getNumDeltas() const noexcept {
            return getNumCheckpoints() - getNumSnapshots();
        }
 
        uint32_t getNumDeadCheckpoints() const noexcept {
            return num_dead_checkpoints_;
        }
 
        std::deque<chkpt_id_t> getCheckpointChain(chkpt_id_t id) override {
            std::deque<chkpt_id_t> results;
            if(!getHead()){
                return results;
            }
            const checkpoint_type* d = findCheckpoint_(id);
            if(!d){
                throw CheckpointError("There is no checkpoint with ID ") << id;
            }
            while(d){
                results.push_back(d->getID());
                d = static_cast<checkpoint_type*>(d->getPrev());
            }
            return results;
        }
 
        checkpoint_type* findLatestCheckpointAtOrBefore(tick_t tick,
                                                        chkpt_id_t from) {
            checkpoint_type* d = findCheckpoint_(from);
            if(!d){
                throw CheckpointError("There is no checkpoint with ID ") << from;
            }
 
            // Search backward
            do{
                if(d->getTick() <= tick){
                    break;
                }
                d = static_cast<checkpoint_type*>(d->getPrev());
            }while(d);
 
            return d;
        }
 
        const checkpoint_type* findCheckpoint(chkpt_id_t id) noexcept {
            auto it = chkpts_.find(id);
            if (it != chkpts_.end()) {
                return static_cast<checkpoint_type*>(it->second.get());
            }
            return nullptr;
        }
 
        bool hasCheckpoint(chkpt_id_t id) noexcept override {
            return chkpts_.find(id) != chkpts_.end();
        }
 
        std::vector<chkpt_id_t> getNextIDs(chkpt_id_t id) override final {
            std::vector<chkpt_id_t> next_ids;
            if (const auto chkpt = findCheckpoint_(id)) {
                for (const auto next : chkpt->getNexts()) {
                    const auto dcp = static_cast<checkpoint_type*>(next);
                    if (!dcp->isFlaggedDeleted()) {
                        next_ids.push_back(next->getID());
                    }
                }
            }
            return next_ids;
        }
 
 
 
        std::string stringize() const override {
            std::stringstream ss;
            ss << "<FastCheckpointer on ";
            for (size_t i = 0; i < getRoots().size(); ++i) {
                TreeNode* root = getRoots()[i];
                if (i != 0) {
                    ss << ", ";
                }
                ss << root->getLocation();
            }
            ss << '>';
            return ss.str();
        }
 
        void dumpList(std::ostream& o) override {
            for(auto& cp : chkpts_){
                o << cp.second->stringize() << std::endl;
            }
        }
 
        void dumpData(std::ostream& o) override {
            for(auto& cp : chkpts_){
                cp.second->dumpData(o);
                o << std::endl;
            }
        }
 
        void dumpAnnotatedData(std::ostream& o) override {
            for(auto& cp : chkpts_){
                o << cp.second->stringize() << std::endl;
                cp.second->dumpData(o);
                o << std::endl;
            }
        }
 
        void traceValue(std::ostream& o, chkpt_id_t id, const ArchData* container, uint32_t offset, uint32_t size) override {
            checkpoint_type* dcp = findCheckpoint_(id);
            o << "trace: Searching for 0x" << std::hex << offset << " (" << std::dec << size
              << " bytes) in ArchData " << (const void*)container << " when loading checkpoint "
              << std::dec << id << std::endl;
            if(!dcp){
                o << "trace: Checkpoint " << id << " not found" << std::endl;
            }else{
                dcp->traceValue(o, getArchDatas(), container, offset, size);
            }
        }
 
 
    protected:
 
        void cleanupChain_(checkpoint_type* d) {
 
            // In order to truly delete any checkpoints, we must traverse back
            // to the previous snapshot (or the head) and forward to the another
            // snapshot or the end of the chain.
            // ONLY if both of those points can be reached without encountering
            // a living checkpoint or the current checkpoint (forward
            // only) can the whole chain (including the leading shapshot) be
            // deleted.
 
 
            if(d == getHead()){
                // Cannot delete head of checkpoint tree
                return;
            }
 
            // Walk forward to another snapshot or current
            const bool needed_later = (getCurrent_() == d) || recursForwardFindAlive_(d);
            if(needed_later) {
                // Cannot delete because a later living checkpoint (or current) depends on this
                if(d->isSnapshot()){
                    // This snapshot is needed later. Move to previous delta and work from there
                    d = static_cast<checkpoint_type*>(d->getPrev());
                }else{
                    return; // This delta is needed. Therefore all preceeding deltas are needed
                }
            }
 
            // Delete backward until current, head, or a non-flagged-deleted checkpoint is hit.
            // It is possible to fracture the checkpoint tree by deleting a segment
            // between two snapshots, so prev can end up with nothing leading up to it
            while(d && d != getHead() && d->isFlaggedDeleted()){
 
                // If the checkpoint to delete is the current checkpoint, then
                // We cannot just set current to the previous checkpoint because
                // we may have run forward and storing a checkpoint in the
                // future would depend on the checkpoint we are about to delete.
                // This could be fixed by requiring the next checkpoint to be a
                // spapshot. Instead, point to the flagged-deleted checkpoint
                // and do not delete
                if(getCurrent_() == d){
                    return;
                }
 
                checkpoint_type* prev = static_cast<checkpoint_type*>(d->getPrev());
 
                // If nothing later in the chain (tree) depends on d's data, it can be deleted.
                // This also patches the checkpoint tree around the deleted checkpoint
                if(d->canDelete()) {
                    // Get checkpoint id regardless of whether alive or dead
                    chkpt_id_t id = d->getID();
                    if (d->isFlaggedDeleted()) {
                        id = d->getDeletedID();
                    }
 
                    num_dead_checkpoints_--;
 
                    // Erase element in the map
                    auto itr = chkpts_.find(id);
                    sparta_assert(itr != chkpts_.end());
                    chkpts_.erase(itr);
                }
 
                d = prev; // Continue until head is reached
            }
        }
 
        bool recursForwardFindAlive_(checkpoint_type* d) const
        {
            const std::vector<Checkpoint*> & nexts = d->getNexts();
            for(const auto & chkpt : nexts)
            {
                checkpoint_type* dc = static_cast<checkpoint_type*>(chkpt);
                // Only check descendants for snapshot-ness
                if(dc->isSnapshot()){
                    // Found a live snapshot that ends this branch. d is not needed
                    // after this
                    return false;
                }
                if(dc == getCurrent_()){
                    // Found current in this search chain
                    return true;
                }
                if(dc->isFlaggedDeleted() == false){
                    // Encountered a checkpoint later in the chain that still
                    // depends on this.
                    return true;
                }
 
                // Continue the search recursively
                if(recursForwardFindAlive_(dc)){
                    return true;
                }
            }
 
            // Found nothing alive.
            return false;
        }
 
        checkpoint_type* findCheckpoint_(chkpt_id_t id) noexcept {
            auto itr = chkpts_.find(id);
            if (itr != chkpts_.end()) {
                return static_cast<checkpoint_type*>(itr->second.get());
            }
            return nullptr;
        }
 
        const checkpoint_type* findCheckpoint_(chkpt_id_t id) const noexcept {
            auto itr = chkpts_.find(id);
            if (itr != chkpts_.end()) {
                return static_cast<checkpoint_type*>(itr->second.get());
            }
            return nullptr;
        }
 
        void dumpCheckpointNode_(const chkpt_id_t id, std::ostream& o) override {
            static std::string SNAPSHOT_NOTICE = "(s)";
            auto cp = findCheckpoint_(id);
 
            // Draw data for this checkpoint
            if(cp->isFlaggedDeleted()){
                o << cp->getDeletedRepr();
            }else{
                o << cp->getID();
            }
            // Show that this is a snapshot
            if(cp->isSnapshot()){
                o << ' ' << SNAPSHOT_NOTICE;
            }
        }
 
    private:
 
        void createHead_() override {
            tick_t tick = 0;
            if(sched_){
                tick = sched_->getCurrentTick();
            }
 
            if(getHead()){
                throw CheckpointError("Cannot create head at ")
                    << tick << " because a head already exists in this checkpointer";
            }
            for (auto root : getRoots()) {
                if(root->isFinalized() == false){
                    CheckpointError exc("Cannot create a checkpoint until the tree is finalized. Attempting to checkpoint from node ");
                    exc << root->getLocation() << " at tick ";
                    if(sched_){
                        exc << tick;
                    }else{
                        exc << "<no scheduler>";
                    }
                    throw exc;
                }
            }
 
            checkpoint_type* dcp = new checkpoint_type(getArchDatas(), next_chkpt_id_++, tick, nullptr, true);
            chkpts_[dcp->getID()].reset(dcp);
            setHead_(dcp);
            num_alive_checkpoints_++;
            num_alive_snapshots_++;
            setCurrent_(dcp);
        }
 
        chkpt_id_t createCheckpoint_(bool force_snapshot=false) override {
            bool is_snapshot;
            checkpoint_type* prev;
 
            if(next_chkpt_id_ == checkpoint_type::UNIDENTIFIED_CHECKPOINT){
                throw CheckpointError("Exhausted all ")
                    << checkpoint_type::UNIDENTIFIED_CHECKPOINT << " possible checkpoint IDs. "
                    << "This is likely a gross misuse of checkpointing";
            }
 
            // Caller guarantees a head
            sparta_assert(getHead() != nullptr);
 
            tick_t tick;
            if(sched_){
                tick = sched_->getCurrentTick();
            }else{
                tick = 0;
            }
 
            if(sched_ && (tick < getHead()->getTick())){
                throw CheckpointError("Cannot create a new checkpoint at tick ")
                    << tick << " because this tick number is smaller than the tick number of the head checkpoint at: "
                    << getHead()->getTick() << ". The head checkpoint cannot be reset once created, so it should be done "
                    << "at the start of simulation before running. The simulator front-end should do this so this must "
                    << "likely be fixed in the simulator.";
            }
 
            if(nullptr == getCurrent_()){
                // Creating a delta from the head
                prev = static_cast<checkpoint_type*>(getHead_());
                is_snapshot = false;
            }else{
                if(sched_ && (tick < getCurrent_()->getTick())){
                    throw CheckpointError("Current tick number from sparta scheduler (")
                        << tick << " ) is less than the current checkpoint's tick number ("
                        << getCurrent_()->getTick() << " To create a checkpoint with an earlier tick number, an "
                        << "older checkpoint having a tick number <= the tick number specified here must first be "
                        << "loaded";
                }
 
                // Find latest checkpoint <= tick
 
                prev = static_cast<checkpoint_type*>(getCurrent_());
                is_snapshot = prev->getDistanceToPrevSnapshot() >= getSnapshotThreshold();
            }
 
            checkpoint_type* dcp = new checkpoint_type(getArchDatas(), // Created during createHead
                                                       next_chkpt_id_++,
                                                       tick,
                                                       prev,
                                                       force_snapshot || is_snapshot);
            chkpts_[dcp->getID()].reset(dcp);
            num_alive_checkpoints_++;
            num_alive_snapshots_ += (dcp->isSnapshot() == true);
            setCurrent_(dcp);
 
            if (dcp->isSnapshot()){
                // Clean up starting with this snapshot and moving back.
                // May have an opportunity to free older deltas right now
                // (instead of upon next deletion)
                cleanupChain_(dcp);
            }
 
            return dcp->getID();
        }
 
        std::map<chkpt_id_t, std::unique_ptr<Checkpoint>> chkpts_;
 
        uint32_t snap_thresh_;
 
        chkpt_id_t next_chkpt_id_;
 
        uint32_t num_alive_checkpoints_;
 
        uint32_t num_alive_snapshots_;
 
        uint32_t num_dead_checkpoints_;
    };
 
} // namespace sparta::serialization::checkpoint