Checkpointer.hpp

// <Checkpointer> -*- C++ -*-
 
#pragma once
 
#include <iostream>
#include <sstream>
#include <stack>
#include <queue>
#include <list>
#include <map>
#include <string>
#include <memory>
 
#include "sparta/simulation/TreeNode.hpp"
#include "sparta/functional/ArchData.hpp"
#include "sparta/utils/SpartaException.hpp"
#include "sparta/utils/SpartaAssert.hpp"
#include "sparta/simulation/TreeNodePrivateAttorney.hpp"
#include "sparta/serialization/checkpoint/Checkpoint.hpp"
 
namespace sparta::serialization::checkpoint
{
    class Checkpointer
    {
    public:
 
 
        typedef Checkpoint::tick_t tick_t;
 
        typedef Checkpoint::chkpt_id_t chkpt_id_t;
 
 
 
        Checkpointer(TreeNode& root, sparta::Scheduler* sched=nullptr):
            sched_(sched)
        {
            roots_.push_back(&root);
        }
 
        Checkpointer(const std::vector<TreeNode*>& roots, sparta::Scheduler* sched=nullptr):
            sched_(sched),
            roots_(roots)
        { }
 
        virtual ~Checkpointer() {
        }
 
 
 
        const std::vector<TreeNode*> & getRoots() const noexcept { return roots_; }
 
        const Scheduler* getScheduler() const noexcept { return sched_; }
 
        virtual uint64_t getTotalMemoryUse() const noexcept = 0;
 
        virtual uint64_t getContentMemoryUse() const noexcept = 0;
 
        uint64_t getTotalCheckpointsCreated() const noexcept {
            return total_chkpts_created_;
        }
 
 
 
        void createHead() {
            tick_t tick = 0;
            if(sched_){
                tick = sched_->getCurrentTick();
            }
            if(head_){
                CheckpointError exc("Cannot create head at ");
                if(sched_){
                    exc << tick;
                }else{
                    exc << "<no scheduler>";
                }
                exc << " because a head already exists in this checkpointer";
                throw exc;
            }
            for (auto root : roots_) {
                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;
                }
            }
 
            enumerateArchDatas_(); // Determines which ArchDatas are required and populates adatas_
 
            createHead_();
 
            sparta_assert(head_ != nullptr, "A call to createHead_ must create a head and invoke setHead_ or throw an exception");
 
            total_chkpts_created_++;
        }
 
        chkpt_id_t createCheckpoint(bool force_snapshot=false) {
            if(!head_){
                createHead();
            }
 
            chkpt_id_t id = createCheckpoint_(force_snapshot);
            total_chkpts_created_++;
            return id;
        }
 
        virtual void deleteCheckpoint(chkpt_id_t id) = 0;
 
        virtual void loadCheckpoint(chkpt_id_t id) = 0;
 
        void forgetCurrent() {
            if(head_){
                current_ = head_;
            }
        }
 
        virtual std::vector<chkpt_id_t> getCheckpointsAt(tick_t t) = 0;
 
        virtual std::vector<chkpt_id_t> getCheckpoints() = 0;
 
        virtual uint32_t getNumCheckpoints() const noexcept = 0;
 
        virtual std::deque<chkpt_id_t> getCheckpointChain(chkpt_id_t id) = 0;
 
        virtual bool hasCheckpoint(chkpt_id_t id) noexcept = 0;
 
        const CheckpointBase* getHead() const noexcept {
            return head_;
        }
 
        chkpt_id_t getHeadID() const noexcept {
            if(!head_){
                return Checkpoint::UNIDENTIFIED_CHECKPOINT;
            }
            return head_->getID();
        }
 
        chkpt_id_t getCurrentID() const {
            if(current_){
                return current_->getID();
            }
            sparta_assert(!head_); // If there was no current, it can only be because there is no head yet
            return Checkpoint::UNIDENTIFIED_CHECKPOINT;
        }
 
        tick_t getCurrentTick() const {
            if(current_){
                return current_->getTick();
            }
            sparta_assert(!head_);
            return 0;
        }
 
        virtual std::vector<chkpt_id_t> getNextIDs(chkpt_id_t id) = 0;
 
 
 
        virtual std::string stringize() const {
            std::stringstream ss;
            ss << "<Checkpointer on ";
            for (size_t i = 0; i < roots_.size(); ++i) {
                TreeNode* root = roots_[i];
                if (i != 0) {
                    ss << ", ";
                }
                ss << root->getLocation();
            }
            ss << ">";
            return ss.str();
        }
 
        virtual void dumpList(std::ostream& o) = 0;
 
        virtual void dumpData(std::ostream& o) = 0;
 
        virtual void dumpAnnotatedData(std::ostream& o) = 0;
 
        virtual void traceValue(std::ostream& o, chkpt_id_t id, const ArchData* container, uint32_t offset, uint32_t size) = 0;
 
        void dumpTree(std::ostream& o) {
            std::deque<uint32_t> c;
            dumpBranch(o, getHeadID(), 0, 0, c);
            o << '\n';
        }
 
        void dumpBranch(std::ostream& o,
                        const chkpt_id_t chkpt,
                        uint32_t indent,
                        uint32_t pos,
                        std::deque<uint32_t>& continues) {
            static const std::string SEP_STR = "-> "; // Normal checkpoint chain
            static const std::string CONT_SEP_STR = "`> "; // Checkpoint branch from higher line
            assert(SEP_STR.size() == CONT_SEP_STR.size()); // Must be the same length so the layout looks OK
 
            // Walk through indent and draw continuations
            uint32_t i = pos;
            std::remove_reference<decltype(continues)>::type::const_iterator next_cont = continues.begin();
            for(; i < indent; ++i){
                if(next_cont != continues.end() && i == *next_cont){
                    o << '|';
                    while(next_cont != continues.end() && *next_cont == i){ // Skip duplicates and move on to next
                        ++next_cont;
                    }
                }else{
                    o << ' ';
                }
            }
 
            auto nexts = getNextIDs(chkpt);
            std::stringstream ss;
 
            // Draw separator between prev checkpoint and this
            if(next_cont != continues.end() && *next_cont == indent && indent != pos){
                ss << CONT_SEP_STR;
            }else{
                ss << SEP_STR;
            }
 
            // Draw box around object if it is current
            if(current_ && current_->getID() == chkpt){
                ss << "[ ";
            }
 
            dumpCheckpointNode_(chkpt, ss);
            ss << ' ';
 
            if(current_ && current_->getID() == chkpt){
                ss << ']';
            }
 
            o << ss.str();
            i += ss.str().size();
 
            // Draw all next checkpoints recursively
            if(nexts.size() > 0){
                if(nexts.size() > 1){
                    continues.push_back(i); // Push continues first
                }
                dumpBranch(o, nexts.front(), i, i, continues); // No indent, no newline
                decltype(nexts)::const_iterator itr = nexts.begin() + 1;
                if(itr == nexts.end()){
                    //continues.pop_back();
                }else{
                    while(itr < nexts.end()){
                        if(itr+1 == nexts.end()){
                            continues.pop_back(); // Do not show this continuation here
                        }
                        o << '\n';
                        dumpBranch(o, *itr, i, 0, continues); // Indent
                        ++itr;
                    }
                }
            }
        }
 
 
    protected:
 
        virtual void createHead_() = 0;
 
        virtual chkpt_id_t createCheckpoint_(bool force_snapshot=false) = 0;
 
        virtual void dumpCheckpointNode_(const chkpt_id_t id, std::ostream& o) {
            o << id;
        }
 
        const std::vector<ArchData*>& getArchDatas() const {
            return adatas_;
        }
 
        CheckpointBase* getHead_() noexcept {
            return head_;
        }
 
        const CheckpointBase* getHead_() const noexcept {
            return head_;
        }
 
        void setHead_(CheckpointBase* head) {
            sparta_assert(head != nullptr, "head argument in setHead_ cannot be nullptr");
            sparta_assert(head_ == nullptr, "Cannot setHead_ again on a Checkpointer once heas is already set");
            head_ = head;
        }
 
        CheckpointBase* getCurrent_() const noexcept {
            return current_;
        }
 
        void setCurrent_(CheckpointBase* current) {
            sparta_assert(current != nullptr,
                        "Can never setCurrent_ to nullptr except. A null current is a valid state at initialization only")
            current_ = current;
        }
 
        Scheduler* const sched_;
 
    private:
 
        void enumerateArchDatas_() {
            sparta_assert(adatas_.size() == 0, "FastCheckpointer already has a vector of ArchDatas. Cannot re-enumerate");
 
            adatas_.clear(); // Clear in case invoked again
 
            // This is a helper for building the adatas_ vector while detecting
            // dupicate ArchDatas.
            // Maps an ArchData key to the TreeNode which pointed to it.
            std::map<ArchData*,TreeNode*> adatas_helper;
 
            // Recursively walk the tree and add all ArchDatas to adatas_
            for(TreeNode* root : roots_) {
                recursAddArchData_(root, adatas_helper);
            }
        }
 
        void recursAddArchData_(TreeNode* n, std::map<ArchData*,TreeNode*>& adatas_helper) {
            assert(n);
            auto adatas = n->getAssociatedArchDatas();
            for(ArchData* ad : adatas){
                if(ad != nullptr){
                    auto itr = adatas_helper.find(ad);
                    if(itr != adatas_helper.end()){
                        throw CheckpointError("Found a second reference to ArchData ")
                            << ad << " in the checkpointer: " << stringize() << " . First reference found throgh "
                            << itr->second->getLocation() << " and second found through " << n->getLocation()
                            << " . An ArchData should be findable throug exactly 1 TreeNode";
                    }
                    adatas_.push_back(ad);
                    adatas_helper[ad] = n; // Add to helper map
                }
            }
            for(TreeNode* child : TreeNodePrivateAttorney::getAllChildren(n)){
                recursAddArchData_(child, adatas_helper);
            }
        }
 
        std::vector<TreeNode*> roots_; 
 
        CheckpointBase* head_ = nullptr;
 
        std::vector<ArchData*> adatas_;
 
        CheckpointBase* current_ = nullptr;
 
        uint64_t total_chkpts_created_ = 0;
    };
 
} // namespace sparta::serialization::checkpoint
 
 
inline std::ostream& operator<<(std::ostream& o, const sparta::serialization::checkpoint::Checkpointer& cpr){
    o << cpr.stringize();
    return o;
}
 
inline std::ostream& operator<<(std::ostream& o, const sparta::serialization::checkpoint::Checkpointer* cpr){
    if(cpr == 0){
        o << "null";
    }else{
        o << cpr->stringize();
    }
    return o;
}