// 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 "io/fs/hdfs_file_writer.h"

#include <fcntl.h>

#include <fmt/core.h>

#include <chrono>

#include <filesystem>

#include <ostream>

#include <random>

#include <string>

#include <thread>

#include <utility>

#include "common/config.h"

#include "common/logging.h"

#include "common/status.h"

#include "cpp/sync_point.h"

#include "io/cache/block_file_cache.h"

#include "io/cache/block_file_cache_factory.h"

#include "io/cache/file_cache_common.h"

#include "io/fs/err_utils.h"

#include "io/fs/file_writer.h"

#include "io/fs/hdfs_file_system.h"

#include "io/hdfs_util.h"

#include "runtime/exec_env.h"

#include "service/backend_options.h"

#include "util/bvar_helper.h"

#include "util/jni-util.h"

namespace doris::io {

bvar::Adder<uint64_t> hdfs_file_writer_total("hdfs_file_writer_total_num");

bvar::Adder<uint64_t> hdfs_bytes_written_total("hdfs_file_writer_bytes_written");

bvar::Adder<uint64_t> hdfs_file_created_total("hdfs_file_writer_file_created");

bvar::Adder<uint64_t> inflight_hdfs_file_writer("inflight_hdfs_file_writer");

bvar::Adder<uint64_t> hdfs_file_writer_async_close_queuing("hdfs_file_writer_async_close_queuing");

bvar::Adder<uint64_t> hdfs_file_writer_async_close_processing(

        "hdfs_file_writer_async_close_processing");

static constexpr size_t MB = 1024 * 1024;

#ifndef USE_LIBHDFS3

static constexpr size_t CLIENT_WRITE_PACKET_SIZE = 64 * 1024; // 64 KB

#endif

inline std::default_random_engine make_random_engine() {

    return std::default_random_engine(

            static_cast<uint32_t>(std::chrono::steady_clock::now().time_since_epoch().count()));

}

// In practice, we've found that if the import frequency to HDFS is too fast,

// it can cause an OutOfMemoryError (OOM) in the JVM started by the JNI.

// For this, we should have a method to monitor how much JVM memory is currently being used.

// The HdfsWriteMemUsageRecorder class increments a recorded value during hdfsWrite when writing to HDFS.

// The HDFS client will blockingly call hdfsHsync or hdfsCloseFile

// which ensures that the client's buffer is sent to the data node and returned with an acknowledgment before returning to the caller.

// HdfsWriteMemUsageRecorder would reduce the mem usage at that time.

// If the current usage exceeds the maximum set by the user, the current mem acquire would return failure.

// The caller could do sleep to wait for free memory.

class HdfsWriteMemUsageRecorder {

public:

    HdfsWriteMemUsageRecorder() = default;

    ~HdfsWriteMemUsageRecorder() = default;

    size_t max_usage() const {

        return static_cast<size_t>(static_cast<double>(max_jvm_heap_size()) *

                                   config::max_hdfs_wirter_jni_heap_usage_ratio);

    }

    Status acquire_memory(size_t memory_size, int try_time) {

#if defined(USE_LIBHDFS3) || defined(BE_TEST)

        return Status::OK();

#else

        if (!config::enable_hdfs_mem_limiter) {

            return Status::OK();

        }

        auto unit = config::hdfs_jni_write_sleep_milliseconds;

        std::default_random_engine rng = make_random_engine();

        std::uniform_int_distribution<int64_t> u(unit, 2 * unit);

        std::uniform_int_distribution<int64_t> u2(2 * unit, 4 * unit);

        auto duration_ms =

                try_time < (config::hdfs_jni_write_max_retry_time / 2) ? u(rng) : u2(rng);

        std::unique_lock lck {cur_memory_latch};

        cv.wait_for(lck, std::chrono::milliseconds(duration_ms),

                    [&]() { return cur_memory_comsuption + memory_size <= max_usage(); });

        if (cur_memory_comsuption + memory_size > max_usage()) {

            lck.unlock();

            return Status::InternalError<false>(

                    "Run out of Jni jvm heap space, current limit size is {}, max heap size is {}, "

                    "ratio is {}",

                    max_usage(), max_jvm_heap_size(), config::max_hdfs_wirter_jni_heap_usage_ratio);

        }

        cur_memory_comsuption += memory_size;

        return Status::OK();

#endif

    }

    void release_memory(size_t memory_size) {

#if defined(USE_LIBHDFS3) || defined(BE_TEST)

#else

        if (!config::enable_hdfs_mem_limiter) {

            return;

        }

        std::unique_lock lck {cur_memory_latch};

        size_t origin_size = cur_memory_comsuption;

        cur_memory_comsuption -= memory_size;

        if (cur_memory_comsuption < max_usage() && origin_size > max_usage()) {

            cv.notify_all();

        }

#endif

    }

private:

    // clang-format off

    size_t max_jvm_heap_size() const {

        return Jni::Util::get_max_jni_heap_memory_size();

    }

    // clang-format on

    [[maybe_unused]] std::size_t cur_memory_comsuption {0};

    std::mutex cur_memory_latch;

    std::condition_variable cv;

};

static HdfsWriteMemUsageRecorder g_hdfs_write_rate_limiter;

HdfsFileWriter::HdfsFileWriter(Path path, std::shared_ptr<HdfsHandler> handler, hdfsFile hdfs_file,

                               std::string fs_name, const FileWriterOptions* opts)

        : _path(std::move(path)),

          _hdfs_handler(std::move(handler)),

          _hdfs_file(hdfs_file),

          _fs_name(std::move(fs_name)),

          _sync_file_data(opts ? opts->sync_file_data : true),

          _batch_buffer(MB * config::hdfs_write_batch_buffer_size_mb) {

    init_cache_builder(opts, _path);

    hdfs_file_writer_total << 1;

    TEST_SYNC_POINT("HdfsFileWriter");

}

HdfsFileWriter::~HdfsFileWriter() {

    if (_async_close_pack != nullptr) {

        // For thread safety

        std::ignore = _async_close_pack->future.get();

        _async_close_pack = nullptr;

    }

    if (_hdfs_file) {

        SCOPED_BVAR_LATENCY(hdfs_bvar::hdfs_close_latency);

        hdfsCloseFile(_hdfs_handler->hdfs_fs, _hdfs_file);

        inflight_hdfs_file_writer << -1;

        _flush_and_reset_approximate_jni_buffer_size();

    }

}

void HdfsFileWriter::_flush_and_reset_approximate_jni_buffer_size() {

    g_hdfs_write_rate_limiter.release_memory(_approximate_jni_buffer_size);

    _approximate_jni_buffer_size = 0;

}

Status HdfsFileWriter::_acquire_jni_memory(size_t size) {

#ifdef USE_LIBHDFS3

    return Status::OK();

#else

    size_t actual_size = std::max(CLIENT_WRITE_PACKET_SIZE, size);

    int try_time = 0;

    if (auto st = g_hdfs_write_rate_limiter.acquire_memory(actual_size, try_time); !st.ok()) {

        if (_approximate_jni_buffer_size > 0) {

            int ret;

            {

                SCOPED_BVAR_LATENCY(hdfs_bvar::hdfs_hflush_latency);

                ret = SYNC_POINT_HOOK_RETURN_VALUE(hdfsHFlush(_hdfs_handler->hdfs_fs, _hdfs_file),

                                                   "HdfsFileWriter::close::hdfsHFlush");

            }

            _flush_and_reset_approximate_jni_buffer_size();

            if (ret != 0) {

                return Status::InternalError(

                        "Write hdfs file failed. (BE: {}) namenode:{}, path:{}, err: {}, "

                        "file_size={}",

                        BackendOptions::get_localhost(), _fs_name, _path.native(), hdfs_error(),

                        bytes_appended());

            }

        }

        // Other hdfs writers might have occupied too much memory, we need to sleep for a while to wait for them

        // releasing their memory

        for (; try_time < config::hdfs_jni_write_max_retry_time; try_time++) {

            if (g_hdfs_write_rate_limiter.acquire_memory(actual_size, try_time).ok()) {

                _approximate_jni_buffer_size += actual_size;

                return Status::OK();

            }

        }

        return st;

    }

    _approximate_jni_buffer_size += actual_size;

    return Status::OK();

#endif

}

Status HdfsFileWriter::close(bool non_block) {

    if (state() == State::CLOSED) {

        return Status::InternalError("HdfsFileWriter already closed, file path {}, fs name {}",

                                     _path.native(), _fs_name);

    }

    if (state() == State::ASYNC_CLOSING) {

        if (non_block) {

            return Status::InternalError("Don't submit async close multi times");

        }

        CHECK(_async_close_pack != nullptr);

        _st = _async_close_pack->future.get();

        _async_close_pack = nullptr;

        // We should wait for all the pre async task to be finished

        _state = State::CLOSED;

        // The next time we call close() with no matter non_block true or false, it would always return the

        // '_st' value because this writer is already closed.

        return _st;

    }

    if (non_block) {

        _state = State::ASYNC_CLOSING;

        _async_close_pack = std::make_unique<AsyncCloseStatusPack>();

        _async_close_pack->future = _async_close_pack->promise.get_future();

        hdfs_file_writer_async_close_queuing << 1;

        return ExecEnv::GetInstance()->non_block_close_thread_pool()->submit_func([&]() {

            hdfs_file_writer_async_close_queuing << -1;

            hdfs_file_writer_async_close_processing << 1;

            _async_close_pack->promise.set_value(_close_impl());

            hdfs_file_writer_async_close_processing << -1;

        });

    }

    _st = _close_impl();

    _state = State::CLOSED;

    return _st;

}

Status HdfsFileWriter::_close_impl() {

    if (_batch_buffer.size() != 0) {

        if (_st = _flush_buffer(); !_st.ok()) {

            return _st;

        }

    }

    int ret;

    if (_sync_file_data) {

        {

            SCOPED_BVAR_LATENCY(hdfs_bvar::hdfs_hsync_latency);

#ifdef USE_LIBHDFS3

            ret = SYNC_POINT_HOOK_RETURN_VALUE(hdfsSync(_hdfs_handler->hdfs_fs, _hdfs_file),

                                               "HdfsFileWriter::close::hdfsHSync");

#else

            ret = SYNC_POINT_HOOK_RETURN_VALUE(hdfsHSync(_hdfs_handler->hdfs_fs, _hdfs_file),

                                               "HdfsFileWriter::close::hdfsHSync");

            _flush_and_reset_approximate_jni_buffer_size();

#endif

        }

        TEST_INJECTION_POINT_RETURN_WITH_VALUE("HdfsFileWriter::hdfsSync",

                                               Status::InternalError("failed to sync hdfs file"));

        if (ret != 0) {

            _st = Status::InternalError(

                    "failed to sync hdfs file. fs_name={} path={} : {}, file_size={}", _fs_name,

                    _path.native(), hdfs_error(), bytes_appended());

            return _st;

        }

    }

    {

        SCOPED_BVAR_LATENCY(hdfs_bvar::hdfs_close_latency);

        // The underlying implementation will invoke `hdfsHFlush` to flush buffered data and wait for

        // the HDFS response, but won't guarantee the synchronization of data to HDFS.

        ret = SYNC_POINT_HOOK_RETURN_VALUE(hdfsCloseFile(_hdfs_handler->hdfs_fs, _hdfs_file),

                                           "HdfsFileWriter::close::hdfsCloseFile");

        inflight_hdfs_file_writer << -1;

        _flush_and_reset_approximate_jni_buffer_size();

    }

    _hdfs_file = nullptr;

    TEST_INJECTION_POINT_RETURN_WITH_VALUE("HdfsFileWriter::hdfsCloseFile",

                                           Status::InternalError("failed to close hdfs file"));

    if (ret != 0) {

        _st = Status::InternalError(

                "Write hdfs file failed. (BE: {}) namenode:{}, path:{}, err: {}, file_size={}",

                BackendOptions::get_localhost(), _fs_name, _path.native(), hdfs_error(),

                bytes_appended());

        return _st;

    }

    hdfs_file_created_total << 1;

    return Status::OK();

}

HdfsFileWriter::BatchBuffer::BatchBuffer(size_t capacity) {

    _batch_buffer.reserve(capacity);

}

bool HdfsFileWriter::BatchBuffer::full() const {

    return size() == capacity();

}

const char* HdfsFileWriter::BatchBuffer::data() const {

    return _batch_buffer.data();

}

size_t HdfsFileWriter::BatchBuffer::capacity() const {

    return _batch_buffer.capacity();

}

size_t HdfsFileWriter::BatchBuffer::size() const {

    return _batch_buffer.size();

}

void HdfsFileWriter::BatchBuffer::clear() {

    _batch_buffer.clear();

}

// TODO(ByteYue): Refactor Upload Buffer to reduce this duplicate code

void HdfsFileWriter::_write_into_local_file_cache() {

    int64_t tablet_id = get_tablet_id(_path.native()).value_or(0);

    auto holder = _cache_builder->allocate_cache_holder(_bytes_appended - _batch_buffer.size(),

                                                        _batch_buffer.capacity(), tablet_id);

    size_t pos = 0;

    size_t data_remain_size = _batch_buffer.size();

    for (auto& block : holder->file_blocks) {

        if (data_remain_size == 0) {

            break;

        }

        size_t block_size = block->range().size();

        size_t append_size = std::min(data_remain_size, block_size);

        if (block->state() == FileBlock::State::EMPTY) {

            block->get_or_set_downloader();

            if (block->is_downloader()) {

                Slice s(_batch_buffer.data() + pos, append_size);

                Status st = block->append(s);

                if (st.ok()) {

                    st = block->finalize();

                }

                if (!st.ok()) {

                    LOG_WARNING("failed to append data to file cache").error(st);

                }

            }

        }

        data_remain_size -= append_size;

        pos += append_size;

    }

}

Status HdfsFileWriter::append_hdfs_file(std::string_view content) {

    RETURN_IF_ERROR(_acquire_jni_memory(content.size()));

    while (!content.empty()) {

        int64_t written_bytes;

        {

            TEST_INJECTION_POINT_CALLBACK("HdfsFileWriter::append_hdfs_file_delay");

            SCOPED_BVAR_LATENCY(hdfs_bvar::hdfs_write_latency);

            int64_t max_to_write = content.size();

            tSize to_write = static_cast<tSize>(std::min(

                    max_to_write, static_cast<int64_t>(std::numeric_limits<tSize>::max())));

            written_bytes = SYNC_POINT_HOOK_RETURN_VALUE(

                    hdfsWrite(_hdfs_handler->hdfs_fs, _hdfs_file, content.data(), to_write),

                    "HdfsFileWriter::append_hdfs_file::hdfsWrite", content);

            {

                TEST_INJECTION_POINT_RETURN_WITH_VALUE(

                        "HdfsFileWriter::append_hdfs_file_error",

                        Status::InternalError(

                                "write hdfs failed. fs_name: {}, path: {}, error: inject error",

                                _fs_name, _path.native()));

            }

        }

        if (written_bytes < 0) {

            return Status::InternalError(

                    "write hdfs failed. fs_name: {}, path: {}, error: {}, file_size={}", _fs_name,

                    _path.native(), hdfs_error(), bytes_appended());

        }

        hdfs_bytes_written_total << written_bytes;

        content.remove_prefix(written_bytes);

    }

    return Status::OK();

}

Status HdfsFileWriter::_flush_buffer() {

    RETURN_IF_ERROR(append_hdfs_file(_batch_buffer.content()));

    if (_cache_builder != nullptr) {

        _write_into_local_file_cache();

    }

    _batch_buffer.clear();

    return Status::OK();

}

size_t HdfsFileWriter::BatchBuffer::append(std::string_view content) {

    size_t append_size = std::min(capacity() - size(), content.size());

    _batch_buffer.append(content.data(), append_size);

    return append_size;

}

std::string_view HdfsFileWriter::BatchBuffer::content() const {

    return _batch_buffer;

}

Status HdfsFileWriter::_append(std::string_view content) {

    while (!content.empty()) {

        if (_batch_buffer.full()) {

            auto error_msg = fmt::format("invalid batch buffer status, capacity {}, size {}",

                                         _batch_buffer.capacity(), _batch_buffer.size());

            return Status::InternalError(error_msg);

        }

        size_t append_size = _batch_buffer.append(content);

        content.remove_prefix(append_size);

        _bytes_appended += append_size;

        if (_batch_buffer.full()) {

            RETURN_IF_ERROR(_flush_buffer());

        }

    }

    return Status::OK();

}

Status HdfsFileWriter::appendv(const Slice* data, size_t data_cnt) {

    if (_state != State::OPENED) [[unlikely]] {

        return Status::InternalError("append to closed file: {}", _path.native());

    }

    for (size_t i = 0; i < data_cnt; i++) {

        RETURN_IF_ERROR(_append({data[i].get_data(), data[i].get_size()}));

    }

    return Status::OK();

}

Result<FileWriterPtr> HdfsFileWriter::create(Path full_path, std::shared_ptr<HdfsHandler> handler,

                                             const std::string& fs_name,

                                             const FileWriterOptions* opts) {

    auto path = convert_path(full_path, fs_name);

#ifdef USE_LIBHDFS3

    std::string hdfs_dir = path.parent_path().string();

    int exists = hdfsExists(handler->hdfs_fs, hdfs_dir.c_str());

    if (exists != 0) {

        VLOG_NOTICE << "hdfs dir doesn't exist, create it: " << hdfs_dir;

        int ret = hdfsCreateDirectory(handler->hdfs_fs, hdfs_dir.c_str());

        if (ret != 0) {

            std::stringstream ss;

            ss << "create dir failed. "

               << " fs_name: " << fs_name << " path: " << hdfs_dir << ", err: " << hdfs_error();

            LOG(WARNING) << ss.str();

            return ResultError(Status::InternalError(ss.str()));

        }

    }

#endif

    // open file

    hdfsFile hdfs_file = nullptr;

    {

        SCOPED_BVAR_LATENCY(hdfs_bvar::hdfs_open_latency);

        hdfs_file = hdfsOpenFile(handler->hdfs_fs, path.c_str(), O_WRONLY, 0, 0, 0);

    }

    if (hdfs_file == nullptr) {

        std::stringstream ss;

        ss << "open file failed. "

           << " fs_name:" << fs_name << " path:" << path << ", err: " << hdfs_error();

        LOG(WARNING) << ss.str();

        return ResultError(Status::InternalError(ss.str()));

    }

    VLOG_NOTICE << "open file. fs_name:" << fs_name << ", path:" << path;

    inflight_hdfs_file_writer << 1;

    return std::make_unique<HdfsFileWriter>(std::move(path), handler, hdfs_file, fs_name, opts);

}

} // namespace doris::io