DETCheckpoint.h

Go to the documentation of this file.

#pragma once
 
#include <string>
#include <map>
#include <vector>
#include <memory>
#include <fstream>
#include <sstream>
#include <iostream>
#include <filesystem>
#include <cstdlib>
 
#include "RestartNodeInfo.h"
#include "../Branch/BranchRecord.h"
#include "../ParallelExec/ParallelGlobal.h"
 
#include <boost/archive/binary_oarchive.hpp>
#include <boost/archive/binary_iarchive.hpp>
#include <boost/serialization/map.hpp>
#include <boost/serialization/vector.hpp>
#include <boost/serialization/string.hpp>
#include <boost/serialization/shared_ptr.hpp>
 
struct PendingBranch {
    int cid = 0;                   
    int branch_index = 0;          
    int sequence_id = 0;           
    int parent_sequence_id = -1;   
 
    template<class Archive>
    void serialize(Archive& ar, unsigned int) {
        ar & cid & branch_index & sequence_id & parent_sequence_id;
    }
};
 
struct DETCheckpoint {
    static constexpr int VERSION = 9;
    int version = VERSION;          
    int sequences_completed = 0;    
    int sequence_counter = 0;       
    uint64_t config_fingerprint = 0;
    bool compress_queue = false;
    bool blobs_compressed = false;
    // RELAP restart blobs are stored ONCE per source and referenced per CID.
    // The MPI writer moves each blob out of its RestartNodeInfo; the serial
    // writer copies it from g_sim.restart_buffers or the backup file.
    std::map<int, std::vector<char>> source_blobs;  
    std::map<int, int> cid_source;                  
    int max_branch_id = -1;
 
    std::string branch_center_blob;
 
    // Crew snapshots are archived once per UNIQUE RestartNodeInfo as an LZ4
    // blob, with index maps that rebuild restart_states and
    // task_restart_states at load. The two containers alias the same
    // shared_ptrs (carried-task registration shares them), so save() dedups by
    // pointer and load() restores the aliasing. These three members are
    // populated only inside save() and load().
    std::vector<std::string> packed_rinfos;       
    std::map<int, int> restart_state_pack_ref;    
    std::vector<int> task_state_pack_ref;         
 
    // CID-keyed maps
    std::map<int, std::shared_ptr<RestartNodeInfo>> restart_states;  
    std::map<int, std::string> branch_record_blobs;                  
 
    // Pending tasks with per-task data (parallel vectors, same index)
    std::vector<PendingBranch> pending_tasks;                          
    std::vector<std::shared_ptr<RestartNodeInfo>> task_restart_states; 
    std::vector<std::string> task_branch_blobs;                        
 
    std::vector<SequenceReport> completed_reports;  
 
    template<class Archive>
    void serialize(Archive& ar, unsigned int) {
        ar & version;
        ar & sequences_completed;
        ar & sequence_counter;
        if (version >= 3) {
            ar & compress_queue;
        }
        if (version >= 4) {
            ar & source_blobs;
            ar & cid_source;
            ar & max_branch_id;
        }
        if (version >= 5) {
            ar & blobs_compressed;
        }
        if (version >= 6) {
            ar & config_fingerprint;
        }
        if (version >= 7) {
            ar & packed_rinfos;
            ar & restart_state_pack_ref;
            ar & task_state_pack_ref;
        }
        ar & branch_center_blob;
        if (version < 7) {
            ar & restart_states;
        }
        ar & branch_record_blobs;
        ar & pending_tasks;
        if (version < 7) {
            ar & task_restart_states;
        }
        ar & task_branch_blobs;
        ar & completed_reports;
    }
 
    void capture_branch_center(const BranchRecord& br) {
        std::ostringstream oss(std::ios::binary);
        boost::archive::binary_oarchive oa(oss);
        oa << br;
        branch_center_blob = oss.str();
    }
 
    void restore_branch_center(BranchRecord& br) const {
        std::istringstream iss(branch_center_blob, std::ios::binary);
        boost::archive::binary_iarchive ia(iss);
        ia >> br;
    }
 
    static bool version_supported(int v) { return v >= 9; }
 
    bool config_matches(uint64_t current) const {
        return version < 6 || config_fingerprint == 0 || current == 0
            || config_fingerprint == current;
    }
 
    static std::string lz4_pack_string(const std::string& s) {
        auto v = lz4_compress(s.data(), static_cast<int>(s.size()));
        return std::string(v.data(), v.size());
    }
    static std::string lz4_unpack_string(const std::string& s) {
        auto v = lz4_decompress(s.data(), static_cast<int>(s.size()));
        return std::string(v.data(), v.size());
    }
 
    static std::string pack_rinfo(const RestartNodeInfo& ri) {
        std::ostringstream oss(std::ios::binary);
        {
            boost::archive::binary_oarchive oa(oss);
            oa << ri;
        }
        return lz4_pack_string(oss.str());
    }
    static std::shared_ptr<RestartNodeInfo> unpack_rinfo(const std::string& blob) {
        auto ri = std::make_shared<RestartNodeInfo>();
        const std::string raw = lz4_unpack_string(blob);
        std::istringstream iss(raw, std::ios::binary);
        boost::archive::binary_iarchive ia(iss);
        ia >> *ri;
        return ri;
    }
 
    void save(const std::string& path) {
        // Pack each UNIQUE crew snapshot once, by pointer. The raw maps stay
        // populated for the caller, but the archive carries only the packed
        // blobs plus index maps (serialize() skips the raw maps).
        {
            packed_rinfos.clear();
            restart_state_pack_ref.clear();
            task_state_pack_ref.clear();
            std::map<const RestartNodeInfo*, int> idx;
            auto pack_of = [&](const std::shared_ptr<RestartNodeInfo>& ri) -> int {
                if (!ri) return -1;
                auto it = idx.find(ri.get());
                if (it != idx.end()) return it->second;
                int i = static_cast<int>(packed_rinfos.size());
                packed_rinfos.push_back(pack_rinfo(*ri));
                idx.emplace(ri.get(), i);
                return i;
            };
            for (auto& [cid, ri] : restart_states)
                restart_state_pack_ref[cid] = pack_of(ri);
            task_state_pack_ref.reserve(task_restart_states.size());
            for (auto& ri : task_restart_states)
                task_state_pack_ref.push_back(pack_of(ri));
        }
        // Compress blob payloads in place before archiving. Callers build the
        // object fresh for each save and discard it afterwards, so mutating
        // it here is safe. Queue-compressed sources are left to that layer.
        if (!blobs_compressed) {
            if (!compress_queue) {
                for (auto& [key, blob] : source_blobs) {
                    if (blob.empty()) continue;
                    auto c = lz4_compress(blob.data(), static_cast<int>(blob.size()));
                    blob = std::move(c);
                }
            }
            for (auto& [key, blob] : branch_record_blobs)
                if (!blob.empty()) blob = lz4_pack_string(blob);
            if (!branch_center_blob.empty())
                branch_center_blob = lz4_pack_string(branch_center_blob);
            blobs_compressed = true;
        }
        // Write to a temp file and rename into place, so a kill mid-save never
        // leaves a truncated checkpoint at the canonical path.
        const std::string tmp_path = path + ".tmp";
        {
            std::ofstream ofs(tmp_path, std::ios::binary);
            if (!ofs) {
                std::cerr << "ERROR: cannot open checkpoint file for writing: " << tmp_path << std::endl;
                return;
            }
            boost::archive::binary_oarchive oa(ofs);
            oa << *this;
            ofs.flush();
            if (!ofs.good()) {
                std::cerr << "ERROR: checkpoint write failed (disk full?): " << tmp_path << std::endl;
                ofs.close();
                std::error_code ec;
                std::filesystem::remove(tmp_path, ec);
                return;
            }
        }
        std::error_code ec;
        std::filesystem::rename(tmp_path, path, ec);
        if (ec) {
            std::cerr << "ERROR: cannot move checkpoint into place: " << ec.message() << std::endl;
            return;
        }
        std::cout << "[CHECKPOINT] Saved: " << path
                  << " (" << sequences_completed << " sequences, "
                  << pending_tasks.size() << " pending)" << std::endl;
    }
 
    static DETCheckpoint load(const std::string& path) {
        DETCheckpoint ckpt;
        std::ifstream ifs(path, std::ios::binary);
        if (!ifs) {
            // A default-constructed checkpoint would make a mistyped --resume
            // path finish at once with exit 0, so exit with an error instead.
            std::cerr << "FATAL: cannot open checkpoint file for reading: " << path << std::endl;
            std::exit(1);
        }
        try {
            boost::archive::binary_iarchive ia(ifs);
            ia >> ckpt;
        } catch (const std::exception& ex) {
            std::cerr << "FATAL: checkpoint file is corrupt or truncated: " << path
                      << " (" << ex.what() << ")" << std::endl;
            std::exit(1);
        }
        const int stored_version = ckpt.version;
        if (!version_supported(stored_version)) {
            std::cerr << "FATAL: checkpoint version " << stored_version
                      << " is older than the format this binary reads (version "
                      << VERSION << ") and cannot\n"
                      << "  be read by this binary. Finish that run with the binary that\n"
                      << "  wrote it, or start a fresh run." << std::endl;
            std::exit(1);
        }
        // Rebuild the rinfo maps from the packed uniques. The indexes restore
        // the shared_ptr aliasing between restart_states and
        // task_restart_states. The guard tests the STORED version and is
        // always true here, because version_supported() has already passed.
        if (stored_version >= 7) {
            std::vector<std::shared_ptr<RestartNodeInfo>> uniques;
            uniques.reserve(ckpt.packed_rinfos.size());
            for (auto& blob : ckpt.packed_rinfos)
                uniques.push_back(unpack_rinfo(blob));
            for (auto& [cid, ref] : ckpt.restart_state_pack_ref)
                ckpt.restart_states[cid] =
                    (ref >= 0 && ref < (int)uniques.size()) ? uniques[ref] : nullptr;
            ckpt.task_restart_states.clear();
            ckpt.task_restart_states.reserve(ckpt.task_state_pack_ref.size());
            for (int ref : ckpt.task_state_pack_ref)
                ckpt.task_restart_states.push_back(
                    (ref >= 0 && ref < (int)uniques.size()) ? uniques[ref] : nullptr);
            ckpt.packed_rinfos.clear();
            ckpt.packed_rinfos.shrink_to_fit();
            ckpt.restart_state_pack_ref.clear();
            ckpt.task_state_pack_ref.clear();
        }
        // Strip the at-rest compression layer so every consumer sees the
        // uncompressed payloads. Queue-layer compression stays intact: it is
        // owned by compress_queue and stripped at restore time.
        if (ckpt.blobs_compressed) {
            if (!ckpt.compress_queue) {
                for (auto& [key, blob] : ckpt.source_blobs) {
                    if (blob.empty()) continue;
                    auto d = lz4_decompress(blob.data(), static_cast<int>(blob.size()));
                    blob = std::move(d);
                }
            }
            for (auto& [key, blob] : ckpt.branch_record_blobs)
                if (!blob.empty()) blob = lz4_unpack_string(blob);
            if (!ckpt.branch_center_blob.empty())
                ckpt.branch_center_blob = lz4_unpack_string(ckpt.branch_center_blob);
            ckpt.blobs_compressed = false;
        }
        return ckpt;
    }
};