~siborgium/prelockdpp

9540c80a2d99d3902513d555015d41a71f4022dd — Sergey Smirnykh 2 years ago d953dcf
Use Boost.Intrusive & plf::colony to use as few allocations as possible
2 files changed, 184 insertions(+), 42 deletions(-)

M CMakeLists.txt
M main.cpp
M CMakeLists.txt => CMakeLists.txt +11 -1
@@ 2,10 2,20 @@ cmake_minimum_required(VERSION 3.0)

project(prelockdpp LANGUAGES CXX)

include(FetchContent)
FetchContent_Declare(
    plf_colony
    GIT_REPOSITORY https://github.com/mattreecebentley/plf_colony.git
    GIT_TAG        7243b4601b5d7e989550c158a9918ea5a05feaf4
)
FetchContent_MakeAvailable(plf_colony)

find_package(Boost 1.70 REQUIRED)
find_package(nlohmann_json 3.10.5 REQUIRED)

add_executable(prelockdpp main.cpp)
target_compile_features(prelockdpp PRIVATE cxx_std_20)
target_link_libraries(prelockdpp PRIVATE nlohmann_json::nlohmann_json)
target_link_libraries(prelockdpp PRIVATE Boost::boost nlohmann_json::nlohmann_json)
target_include_directories(prelockdpp PRIVATE "${plf_colony_SOURCE_DIR}")

install(TARGETS prelockdpp RUNTIME DESTINATION bin)

M main.cpp => main.cpp +173 -41
@@ 1,15 1,20 @@

#include <iostream>
#include <fstream>
#include <limits>
#include <charconv>
#include <string>
#include <string_view>
#include <vector>
#include <unordered_set>
#include <vector>

#include <cassert>
#include <cstring>

#include <boost/intrusive/set.hpp>
#include <boost/intrusive/set_hook.hpp>
#include <plf_colony.h>

#include <dirent.h>
#include <fcntl.h>
#include <sys/epoll.h>


@@ 63,12 68,6 @@ struct errno_exception: public std::runtime_error {
        // intentionally left blank
    }
};
struct file_not_found_error: public std::runtime_error {
    using std::runtime_error::runtime_error;
};
struct process_lookup_error: public std::runtime_error {
    using std::runtime_error::runtime_error;
};

constexpr auto tspec_diff(const timespec& t0, const timespec& t1) {
    timespec t {


@@ 147,17 146,47 @@ struct timer_fd: public base_fd {
    }
};

struct uniq_id {
namespace intr = boost::intrusive;

namespace tags {
    struct alive{};
    struct dead{};
};

using uniq_alive_hook =
    intr::set_base_hook<intr::optimize_size<true>, intr::tag<tags::alive>>;
using uniq_temp_hook =
    intr::set_base_hook<intr::optimize_size<true>, intr::tag<tags::dead>>;

struct uniq_id: public uniq_alive_hook, public uniq_temp_hook {
    timespec start_time;
    int      pid;

    uniq_id(timespec time, int pid): start_time{ time }, pid{ pid } {}
    constexpr bool operator == (const uniq_id& id) const {
        return start_time.tv_sec == id.start_time.tv_sec
            && start_time.tv_nsec == id.start_time.tv_nsec
            && pid == id.pid;
    }
    constexpr bool operator < (const uniq_id& id) const {
        if (start_time.tv_sec != id.start_time.tv_sec) {
            return start_time.tv_sec < id.start_time.tv_sec;
        }
        if (start_time.tv_nsec != id.start_time.tv_nsec) {
            return start_time.tv_nsec < id.start_time.tv_nsec;
        }
        return pid < id.pid;
    }
};

using uniq_alive_set = intr::set<
    uniq_id, intr::compare<std::less<uniq_id>>, intr::base_hook<uniq_alive_hook>
>;

using uniq_temp_set = intr::set<
    uniq_id, intr::compare<std::less<uniq_id>>, intr::base_hook<uniq_temp_hook>
>;

namespace std {
    template<>
    struct hash<uniq_id> {


@@ 173,6 202,20 @@ namespace std {
    };
}

struct linebuf {
    char*  ptr{ nullptr };
    size_t capacity{ 0 };

    ssize_t getline(FILE* file) {
        return ::getline(&ptr, &capacity, file);
    }
    ~linebuf() {
        if (ptr) {
            free(ptr);
        }
    }
};

struct context {
    timespec start_time;
    int      self_pid;


@@ 200,13 243,21 @@ struct context {
    char*       pid_end;

    // buffers used to avoid re-allocating the same space
    linebuf                         line_buffer;
    std::string                     stat_buffer;
    std::string                     cgroup_buffer;

    // names (comms, actually, see comm in proc(5), /proc/[pid]/stat) of
    // programs we should lock
    std::unordered_set<std::string> name_set;

    // storage for ids
    plf::colony<uniq_id> ids;
    uniq_alive_set       alive_ids;

    // FIXME: I could use Boost.Intrusive & build something decent
    //        instead of endless maps & sets
    // uniq_id -> mappings list

    void mlockall() {
        if (::mlockall(MCL_FUTURE)) {
            throw std::runtime_error{ concat("mlockall failed ", errno) };


@@ 306,19 357,24 @@ struct context {
                // FIXME: fsck unordered_set
                //        imagine not being able to lookup strings by string_view key
                //        in 2022
                bool lock_ok = name_set.contains(std::string(comm));
                bool lock_ok = std::find_if(name_set.begin(), name_set.end(), [=](auto&& c){
                    return comm == c;
                }) != name_set.end();

                if (!lock_ok && cfg.check_cgroup) {
                    fmt_proc_pid_file(cgroup_path);

                    std::ifstream stream{ path_buf };
                    auto file = fopen(path_buf, "r");
                    if (!file) {
                        throw errno_exception("fopen failed", errno);
                    }
                    DEFER{ fclose(file); };

                    std::size_t index{ 0 };
                    while (std::getline(stream, cgroup_buffer)) {
                    ssize_t n;
                    while ((n = line_buffer.getline(file)) != -1) {
                        if (index == cgroup_v2_index) {
                            std::string_view cgroup{
                                cgroup_buffer.c_str() + 3,
                                cgroup_buffer.c_str() + cgroup_buffer.size()
                            };
                            std::cout << '\n';
                            std::string_view cgroup( line_buffer.ptr, n );
                            lock_ok = match_cgroup(cgroup);
                            break;
                        }


@@ 339,7 395,21 @@ struct context {
                return stat_view.substr(begin, end - begin);
            }();

            callback(start_time);
            auto start_time_ = [=,this] {
                unsigned long long ticks;
                auto [p, ec] = std::from_chars(
                    start_time.begin(),
                    start_time.end(),
                    ticks
                );
                if (ec != std::errc()) {
                    throw std::logic_error{ "Failed to parse start_time value" };
                }
                return tspec_from_ticks(ticks);
            }();

            callback(start_time_);

        } catch (const std::system_error& e) {
            auto code = e.code();
            if (code.value() != ENOENT) {


@@ 347,10 417,32 @@ struct context {
            }
        }
    }
    auto add_mappings(auto & vec) {
        try {
            fmt_proc_pid_file(maps_path);
            auto file = fopen(path_buf, "r");
            if (!file) {
                throw errno_exception("fopen(/proc/PID/maps) failed", errno);
            }
            ssize_t n = 0;
            while ((n = line_buffer.getline(file)) != -1) {
                auto begin = std::find(line_buffer.ptr, line_buffer.ptr + n, '/');
                if (begin != line_buffer.ptr + n) {
                    vec.emplace_back(begin, line_buffer.ptr + n - begin);
                }
            }
        } catch (const std::system_error& e) {
            if (e.code().value() != ENOENT) {
                throw;
            }
        }
    }
    auto get_current_set() {
        using namespace std::string_view_literals;

        std::unordered_set<uniq_id> uniq_set;
        // this set will contain ids (see later)
        // that are alive at this iteration
        uniq_temp_set current_ids;

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


@@ 372,7 464,7 @@ struct context {
            auto name_end = entry->d_name + strlen(entry->d_name);
            int pid;
            auto [p, ec] = std::from_chars(entry->d_name, name_end, pid);
            if (ec != std::errc() || p != name_end) {
            if (ec != std::errc() || p != name_end || pid == self_pid) {
                continue;
            }



@@ 380,37 472,73 @@ struct context {
            fmt_proc_pid_file(stat_path);

            with_uniq_id(path_buf, [&](auto start_time) {
                auto start_time_ = [=,this] {
                    unsigned long long ticks;
                    auto [p, ec] = std::from_chars(
                        start_time.begin(),
                        start_time.end(),
                        ticks
                    );
                    if (ec != std::errc()) {
                        throw std::logic_error{ "Failed to parse start_time value" };
                    }
                    return tspec_from_ticks(ticks);
                }();

                auto [iter, inserted] = uniq_set.emplace(uniq_id{ start_time_, pid });
                // if uniq_id is already exists
                // (e.g. process was alive on the prev iteration)
                // we migrate it to the `current_ids` set
                uniq_id id{ start_time, pid };
                if (auto iter = alive_ids.find(id); iter != alive_ids.end()) {
                    current_ids.insert(*iter);
                    alive_ids.erase(iter);
                } else {
                    current_ids.insert(*ids.emplace(id));
                }

#if 0

                auto [iter, inserted] = uniq_set.emplace(uniq_id{ start_time, pid });
                assert(inserted);


                if (!uniq_mappings.contains(*iter)) {
                    add_mappings(uniq_mappings[*iter]);
                }
#endif
            });
        }
        if (errno) {
            throw errno_exception("readdir(proc) failed", errno);
        }

        // all alive ids are in `current_ids`
        // and alive_ids only store dead ones
        // lets clear alive_ids and re-populate it with actually alive ids
        // FIXME: figure out how to store iterators
        //        and enjoy cheap removal
        {
            auto maybe_dead = ids.begin();
            while (maybe_dead != ids.end()) {
                if (auto dead = alive_ids.find(*maybe_dead); dead != alive_ids.end()) {
                    alive_ids.erase(dead);
                    auto iter = maybe_dead;
                    ids.erase(iter);
                    ++maybe_dead;
                    continue;
                }
                ++maybe_dead;
            }
        }
        for (auto&& id: current_ids) {
            alive_ids.insert(id);
        }

        for (auto && id: alive_ids) {
            std::cout << id.start_time.tv_sec << '.' << id.start_time.tv_nsec << ' ' << id.pid << '\n';
        }

#if 0
        for (auto && [k, v]: uniq_mappings) {
            for (auto && m: v) {
                std::cout << m << '\n';
            }
        }
#endif
        auto m = tspec_diff(m0, get_clock_time(CLOCK_MONOTONIC));
        auto p = tspec_diff(p0, get_clock_time(CLOCK_PROCESS_CPUTIME_ID));

        std::cout << "Uptime:         " << uptime.tv_sec << '.' << uptime.tv_nsec << '\n';
        std::cout << "Monotonic time: " << m.tv_sec << '.' << m.tv_nsec << '\n';
        std::cout << "Proc CPU  time: " << p.tv_sec << '.' << p.tv_nsec << '\n';

        for (auto && id: uniq_set) {
            std::cout << id.start_time.tv_sec << '.' << id.start_time.tv_nsec << ' ' << id.pid << '\n';
        }
    }

    context(config& cfg):


@@ 422,10 550,14 @@ struct context {

        // get cgroup v2 index
        {
            std::ifstream file{ "/proc/self/cgroup" };
            std::string buf;
            while (std::getline(file, buf)) {
                std::string_view line{ buf };
            auto file = fopen("/proc/self/cgroup", "r");
            if (!file) {
                throw errno_exception("fopen(/proc/self/cgroup) failed", errno);
            }
            DEFER{ fclose(file); };
            ssize_t n;
            while ((n = line_buffer.getline(file)) != -1) {
                std::string_view line{ line_buffer.ptr, size_t(n) };
                if (line.starts_with("0::")) {
                    break;
                }