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
26using namespace OHOS::HiviewDFX;
27using namespace std;
28
29#ifdef LOCK_TO_CPU
30static 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            DFXLOGW("Invalid cpu %{public}d", lockCpu);
38        } else {
39            DFXLOGE("sched_setaffinity failed");
40        }
41        return false;
42    }
43
44    DFXLOGI("Locked to cpu %{public}d\n", lockCpu);
45    return true;
46}
47
48int 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                DFXLOGW("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
83BENCHMARK_MAIN();
84#endif