// 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 "runtime/memory/heap_profiler.h"

#ifdef USE_JEMALLOC

#include "jemalloc/jemalloc.h"

#endif

#include "agent/utils.h"

#include "common/config.h"

#include "io/fs/local_file_system.h"

namespace doris {

void HeapProfiler::set_prof_active(bool prof) {

#ifdef USE_JEMALLOC

    std::lock_guard guard(_mutex);

    try {

        int err = jemallctl("prof.active", nullptr, nullptr, &prof, 1);

        err |= jemallctl("prof.thread_active_init", nullptr, nullptr, &prof, 1);

        if (err) {

            LOG(WARNING) << "jemalloc heap profiling start failed, " << err;

        } else {

            LOG(INFO) << "jemalloc heap profiling started";

        }

    } catch (...) {

        LOG(WARNING) << "jemalloc heap profiling start failed";

    }

#endif

}

bool HeapProfiler::get_prof_dump(const std::string& profile_file_name) {

#ifdef USE_JEMALLOC

    std::lock_guard guard(_mutex);

    const char* file_name_ptr = profile_file_name.c_str();

    try {

        int err = jemallctl("prof.dump", nullptr, nullptr, &file_name_ptr, sizeof(const char*));

        if (err) {

            LOG(WARNING) << "dump heap profile failed, " << err;

            return false;

        } else {

            LOG(INFO) << "dump heap profile to " << profile_file_name;

            return true;

        }

    } catch (...) {

        LOG(WARNING) << "dump heap profile failed";

        return false;

    }

#else

    return false;

#endif

}

static std::string jeprof_profile_to_dot(const std::string& profile_file_name) {

    AgentUtils util;

    const static std::string jeprof_path = fmt::format("{}/bin/jeprof", std::getenv("DORIS_HOME"));

    const static std::string binary_path =

            fmt::format("{}/lib/doris_be", std::getenv("DORIS_HOME"));

    // https://doris.apache.org/community/developer-guide/debug-tool/#3-jeprof-parses-heap-profile

    std::string jeprof_cmd =

            fmt::format("{} --dot {} {}", jeprof_path, binary_path, profile_file_name);

    std::string msg;

    bool rc = util.exec_cmd(jeprof_cmd, &msg);

    if (!rc) {

        LOG(WARNING) << "jeprof profile to dot failed: " << msg;

    }

    return msg;

}

void HeapProfiler::heap_profiler_start() {

    set_prof_active(true);

}

void HeapProfiler::heap_profiler_stop() {

    set_prof_active(false);

}

bool HeapProfiler::check_active_heap_profiler() {

#ifdef USE_JEMALLOC

    size_t value = 0;

    size_t sz = sizeof(value);

    // A Bug in Jemalloc 5.3, `jemallctl("prof", &value, &sz, nullptr, 0)` always returns `0`,

    // the real conf `prof` value cannot be checked.

    jemallctl("prof.active", &value, &sz, nullptr, 0);

    return value;

#else

    return false;

#endif

}

bool HeapProfiler::check_enable_heap_profiler() {

#ifdef USE_JEMALLOC

    bool prof = false;

    size_t sz = sizeof(prof);

    if (jemallctl("opt.prof", &prof, &sz, nullptr, 0) != 0) {

        LOG(WARNING) << "Failed to get option: opt.prof";

    }

    return prof;

#else

    return false;

#endif

}

std::string HeapProfiler::dump_heap_profile() {

    if (!config::jeprofile_dir.empty()) {

        auto st = io::global_local_filesystem()->create_directory(config::jeprofile_dir);

        if (!st.ok()) {

            LOG(WARNING) << "create jeprofile dir failed.";

            return "";

        }

    }

    std::string profile_file_name =

            fmt::format("{}/jeheap_dump.{}.{}.{}.heap", config::jeprofile_dir, std::time(nullptr),

                        getpid(), rand());

    if (get_prof_dump(profile_file_name)) {

        return profile_file_name;

    } else {

        return "";

    }

}

std::string HeapProfiler::dump_heap_profile_to_dot() {

    std::string profile_file_name = dump_heap_profile();

    if (!profile_file_name.empty()) {

        return jeprof_profile_to_dot(profile_file_name);

    } else {

        return "";

    }

}

bool HeapProfiler::heap_profiler_reset(size_t lg_sample) {

#ifdef USE_JEMALLOC

    int ret = jemallctl("prof.reset", nullptr, nullptr, &lg_sample, sizeof(lg_sample));

    if (ret != 0) {

        LOG(WARNING) << "jemallctl failed to set prof.reset:" << lg_sample

                     << ".Check whether JEMALLOC_CONF is configured with prof:true.";

        return false;

    }

#endif

    return true;

}

} // namespace doris