1 /*
2  * Copyright (c) 2023 Huawei Device Co., Ltd.
3  * Licensed under the Apache License, Version 2.0 (the "License");
4  * you may not use this file except in compliance with the License.
5  * You may obtain a copy of the License at
6  *
7  *     http://www.apache.org/licenses/LICENSE-2.0
8  *
9  * Unless required by applicable law or agreed to in writing, software
10  * distributed under the License is distributed on an "AS IS" BASIS,
11  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12  * See the License for the specific language governing permissions and
13  * limitations under the License.
14  */
15 
16 #include <benchmark/benchmark.h>
17 #ifdef LOCK_TO_CPU
18 #include <limits.h>
19 #include <sched.h>
20 #include <string>
21 #include <vector>
22 #include "dfx_log.h"
23 #include "string_util.h"
24 #endif
25 
26 using namespace OHOS::HiviewDFX;
27 using namespace std;
28 
29 #ifdef LOCK_TO_CPU
LockToCpu(int lockCpu)30 static bool LockToCpu(int lockCpu)
31 {
32     cpu_set_t cpuSet;
33     CPU_ZERO(&cpuSet);
34     CPU_SET(lockCpu, &cpuSet);
35     if (sched_setaffinity(0, sizeof(cpuSet), &cpuSet) != 0) {
36         if (errno == EINVAL) {
37             LOGW("Invalid cpu %d", lockCpu);
38         } else {
39             LOGE("sched_setaffinity failed");
40         }
41         return false;
42     }
43 
44     LOGI("Locked to cpu %d\n", lockCpu);
45     return true;
46 }
47 
main(int argc,char ** argv)48 int main(int argc, char** argv)
49 {
50     std::vector<char*> newArgv;
51     newArgv.push_back(argv[0]);
52     // Look for the special option --benchmark_lock_cpu.
53     int lockCpu = -1;
54     for (int i = 1; i < argc; i++) {
55         if (StartsWith(argv[i], "--benchmark_cpu=")) {
56             char* endptr = nullptr;
57             long cpu = strtol(&argv[i][16], &endptr, 10);
58             if (endptr == nullptr || *endptr != '\0' || cpu > INT_MAX || cpu < 0) {
59                 LOGW("Malformed value for --benchmark_cpu, requires a valid positive number.");
60                 return 1;
61             }
62             lockCpu = cpu;
63         } else {
64             newArgv.push_back(argv[i]);
65         }
66     }
67     newArgv.push_back(nullptr);
68     if (lockCpu != -1 && !LockToCpu(lockCpu)) {
69         return 1;
70     }
71 
72     int newArgc = newArgv.size() - 1;
73     ::benchmark::Initialize(&newArgc, newArgv.data());
74     if (::benchmark::ReportUnrecognizedArguments(newArgc, newArgv.data())) {
75         return 1;
76     }
77     ::benchmark::RunSpecifiedBenchmarks();
78     ::benchmark::Shutdown();
79 }
80 
81 #else
82 
83 BENCHMARK_MAIN();
84 #endif