208 lines
7.5 KiB
C++
208 lines
7.5 KiB
C++
// 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.
|
|
|
|
#include <algorithm>
|
|
#include <fstream>
|
|
#include "gtest/gtest.h"
|
|
#include "gmock/gmock.h"
|
|
#include "agent/cgroups_mgr.h"
|
|
#include "boost/filesystem.hpp"
|
|
#include "util/logging.h"
|
|
|
|
#ifndef BE_TEST
|
|
#define BE_TEST
|
|
#endif
|
|
|
|
using ::testing::_;
|
|
using ::testing::Return;
|
|
using ::testing::SetArgPointee;
|
|
using std::string;
|
|
|
|
namespace doris {
|
|
|
|
class CgroupsMgrTest : public testing::Test {
|
|
public:
|
|
// create a mock cgroup folder
|
|
static void SetUpTestCase() {
|
|
ASSERT_FALSE(boost::filesystem::exists(_s_cgroup_path));
|
|
// create a mock cgroup path
|
|
ASSERT_TRUE(boost::filesystem::create_directory(_s_cgroup_path));
|
|
}
|
|
|
|
// delete the mock cgroup folder
|
|
static void TearDownTestCase() {
|
|
ASSERT_TRUE(boost::filesystem::remove_all(_s_cgroup_path));
|
|
}
|
|
|
|
// test if a file contains specific number
|
|
static bool does_contain_number(const std::string& file_path, int32_t number) {
|
|
std::ifstream input_file(file_path.c_str());
|
|
int32_t task_id;
|
|
while (input_file >> task_id) {
|
|
if (task_id == number) {
|
|
return true;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static std::string _s_cgroup_path;
|
|
static CgroupsMgr _s_cgroups_mgr;
|
|
};
|
|
|
|
std::string CgroupsMgrTest::_s_cgroup_path = "./doris_cgroup_testxxxx123";
|
|
CgroupsMgr CgroupsMgrTest::_s_cgroups_mgr(NULL, CgroupsMgrTest::_s_cgroup_path);
|
|
|
|
TEST_F(CgroupsMgrTest, TestIsDirectory) {
|
|
// test folder exist
|
|
bool exist = _s_cgroups_mgr.is_directory(CgroupsMgrTest::_s_cgroup_path.c_str());
|
|
ASSERT_TRUE(exist);
|
|
// test folder not exist
|
|
bool not_exist = _s_cgroups_mgr.is_directory("./abc");
|
|
ASSERT_FALSE(not_exist);
|
|
// test file exist, but not folder
|
|
bool not_folder = _s_cgroups_mgr.is_directory("/etc/profile");
|
|
ASSERT_FALSE(not_folder);
|
|
|
|
}
|
|
|
|
TEST_F(CgroupsMgrTest, TestIsFileExist) {
|
|
// test file exist
|
|
bool exist = _s_cgroups_mgr.is_file_exist(CgroupsMgrTest::_s_cgroup_path.c_str());
|
|
ASSERT_TRUE(exist);
|
|
// test file not exist
|
|
bool not_exist = _s_cgroups_mgr.is_file_exist("./abc");
|
|
ASSERT_FALSE(not_exist);
|
|
}
|
|
|
|
TEST_F(CgroupsMgrTest, TestInitCgroups) {
|
|
// test for task file not exist
|
|
AgentStatus op_status = _s_cgroups_mgr.init_cgroups();
|
|
ASSERT_EQ(AgentStatus::DORIS_ERROR, op_status);
|
|
|
|
// create task file, then init should success
|
|
std::string task_file_path = _s_cgroup_path + "/tasks";
|
|
std::ofstream outfile(task_file_path.c_str());
|
|
outfile << 1111111 << std::endl;
|
|
outfile.close();
|
|
|
|
// create a mock user under cgroup path
|
|
ASSERT_TRUE(boost::filesystem::create_directory(_s_cgroup_path + "/yiguolei"));
|
|
std::ofstream user_out_file(_s_cgroup_path + "/yiguolei/tasks");
|
|
user_out_file << 123 << std::endl;
|
|
user_out_file.close();
|
|
|
|
// create a mock user group under cgroup path
|
|
ASSERT_TRUE(boost::filesystem::create_directory(_s_cgroup_path + "/yiguolei/low"));
|
|
std::ofstream group_out_file(CgroupsMgrTest::_s_cgroup_path + "/yiguolei/low/tasks");
|
|
group_out_file << 456 << std::endl;
|
|
group_out_file.close();
|
|
|
|
op_status = _s_cgroups_mgr.init_cgroups();
|
|
// init should be successful
|
|
ASSERT_EQ(AgentStatus::DORIS_SUCCESS, op_status);
|
|
// all tasks should be migrated to root cgroup path
|
|
ASSERT_TRUE(does_contain_number(task_file_path, 1111111));
|
|
ASSERT_TRUE(does_contain_number(task_file_path, 123));
|
|
ASSERT_TRUE(does_contain_number(task_file_path, 456));
|
|
}
|
|
|
|
TEST_F(CgroupsMgrTest, TestAssignThreadToCgroups) {
|
|
// default cgroup not exist, so that assign to an unknown user will fail
|
|
AgentStatus op_status = _s_cgroups_mgr.assign_thread_to_cgroups(111,
|
|
"abc",
|
|
"low");
|
|
ASSERT_EQ(AgentStatus::DORIS_ERROR, op_status);
|
|
// user cgroup exist
|
|
// create a mock user under cgroup path
|
|
ASSERT_TRUE(boost::filesystem::create_directory(_s_cgroup_path + "/yiguolei2"));
|
|
std::ofstream user_out_file(_s_cgroup_path + "/yiguolei2/tasks");
|
|
user_out_file << 123 << std::endl;
|
|
user_out_file.close();
|
|
|
|
op_status = _s_cgroups_mgr.assign_thread_to_cgroups(111,
|
|
"yiguolei2",
|
|
"aaaa");
|
|
ASSERT_EQ(AgentStatus::DORIS_SUCCESS, op_status);
|
|
ASSERT_TRUE(does_contain_number(_s_cgroup_path + "/yiguolei2/tasks", 111));
|
|
|
|
// user,level cgroup exist
|
|
// create a mock user group under cgroup path
|
|
ASSERT_TRUE(boost::filesystem::create_directory(_s_cgroup_path + "/yiguolei2/low"));
|
|
std::ofstream group_out_file(_s_cgroup_path + "/yiguolei2/low/tasks");
|
|
group_out_file << 456 << std::endl;
|
|
group_out_file.close();
|
|
|
|
op_status = _s_cgroups_mgr.assign_thread_to_cgroups(111,
|
|
"yiguolei2",
|
|
"low");
|
|
ASSERT_EQ(AgentStatus::DORIS_SUCCESS, op_status);
|
|
ASSERT_TRUE(does_contain_number(_s_cgroup_path + "/yiguolei2/low/tasks", 111));
|
|
}
|
|
|
|
TEST_F(CgroupsMgrTest, TestModifyUserCgroups) {
|
|
std::map<std::string, int32_t> user_share;
|
|
std::map<std::string, int32_t> level_share;
|
|
user_share["cpu.shares"] = 100;
|
|
level_share["low"] = 100;
|
|
std::string user_name = "user_modify";
|
|
AgentStatus op_status = _s_cgroups_mgr.modify_user_cgroups(user_name,
|
|
user_share,
|
|
level_share);
|
|
|
|
ASSERT_EQ(AgentStatus::DORIS_SUCCESS, op_status);
|
|
|
|
ASSERT_TRUE(does_contain_number(_s_cgroup_path + "/user_modify/cpu.shares", 100));
|
|
ASSERT_TRUE(does_contain_number(_s_cgroup_path + "/user_modify/low/cpu.shares", 100));
|
|
}
|
|
|
|
TEST_F(CgroupsMgrTest, TestUpdateLocalCgroups) {
|
|
// mock TFetchResourceResult from fe
|
|
TFetchResourceResult user_resource_result;
|
|
|
|
TUserResource user_resource;
|
|
user_resource.shareByGroup["low"] = 123;
|
|
user_resource.shareByGroup["normal"] = 234;
|
|
TResourceGroup resource_group;
|
|
resource_group.resourceByType[TResourceType::type::TRESOURCE_CPU_SHARE] = 100;
|
|
user_resource.resource = resource_group;
|
|
|
|
user_resource_result.resourceVersion = 2;
|
|
user_resource_result.resourceByUser["yiguolei3"] = user_resource;
|
|
|
|
AgentStatus op_status = _s_cgroups_mgr.update_local_cgroups(user_resource_result);
|
|
ASSERT_EQ(AgentStatus::DORIS_SUCCESS, op_status);
|
|
ASSERT_EQ(2, _s_cgroups_mgr._cur_version);
|
|
ASSERT_TRUE(does_contain_number(_s_cgroup_path + "/yiguolei3/cpu.shares", 100));
|
|
ASSERT_TRUE(does_contain_number(_s_cgroup_path + "/yiguolei3/low/cpu.shares", 123));
|
|
ASSERT_TRUE(does_contain_number(_s_cgroup_path + "/yiguolei3/normal/cpu.shares", 234));
|
|
}
|
|
|
|
TEST_F(CgroupsMgrTest, TestRelocateTasks) {
|
|
// create a source cgroup, add some taskid into it
|
|
AgentStatus op_status = _s_cgroups_mgr.relocate_tasks("/a/b/c/d", _s_cgroup_path);
|
|
ASSERT_EQ(AgentStatus::DORIS_ERROR, op_status);
|
|
}
|
|
|
|
} // namespace doris
|
|
|
|
int main(int argc, char **argv) {
|
|
doris::init_glog("be-test");
|
|
::testing::InitGoogleTest(&argc, argv);
|
|
return RUN_ALL_TESTS();
|
|
}
|