// 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/jemalloc_control.h"

#include <atomic>

#include <condition_variable>

namespace doris {

#include "common/compile_check_begin.h"

std::mutex JemallocControl::je_purge_dirty_pages_lock;

std::atomic<bool> JemallocControl::je_purge_dirty_pages_notify {false};

std::mutex JemallocControl::je_reset_dirty_decay_lock;

std::atomic<bool> JemallocControl::je_enable_dirty_page {true};

std::condition_variable JemallocControl::je_reset_dirty_decay_cv;

std::atomic<bool> JemallocControl::je_reset_dirty_decay_notify {false};

std::atomic<int64_t> JemallocControl::je_cache_bytes_ = 0;

std::atomic<int64_t> JemallocControl::je_tcache_mem_ = 0;

std::atomic<int64_t> JemallocControl::je_metadata_mem_ = 0;

std::atomic<int64_t> JemallocControl::je_dirty_pages_mem_ = std::numeric_limits<int64_t>::min();

std::atomic<int64_t> JemallocControl::je_virtual_memory_used_ = 0;

void JemallocControl::refresh_allocator_mem() {

#if defined(ADDRESS_SANITIZER) || defined(LEAK_SANITIZER) || defined(THREAD_SANITIZER)

#elif defined(USE_JEMALLOC)

    // jemalloc mallctl refer to : https://jemalloc.net/jemalloc.3.html

    // https://www.bookstack.cn/read/aliyun-rds-core/4a0cdf677f62feb3.md

    //  Check the Doris BE web page `http://ip:webserver_port/memory` to get the Jemalloc Profile.

    // 'epoch' is a special mallctl -- it updates the statistics. Without it, all

    // the following calls will return stale values. It increments and returns

    // the current epoch number, which might be useful to log as a sanity check.

    uint64_t epoch = 0;

    size_t sz = sizeof(epoch);

    jemallctl("epoch", &epoch, &sz, &epoch, sz);

    // Number of extents of the given type in this arena in the bucket corresponding to page size index.

    // Large size class starts at 16384, the extents have three sizes before 16384: 4096, 8192, and 12288, so + 3

    int64_t dirty_pages_bytes = 0;

    for (unsigned i = 0; i < get_jemallctl_value<unsigned>("arenas.nlextents") + 3; i++) {

        dirty_pages_bytes += get_je_all_arena_extents_metrics(i, "dirty_bytes");

    }

    je_dirty_pages_mem_.store(dirty_pages_bytes, std::memory_order_relaxed);

    je_tcache_mem_.store(get_je_all_arena_metrics("tcache_bytes"));

    // Doris uses Jemalloc as default Allocator, Jemalloc Cache consists of two parts:

    // - Thread Cache, cache a specified number of Pages in Thread Cache.

    // - Dirty Page, memory Page that can be reused in all Arenas.

    je_cache_bytes_.store(je_tcache_mem_ + dirty_pages_bytes, std::memory_order_relaxed);

    // Total number of bytes dedicated to metadata, which comprise base allocations used

    // for bootstrap-sensitive allocator metadata structures.

    je_metadata_mem_.store(get_jemallctl_value<int64_t>("stats.metadata"),

                           std::memory_order_relaxed);

    je_virtual_memory_used_.store(get_jemallctl_value<int64_t>("stats.mapped"),

                                  std::memory_order_relaxed);

#else

    je_cache_bytes_.store(get_tc_metrics("tcmalloc.pageheap_free_bytes") +

                                  get_tc_metrics("tcmalloc.central_cache_free_bytes") +

                                  get_tc_metrics("tcmalloc.transfer_cache_free_bytes") +

                                  get_tc_metrics("tcmalloc.thread_cache_free_bytes"),

                          std::memory_order_relaxed);

    je_virtual_memory_used_.store(get_tc_metrics("generic.total_physical_bytes") +

                                          get_tc_metrics("tcmalloc.pageheap_unmapped_bytes"),

                                  std::memory_order_relaxed);

#endif

}

#ifdef USE_JEMALLOC

void JemallocControl::action_jemallctl(const std::string& name) {

    try {

        int err = jemallctl(name.c_str(), nullptr, nullptr, nullptr, 0);

        if (err) {

            LOG(WARNING) << fmt::format("Failed, jemallctl action {}", name);

        }

    } catch (...) {

        LOG(WARNING) << fmt::format("Exception, jemallctl action {}", name);

    }

}

int64_t JemallocControl::get_je_all_arena_metrics(const std::string& name) {

    return get_jemallctl_value<int64_t>(

            fmt::format("stats.arenas.{}.{}", MALLCTL_ARENAS_ALL, name));

}

int64_t JemallocControl::get_je_all_arena_extents_metrics(int64_t page_size_index,

                                                          const std::string& extent_type) {

    return get_jemallctl_value<int64_t>(fmt::format(

            "stats.arenas.{}.extents.{}.{}", MALLCTL_ARENAS_ALL, page_size_index, extent_type));

}

void JemallocControl::je_purge_all_arena_dirty_pages() {

    // https://github.com/jemalloc/jemalloc/issues/2470

    // If there is a core dump here, it may cover up the real stack, if stack trace indicates heap corruption

    // (which led to invalid jemalloc metadata), like double free or use-after-free in the application.

    // Try sanitizers such as ASAN, or build jemalloc with --enable-debug to investigate further.

    action_jemallctl(fmt::format("arena.{}.purge", MALLCTL_ARENAS_ALL));

}

void JemallocControl::je_reset_all_arena_dirty_decay_ms(ssize_t dirty_decay_ms) {

    // Each time this interface is set, all currently unused dirty pages are considered

    // to have fully decayed, which causes immediate purging of all unused dirty pages unless

    // the decay time is set to -1

    //

    // NOTE: Using "arena.MALLCTL_ARENAS_ALL.dirty_decay_ms" to modify all arenas will fail or even crash,

    // which may be a bug.

    for (unsigned i = 0; i < get_jemallctl_value<unsigned>("arenas.narenas"); i++) {

        set_jemallctl_value<ssize_t>(fmt::format("arena.{}.dirty_decay_ms", i), dirty_decay_ms);

    }

}

void JemallocControl::je_decay_all_arena_dirty_pages() {

    // Trigger decay-based purging of unused dirty/muzzy pages for arena <i>, or for all arenas if <i> equals

    // MALLCTL_ARENAS_ALL. The proportion of unused dirty/muzzy pages to be purged depends on the

    // current time; see opt.dirty_decay_ms and opt.muzy_decay_ms for details.

    action_jemallctl(fmt::format("arena.{}.decay", MALLCTL_ARENAS_ALL));

}

void JemallocControl::je_thread_tcache_flush() {

    constexpr size_t TCACHE_LIMIT = (1ULL << 30); // 1G

    if (je_tcache_mem() > TCACHE_LIMIT) {

        action_jemallctl("thread.tcache.flush");

    }

}

#else

void JemallocControl::action_jemallctl(const std::string& name) {}

int64_t JemallocControl::get_je_all_arena_metrics(const std::string& name) {

    return 0;

}

int64_t JemallocControl::get_je_all_arena_extents_metrics(int64_t page_size_index,

                                                          const std::string& extent_type) {

    return 0;

}

void JemallocControl::je_purge_all_arena_dirty_pages() {}

void JemallocControl::je_reset_all_arena_dirty_decay_ms(ssize_t dirty_decay_ms) {}

void JemallocControl::je_decay_all_arena_dirty_pages() {}

void JemallocControl::je_thread_tcache_flush() {}

#endif

#include "common/compile_check_end.h"

} // namespace doris