PersistentFastCheckpointer.hpp

// <PersistentFastCheckpointer> -*- C++ -*-
 
#pragma once
 
#include <iostream>
#include <sstream>
#include <stack>
#include <queue>
#include <string>
 
#include "sparta/simulation/TreeNode.hpp"
#include "sparta/functional/ArchData.hpp"
#include "sparta/utils/SpartaException.hpp"
#include "sparta/utils/SpartaAssert.hpp"
#include "sparta/serialization/checkpoint/FastCheckpointer.hpp"
 
namespace sparta::serialization::checkpoint
{
    class PersistentFastCheckpointer : public FastCheckpointer
    {
    public:
 
        class FileWriteAdapter
        {
            std::ostream& fs_;
 
        public:
            FileWriteAdapter(std::ostream& out) :
                fs_(out)
            {;}
 
            void dump(std::ostream& o) const {
                o << "<dump not supported on checkpoint file storage adapter>";
            }
 
            uint32_t getSize() const {
                // Return MEMORY size. Assume file is on disk.
                return sizeof(decltype(*this));
            }
 
            void beginLine(ArchData::line_idx_type idx) {
                fs_ << 'L'; // Line start char
 
                ArchData::line_idx_type idx_repr = reorder<ArchData::line_idx_type, LE>(idx);
                fs_.write((char*)&idx_repr, sizeof(ArchData::line_idx_type));
            }
 
            void writeLineBytes(const char* data, size_t size) {
                fs_.write(data, size);
            }
 
            void endArchData() {
                fs_ << "E"; // Indicates end of this checkpoint data
 
                sparta_assert(fs_.good(),
                              "Ostream error while writing checkpoint data");
            }
 
            bool good() const {
                return fs_.good();
            }
        };
 
        class FileReadAdapter
        {
            std::istream& fs_;
 
        public:
            FileReadAdapter(std::istream& in) :
                fs_(in)
            {;}
 
            void dump(std::ostream& o) const {
                o << "<dump not supported on checkpoint file storage adapter>";
            }
 
            uint32_t getSize() const {
                // Return MEMORY size. Assume file is on disk.
                return sizeof(decltype(*this));
            }
 
            void prepareForLoad() {
                fs_.seekg(0); // Seek to start with get pointer before consuming
            }
 
            bool good() const {
                return fs_.good();
            }
 
            ArchData::line_idx_type getNextRestoreLine() {
                char ctrl;
                fs_ >> ctrl;
                sparta_assert(fs_.good(),
                              "Encountered checkpoint data stream error or eof");
                if(ctrl == 'L'){
                    ArchData::line_idx_type ln_idx = 0;
                    fs_.read((char*)&ln_idx, sizeof(ln_idx)); // Presumed LE encoding
                    return ln_idx;
                }else if(ctrl == 'E'){
                    return ArchData::INVALID_LINE_IDX; // Done with restore
                }else{
                    throw SpartaException("Failed to restore a checkpoint because a '")
                        << ctrl << "' control character was found where an 'L' or 'E' was found";
                }
            };
 
            void copyLineBytes(char* buf, uint32_t size) {
                fs_.read(buf, size);
            }
        };
 
 
        PersistentFastCheckpointer(TreeNode& root,
                                   sparta::Scheduler* sched=nullptr) :
            FastCheckpointer(root, sched),
            prefix_("chkpt"),
            suffix_("data")
        { }
 
        virtual ~PersistentFastCheckpointer() {
 
        }
 
        FastCheckpointer::chkpt_id_t save(std::ostream& outf) {
            FastCheckpointer::chkpt_id_t checkpoint_id = createCheckpoint(true);
            save_(outf);
            return checkpoint_id;
        }
 
        FastCheckpointer::chkpt_id_t save(std::string filename) {
            FastCheckpointer::chkpt_id_t checkpoint_id = createCheckpoint(true);
            std::ofstream outf(filename, std::ofstream::out | std::ofstream::binary | std::ofstream::trunc);
            save_(outf);
            outf.close();
            return checkpoint_id;
        }
 
        FastCheckpointer::chkpt_id_t save() {
            const bool force_snapshot = true;
            FastCheckpointer::chkpt_id_t checkpoint_id = createCheckpoint(force_snapshot);
            std::ostringstream chkpt_filename;
            chkpt_filename << prefix_ << "."
                           << checkpoint_id
                           << "." << suffix_;
            std::ofstream outf(chkpt_filename.str(), std::ofstream::binary | std::ofstream::trunc);
            save_(outf);
            outf.close();
            return checkpoint_id;
        }
 
        void restore(std::istream& in) {
            auto adatas = getArchDatas();
            FileReadAdapter fsa(in);
            for (auto aditr=adatas.begin(); aditr!= adatas.end(); aditr++) {
                (*aditr)->restoreAll(fsa);
            }
        }
 
        void restore(const std::string& filename) {
            std::ifstream in(filename);
            restore(in);
            in.close();
        }
 
    private:
 
        void save_(std::ostream& outf) {
            // Throw on write failure
            outf.exceptions(std::ostream::eofbit | std::ostream::badbit |
                            std::ostream::failbit | std::ostream::goodbit);
            FileWriteAdapter fsa(outf);
            auto adatas = getArchDatas();
            for (auto aditr=adatas.begin(); aditr!= adatas.end(); aditr++) {
                (*aditr)->saveAll(fsa);
            }
        }
 
        std::string prefix_;
        std::string suffix_;
 
    };
 
} // namespace sparta::serialization::checkpoint