be/src/agent/cgroup_cpu_ctl.h

// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

#pragma once

#include <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

#include <set>
#include <shared_mutex>

#include "common/config.h"
#include "common/status.h"
#include "util/cpu_info.h"

namespace doris {

class CgroupCpuCtl {
public:
    virtual ~CgroupCpuCtl() = default;
    CgroupCpuCtl(uint64_t wg_id) { _wg_id = wg_id; }

    virtual Status init() = 0;

    virtual Status add_thread_to_cgroup() = 0;

    void update_cpu_hard_limit(int cpu_hard_limit);

    void update_cpu_soft_limit(int cpu_shares);

    // for log
    void get_cgroup_cpu_info(uint64_t* cpu_shares, int* cpu_hard_limit);

    static void init_doris_cgroup_path();

    static Status delete_unused_cgroup_path(std::set<uint64_t>& used_wg_ids);

    static std::shared_ptr<CgroupCpuCtl> create_cgroup_cpu_ctl(uint64_t wg_id);

    static bool is_a_valid_cgroup_path(std::string cg_path);

    static uint64_t cpu_soft_limit_default_value();

protected:
    virtual Status modify_cg_cpu_hard_limit_no_lock(int cpu_hard_limit) = 0;

    virtual Status modify_cg_cpu_soft_limit_no_lock(int cpu_shares) = 0;

    Status add_thread_to_cgroup(std::string task_file);

    static Status write_cg_sys_file(std::string file_path, std::string value, std::string msg,
                                    bool is_append);

    static Status init_cgroup_v2_query_path_public_file(std::string home_path,
                                                        std::string query_path);

protected:
    inline static uint64_t _cpu_core_num;
    // TODO (yiguolei): Should echo this period to cpu.cfs_period_us???
    inline const static uint64_t _cpu_cfs_period_us = 100000;
    inline static std::string _doris_cgroup_cpu_path = "";
    inline static std::string _doris_cgroup_cpu_query_path = "";
    inline static bool _is_enable_cgroup_v1_in_env = false;
    inline static bool _is_enable_cgroup_v2_in_env = false;
    inline static bool _is_cgroup_query_path_valid = false;

    // cgroup v2 public file
    inline static std::string _doris_cgroup_cpu_path_subtree_ctl_file = "";
    inline static std::string _cgroup_v2_query_path_subtree_ctl_file = "";
    inline static std::string _doris_cg_v2_procs_file = "";

protected:
    int _cpu_hard_limit = 0;
    std::shared_mutex _lock_mutex;
    bool _init_succ = false;
    uint64_t _wg_id = -1; // workload group id
    uint64_t _cpu_shares = 0;
};

/*
    NOTE: directory structure
    1 sys cgroup root path:
        /sys/fs/cgroup
    
    2 sys cgroup cpu controller path:
        /sys/fs/cgroup/cpu
    
    3 doris home path:
        /sys/fs/cgroup/cpu/{doris_home}/
    
    4 doris query path
        /sys/fs/cgroup/cpu/{doris_home}/query
    
    5 workload group path
        /sys/fs/cgroup/cpu/{doris_home}/query/{workload group id}
    
    6 workload group quota file:
        /sys/fs/cgroup/cpu/{doris_home}/query/{workload group id}/cpu.cfs_quota_us
    
    7 workload group tasks file:
        /sys/fs/cgroup/cpu/{doris_home}/query/{workload group id}/tasks
    
    8 workload group cpu.shares file:
    /sys/fs/cgroup/cpu/{doris_home}/query/{workload group id}/cpu.shares
*/
class CgroupV1CpuCtl : public CgroupCpuCtl {
public:
    CgroupV1CpuCtl(uint64_t tg_id) : CgroupCpuCtl(tg_id) {}
    Status init() override;
    Status modify_cg_cpu_hard_limit_no_lock(int cpu_hard_limit) override;
    Status modify_cg_cpu_soft_limit_no_lock(int cpu_shares) override;
    Status add_thread_to_cgroup() override;

private:
    std::string _cgroup_v1_cpu_tg_path; // workload group path
    std::string _cgroup_v1_cpu_tg_quota_file;
    std::string _cgroup_v1_cpu_tg_shares_file;
    std::string _cgroup_v1_cpu_tg_task_file;
};

/*
    NOTE: cgroup v2 directory structure
    1 root path:
        /sys/fs/cgroup
    
    2 doris home path:
        /sys/fs/cgroup/{doris_home}/

    3 doris home subtree_control file:
        /sys/fs/cgroup/{doris_home}/cgroup.subtree_control
    
    4 query path:
        /sys/fs/cgroup/{doris_home}/query/

    5 query path subtree_control file:
        /sys/fs/cgroup/{doris_home}/query/cgroup.subtree_control

    6 query path procs file:
        /sys/fs/cgroup/{doris_home}/query/cgroup.procs

    7 workload group path:
        /sys/fs/cgroup/{doris_home}/query/{workload_group_id}

    8 workload grou cpu.max file:
        /sys/fs/cgroup/{doris_home}/query/{workload_group_id}/cpu.max

    9 workload grou cpu.weight file:
        /sys/fs/cgroup/{doris_home}/query/{workload_group_id}/cpu.weight

    10 workload group cgroup type file:
        /sys/fs/cgroup/{doris_home}/query/{workload_group_id}/cgroup.type

*/
class CgroupV2CpuCtl : public CgroupCpuCtl {
public:
    CgroupV2CpuCtl(uint64_t tg_id) : CgroupCpuCtl(tg_id) {}
    Status init() override;
    Status modify_cg_cpu_hard_limit_no_lock(int cpu_hard_limit) override;
    Status modify_cg_cpu_soft_limit_no_lock(int cpu_shares) override;
    Status add_thread_to_cgroup() override;

private:
    std::string _cgroup_v2_query_wg_path;
    std::string _cgroup_v2_query_wg_cpu_max_file;
    std::string _cgroup_v2_query_wg_cpu_weight_file;
    std::string _cgroup_v2_query_wg_thread_file;
    std::string _cgroup_v2_query_wg_type_file;
};

} // namespace doris