~siborgium/prelockdpp

f4f2867100fe710c09feacb6a6be8fdf98536d62 — Sergey Smirnykh 1 year, 11 months ago c2ca3b6
Files rotation
1 files changed, 116 insertions(+), 52 deletions(-)

M main.cpp
M main.cpp => main.cpp +116 -52
@@ 234,40 234,48 @@ namespace std {
    };
}

// each `mapping` holds filename it points to
// `mapping` is to be allocated via `mapping::create`, see it's description
struct mapping: intr::set_base_hook<intr::optimize_size<true>> {
    std::size_t      refs;
// each `mapped_file` holds filename it points to
// `mapped_file` is to be allocated via `mapped_file::create`, see it's description
struct mapped_file: intr::set_base_hook<intr::optimize_size<true>> {
    using state_t = int8_t;

    std::string_view filename;
    state_t          state;

    bool operator < (const mapping& m) const { return filename < m.filename; }
    bool operator == (const mapping& m) const { return filename == m.filename; }
    bool operator < (const mapped_file& m) const { return filename < m.filename; }
    bool operator == (const mapped_file& m) const { return filename == m.filename; }

    // allocates mapping and it's filename as contigious chunk of memory
    // so you can free both mapping and filename at once
    static mapping* create(std::string_view filename) {
        auto size = sizeof(mapping) + filename.size();
        auto* ptr = new (std::align_val_t{ alignof(mapping) }) char[size];
    // creates mapping with state = lock
    static mapped_file* create(std::string_view filename) {
        auto size = sizeof(mapped_file) + filename.size();
        auto* ptr = new (std::align_val_t{ alignof(mapped_file) }) char[size];

        auto* begin = ptr + sizeof(mapping);
        auto* begin = ptr + sizeof(mapped_file);
        auto* end = begin + filename.size();
        assert(end - begin == filename.size());
        auto* end1 = std::copy(filename.begin(), filename.end(), begin);
        assert(end1 == end);

        auto* map = new (ptr) mapping{ .refs = 0, .filename = { begin, end } };
        auto* map = new (ptr) mapped_file{ .filename = { begin, end }, .state = 0 };
        assert((void*)map == (void*)ptr);
        return map;
    }
    static void operator delete (void *p) { return delete[] static_cast<char*>(p); }
};

struct mapping_sv_comp {
    constexpr bool operator ()(const mapping& m, std::string_view view) const { return m.filename < view; }
    constexpr bool operator ()(std::string_view view, const mapping& m) const { return view < m.filename; }
struct mapped_file_sv_comp {
    constexpr bool operator ()(const mapped_file& m, std::string_view view) const { return m.filename < view; }
    constexpr bool operator ()(std::string_view view, const mapped_file& m) const { return view < m.filename; }
};

using mapping_set = intr::set<mapping, intr::compare<std::less<mapping>>>;
using mapped_file_set = intr::set<mapped_file, intr::compare<std::less<mapped_file>>>;

struct process_maps: intr::set_base_hook<intr::optimize_multikey<true>> {
    int fd;
    std::vector<mapped_file*> mapped_files;
};

struct linebuf {
    char*  ptr{ nullptr };


@@ 287,6 295,46 @@ struct linebuf {
    }
};

struct mapping_set {
    mapped_file_set maps;

    void add(std::string_view path) {
        mapped_file_set::insert_commit_data commit_data;
        auto [iter, can_insert] = maps.insert_check(path, mapped_file_sv_comp{}, commit_data);
        if (can_insert) {
            // TODO: maybe use pmr::monotonic_buffer_resource to store mapped_files
            iter = maps.insert_commit(*mapped_file::create(path), commit_data);
        }
        iter->state = can_insert;
    }

    void rotate() {
        mapped_file_set to_unlock;
        DEFER{ to_unlock.clear_and_dispose([](auto && x){ delete x; }); };
        for (auto begin = maps.begin(); begin != maps.end();) {
            auto& m = *begin;
            if (m.state < 0) {
                ++begin;
                maps.erase(maps.iterator_to(m));
                to_unlock.push_front(m);
            } else {
                --m.state;
                if (m.state >= 0) {
                    fmt::print("locking '{}'\n", m.filename);
                }
                ++begin;
            }
        }
        for (auto&& t: to_unlock) {
            fmt::print("unlocking '{}'\n", t.filename);
        }
    }

    ~mapping_set() {
        maps.clear_and_dispose([](auto && m){ delete m; });
    }
};

struct context {
    timespec start_time;
    int      self_pid;


@@ 406,7 454,7 @@ struct context {
#if 0
        std::cout << "match_cgroup(" << cgroup << "): unimplemented\n";
#endif
        return true;
        return false;
    }
    auto with_uniq_id(const char* stat_path, auto && callback) {
        using namespace std::string_view_literals;


@@ 506,12 554,9 @@ struct context {
                auto end = line_buffer.ptr + n;
                if (auto begin = std::find(line_buffer.ptr, end, '/'); begin != end) {
                    std::string_view path{ begin, end };
                    path.remove_suffix(path.ends_with('\n'));

                    mapping_set::insert_commit_data commit_data;
                    auto [_, can_insert] = set.insert_check( path, mapping_sv_comp{}, commit_data);
                    if (can_insert) {
                        set.insert_commit(*mapping::create(path), commit_data);
                    }
                    set.add(path);
                }
            }
        } catch (const errno_exception& e) {


@@ 520,7 565,8 @@ struct context {
            }
        }
    }
    auto get_current_set() {
    /// @brief Polls procfs, updating stored process & mappings info
    auto& acquire_current_set() {
        using namespace std::string_view_literals;

        // this set will contain ids (see later)


@@ 531,11 577,6 @@ struct context {
        //       when inserting to `current_uniqs_set`, just = move(new_ids)
        //       plus any_hooks are pain
        uniq_alive_set new_ids;
        mapping_set    uniq_mappings;

        auto m0 = get_clock_time(CLOCK_MONOTONIC);
        auto p0 = get_clock_time(CLOCK_PROCESS_CPUTIME_ID);
        auto uptime = get_uptime();

        DIR* proc_dir = opendir("/proc");
        if (!proc_dir) {


@@ 586,7 627,7 @@ struct context {
                    }
                }

                add_mappings(uniq_mappings);
                add_mappings(mappings);
            });
        }
        if (errno) {


@@ 596,30 637,16 @@ struct context {
        // all alive ids are in `new_ids`
        // and `current_uniqs_set` only stores dead ones
        // lets clear dead ids and re-populate it with actually alive ids
        {
            while (!current_uniqs_set.empty()) {
                auto* dead = current_uniqs_set.unlink_leftmost_without_rebalance();
                auto orig = get_self_iter(*dead);
                ids.erase(orig);
            }
            current_uniqs_set = std::move(new_ids);
        }

        for (auto && id: current_uniqs_set) {
            auto lifetime = tspec_diff(id.start_time, uptime);
            fmt::print("{}.{} {} ({}.{})\n", id.start_time.tv_sec, id.start_time.tv_nsec, id.pid, lifetime.tv_sec, lifetime.tv_nsec);
        }

        for (auto && k: uniq_mappings) {
            fmt::print("Mapped file: {}\n", k.filename);
        }

        auto m = tspec_diff(m0, get_clock_time(CLOCK_MONOTONIC));
        auto p = tspec_diff(p0, get_clock_time(CLOCK_PROCESS_CPUTIME_ID));
        current_uniqs_set.clear_and_dispose([this](auto && dead) {
            auto orig = get_self_iter(*dead);
            ids.erase(orig);
        });
        current_uniqs_set = std::move(new_ids);
        return current_uniqs_set;
    }

        fmt::print("Uptime:         {}.{}\n", uptime.tv_sec, uptime.tv_nsec);
        fmt::print("Monotonic time: {}.{}\n", m.tv_sec, m.tv_nsec);
        fmt::print("Proc CPU time:  {}.{}\n", p.tv_sec, p.tv_nsec);
    void add_critical_name(std::string_view name) {
        name_set.emplace(name.data(), name.size());
    }

    context(config& cfg):


@@ 684,7 711,19 @@ struct lock_context {
        if (read(timer.fd, &time, sizeof(time)) != sizeof(time)) {
            throw std::logic_error{ "Reading from timer fd returned fewer bytes than expected" };
        }
        ctx.get_current_set();
        auto m0 = ctx.get_clock_time(CLOCK_MONOTONIC);
        auto p0 = ctx.get_clock_time(CLOCK_PROCESS_CPUTIME_ID);
        auto uptime = ctx.get_uptime();

        ctx.acquire_current_set();
        ctx.mappings.rotate();

        auto m = tspec_diff(m0, ctx.get_clock_time(CLOCK_MONOTONIC));
        auto p = tspec_diff(p0, ctx.get_clock_time(CLOCK_PROCESS_CPUTIME_ID));

        fmt::print("Uptime:         {}.{}\n", uptime.tv_sec, uptime.tv_nsec);
        fmt::print("Monotonic time: {}.{}\n", m.tv_sec, m.tv_nsec);
        fmt::print("Proc CPU time:  {}.{}\n", p.tv_sec, p.tv_nsec);
    }
    void on_signal() {
        signalfd_siginfo info;


@@ 752,9 791,34 @@ int main(int argc, char **argv) {
        }
#endif

#if 0
        auto arg = [begin = argv + 1, end = argv + argc](std::string_view name) {
            auto iter = std::find_if(begin, end, [=](auto* a){ return name == a; });
            if (iter == end) {
                return std::string_view{};
            }
            ++iter;
            if (iter == end) {
                throw std::invalid_argument{ "No argument value specified" };
            }
            return std::string_view{ *iter };
        };

        auto conf_path = arg("-c");
        if (conf_path.empty()) {
            conf_path = "/usr/share/prelockdpp/prelockd.conf";
        }
#endif

        config  cfg;
        context ctx{ cfg };

        for (int i = 1; i < argc; ++i) {
            ctx.add_critical_name(argv[i]);
        }

        fmt::print("Name set: {}\n", ctx.name_set);

        ctx.mlockall();

        lock_context runner{ ctx };