file_cache.cpp

/**
 * \file lite/src/mge/algo_cache/file_cache.cpp
 * MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 *
 * Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 */

#include "lite_build_config.h"

#if LITE_BUILD_WITH_MGE
#include "../common.h"
#include "file_cache.h"

using namespace lite;

//////////////////////// InFilePersistentCache::InputMemory ///////////////
class InFilePersistentCache::InputMemory {
    const uint8_t* m_ptr;
    size_t m_offset = 0;
    size_t m_size;

public:
    InputMemory(const uint8_t* bin, size_t size) : m_ptr{bin}, m_size{size} {}

    template <typename T>
    void read(T& val) {
        static_assert(
                std::is_trivially_copyable<T>::value,
                "only support trivially copyable type");
        LITE_ASSERT(m_offset + sizeof(T) <= m_size);
        memcpy(&val, m_ptr, sizeof(T));
        m_offset += sizeof(T);
        m_ptr += sizeof(T);
    }

    template <typename T>
    void read(T* buf, size_t size) {
        static_assert(
                std::is_trivially_copyable<T>::value && sizeof(T) == 1,
                "only support read bytes");
        LITE_ASSERT(m_offset + size <= m_size);
        memcpy(buf, m_ptr, size);
        m_offset += size;
        m_ptr += size;
    }
};

//////////////////////// InFilePersistentCache::InputFile ///////////////
class InFilePersistentCache::InputFile {
    FILE* m_fp;

public:
    InputFile(const char* path) : m_fp{fopen(path, "rb")} {
        LITE_ASSERT(m_fp, "failed to open %s: %s", path, strerror(errno));
    }
    ~InputFile() {
        if (m_fp) {
            fclose(m_fp);
        }
    }

    template <typename T>
    void read(T& val) {
        static_assert(
                std::is_trivially_copyable<T>::value,
                "only support trivially copyable type");
        auto ret = fread(&val, sizeof(T), 1, m_fp);
        LITE_ASSERT(ret == 1);
    }

    template <typename T>
    void read(T* buf, size_t size) {
        static_assert(
                std::is_trivially_copyable<T>::value && sizeof(T) == 1,
                "only support read bytes");
        auto ret = fread(buf, size, 1, m_fp);
        LITE_ASSERT(ret == 1);
    }
};

//////////////////////// InFilePersistentCache::OutputFile ///////////////
class InFilePersistentCache::OutputFile {
    FILE* m_fp;

public:
    OutputFile(const char* path) : m_fp{fopen(path, "wb")} {
        LITE_ASSERT(m_fp, "failed to open %s: %s", path, strerror(errno));
    }
    ~OutputFile() {
        if (m_fp) {
            fclose(m_fp);
        }
    }

    template <typename T>
    void write(T val) {
        auto ret = fwrite(&val, sizeof(T), 1, m_fp);
        LITE_ASSERT(ret == 1);
    }

    template <typename T>
    void write(const T* buf, size_t size) {
        static_assert(sizeof(T) == 1, "only support write bytes");
        auto ret = fwrite(buf, size, 1, m_fp);
        LITE_ASSERT(ret == 1);
    }

    void flush() { fflush(m_fp); }

    void set_head() { fseek(m_fp, 0, SEEK_SET); }
};

//////////////////////// InFilePersistentCache::BlobStorage ///////////////

template <typename Input>
InFilePersistentCache::BlobStorage& InFilePersistentCache::BlobStorage::init_from_input(
        Input& inp) {
    uint32_t data_size;
    inp.read(data_size);
    size = data_size;
    data_refhold = std::make_unique<uint8_t[]>(size);
    inp.read(data_refhold.get(), size);
    ptr = data_refhold.get();
    return *this;
}

void InFilePersistentCache::BlobStorage::write_to_file(OutputFile& out_file) const {
    uint32_t u_size = size;
    out_file.write(u_size);
    out_file.write(data_refhold.get(), u_size);
}

InFilePersistentCache::BlobStorage& InFilePersistentCache::BlobStorage::init_data_ref(
        const Blob& b) {
    data_refhold = std::make_unique<uint8_t[]>(b.size + 1);
    memcpy(data_refhold.get(), b.ptr, b.size);
    data_refhold.get()[b.size] = 0;  // for C-string safety
    ptr = data_refhold.get();
    size = b.size;
    return *this;
}

//////////////////////// InFilePersistentCache //////////////////////

template <typename Input>
void InFilePersistentCache::read_cache(Input& inp) {
    uint32_t nr_category;
    inp.read(nr_category);
    char category_buf[256];
    for (uint32_t i = 0; i < nr_category; i++) {
        uint32_t category_size;
        inp.read(category_size);
        inp.read(category_buf, category_size);
        category_buf[category_size] = '\0';

        std::string category(category_buf);
        mgb_log_debug("load new category: %s", category_buf);

        // read bobs
        uint32_t nr_bobs;
        inp.read(nr_bobs);
        for (uint32_t j = 0; j < nr_bobs; j++) {
            BlobStorage key_storage;
            key_storage.init_from_input(inp).init_hash();
            mgb_log_debug("read key: %zu", key_storage.hash);
            m_cache[category][std::move(key_storage)].init_from_input(inp);
        }
    }
}

InFilePersistentCache::InFilePersistentCache(const char* path, bool always_open) {
    if (!access(path, F_OK)) {
        mgb_log_debug("use fastrun cache: %s", path);
        InputFile inp(path);
        read_cache<InputFile>(inp);
    }
    if (always_open) {
        m_always_open_file = std::make_shared<OutputFile>(path);
    }
}

InFilePersistentCache::InFilePersistentCache(const uint8_t* bin, size_t size) {
    LITE_ASSERT(bin);
    InputMemory inp(bin, size);
    read_cache<InputMemory>(inp);
}

void InFilePersistentCache::dump_cache(const char* path) {
    OutputFile out_file(path);
    dump_cache(&out_file);
}

void InFilePersistentCache::dump_cache(OutputFile* out_file) {
    uint32_t nr_category = m_cache.size();
    out_file->write(nr_category);

    for (const auto& cached_category : m_cache) {
        uint32_t category_size = cached_category.first.size();
        out_file->write(category_size);
        out_file->write(cached_category.first.data(), category_size);
        mgb_log_debug("write new category: %s", cached_category.first.c_str());

        uint32_t nr_bobs = cached_category.second.size();
        out_file->write(nr_bobs);
        for (const auto& item : cached_category.second) {
            mgb_log_debug("dump key: %zu", item.first.hash);
            item.first.write_to_file(*out_file);
            item.second.write_to_file(*out_file);
        }
    }
}

mgb::Maybe<InFilePersistentCache::Blob> InFilePersistentCache::get(
        const std::string& category, const Blob& key) {
    decltype(m_cache.begin()) iter0;
    {
        MGB_LOCK_GUARD(m_mtx);
        iter0 = m_cache.find(category);
        if (iter0 == m_cache.end())
            return mgb::None;
    }

    BlobStorage key_storage;
    key_storage.Blob::operator=(key);
    key_storage.init_hash();

    MGB_LOCK_GUARD(m_mtx);

    auto iter1 = iter0->second.find(key_storage);
    if (iter1 == iter0->second.end())
        return mgb::None;
    return iter1->second;
}

void InFilePersistentCache::put(
        const std::string& category, const Blob& key, const Blob& value) {
    BlobStorage key_storage;
    key_storage.init_data_ref(key).init_hash();

    MGB_LOCK_GUARD(m_mtx);
    auto size0 = m_cache.size();
    m_cache[category][std::move(key_storage)].init_data_ref(value);
    if (m_cache.size() > size0) {
        mgb_log_debug("new cache category: %s", category.c_str());
    }
    if (m_always_open_file) {
        m_always_open_file->set_head();
        dump_cache(m_always_open_file.get());
        m_always_open_file->flush();
    }
}
#endif

// vim: syntax=cpp.doxygen foldmethod=marker foldmarker=f{{{,f}}}