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

#include <bthread/bthread.h>

#include <glog/logging.h> // IWYU pragma: export

#include <atomic>

#include <chrono>

#include <mutex>

#include <thread>

#if defined(__APPLE__)

#include <ctime>

#endif

namespace doris {

// Just 10^6.

static constexpr auto NS = 1000000000UL;

bvar::Adder<int64_t> s3_get_rate_limit_sleep_ns("s3_get_rate_limit_sleep_ns");

bvar::Adder<int64_t> s3_get_rate_limit_sleep_count("s3_get_rate_limit_sleep_count");

bvar::Adder<int64_t> s3_get_rate_limit_rejected_count("s3_get_rate_limit_rejected_count");

bvar::Adder<int64_t> s3_put_rate_limit_sleep_ns("s3_put_rate_limit_sleep_ns");

bvar::Adder<int64_t> s3_put_rate_limit_sleep_count("s3_put_rate_limit_sleep_count");

bvar::Adder<int64_t> s3_put_rate_limit_rejected_count("s3_put_rate_limit_rejected_count");

static std::atomic<int64_t> s3_get_rate_limit_sleep_log_count {0};

static std::atomic<int64_t> s3_get_rate_limit_rejected_log_count {0};

static std::atomic<int64_t> s3_put_rate_limit_sleep_log_count {0};

static std::atomic<int64_t> s3_put_rate_limit_rejected_log_count {0};

class TokenBucketRateLimiter::SimpleSpinLock {

public:

    SimpleSpinLock() = default;

    ~SimpleSpinLock() = default;

    void lock() {

        int spin_count = 0;

        static constexpr int MAX_SPIN_COUNT = 50;

        while (_flag.test_and_set(std::memory_order_acq_rel)) {

            spin_count++;

            if (spin_count >= MAX_SPIN_COUNT) {

                LOG(WARNING) << "Warning: Excessive spinning detected while acquiring lock. Spin "

                                "count: "

                             << spin_count;

                spin_count = 0;

            }

            // Spin until we acquire the lock

        }

    }

    void unlock() { _flag.clear(std::memory_order_release); }

private:

    std::atomic_flag _flag = ATOMIC_FLAG_INIT;

};

TokenBucketRateLimiter::TokenBucketRateLimiter(size_t max_speed, size_t max_burst, size_t limit)

        : _max_speed(max_speed),

          _max_burst(max_burst),

          _limit(limit),

          _mutex(std::make_unique<TokenBucketRateLimiter::SimpleSpinLock>()),

          _remain_tokens(max_burst) {}

TokenBucketRateLimiter::~TokenBucketRateLimiter() = default;

TokenBucketRateLimiterHolder::~TokenBucketRateLimiterHolder() = default;

std::pair<size_t, double> TokenBucketRateLimiter::_update_remain_token(long now, size_t amount) {

    // Values obtained under lock to be checked after release

    size_t count_value;

    double tokens_value;

    {

        std::lock_guard<SimpleSpinLock> lock(*_mutex);

        now = (now < _prev_ns_count) ? _prev_ns_count : now;

        if (_max_speed) {

            double delta_seconds =

                    _prev_ns_count ? static_cast<double>(now - _prev_ns_count) / NS : 0;

            _remain_tokens = std::min<double>(_remain_tokens + _max_speed * delta_seconds - amount,

                                              _max_burst);

        }

        _count += amount;

        count_value = _count;

        tokens_value = _remain_tokens;

        _prev_ns_count = now;

    }

    return {count_value, tokens_value};

}

int64_t TokenBucketRateLimiter::add(size_t amount) {

    // Values obtained under lock to be checked after release

    auto duration = std::chrono::steady_clock::now().time_since_epoch();

    auto time_nano_count = std::chrono::duration_cast<std::chrono::nanoseconds>(duration).count();

    auto [count_value, tokens_value] = _update_remain_token(time_nano_count, amount);

    if (_limit && count_value > _limit) {

        // CK would throw exception

        return -1;

    }

    // Wait unless there is positive amount of remain_tokens - throttling

    int64_t sleep_time_ns = 0;

    if (_max_speed && tokens_value < 0) {

        sleep_time_ns = static_cast<int64_t>(-tokens_value / _max_speed * NS);

        bthread_usleep(sleep_time_ns / 1000);

    }

    return sleep_time_ns;

}

TokenBucketRateLimiterHolder::TokenBucketRateLimiterHolder(size_t max_speed, size_t max_burst,

                                                           size_t limit,

                                                           std::function<void(int64_t)> metric_func)

        : rate_limiter(std::make_unique<TokenBucketRateLimiter>(max_speed, max_burst, limit)),

          metric_func(std::move(metric_func)) {}

int64_t TokenBucketRateLimiterHolder::add(size_t amount) {

    return add_with_config(amount).sleep_duration;

}

TokenBucketRateLimiterResult TokenBucketRateLimiterHolder::add_with_config(size_t amount) {

    TokenBucketRateLimiterResult result;

    {

        std::shared_lock read {rate_limiter_rw_lock};

        result = {.sleep_duration = rate_limiter->add(amount),

                  .max_speed = rate_limiter->get_max_speed(),

                  .max_burst = rate_limiter->get_max_burst(),

                  .limit = rate_limiter->get_limit()};

    }

    metric_func(result.sleep_duration);

    return result;

}

int TokenBucketRateLimiterHolder::reset(size_t max_speed, size_t max_burst, size_t limit) {

    {

        std::unique_lock write {rate_limiter_rw_lock};

        rate_limiter = std::make_unique<TokenBucketRateLimiter>(max_speed, max_burst, limit);

    }

    return 0;

}

size_t TokenBucketRateLimiterHolder::get_max_speed() const {

    std::shared_lock read {rate_limiter_rw_lock};

    return rate_limiter->get_max_speed();

}

size_t TokenBucketRateLimiterHolder::get_max_burst() const {

    std::shared_lock read {rate_limiter_rw_lock};

    return rate_limiter->get_max_burst();

}

size_t TokenBucketRateLimiterHolder::get_limit() const {

    std::shared_lock read {rate_limiter_rw_lock};

    return rate_limiter->get_limit();

}

std::string to_string(S3RateLimitType type) {

    switch (type) {

    case S3RateLimitType::GET:

        return "get";

    case S3RateLimitType::PUT:

        return "put";

    default:

        return std::to_string(static_cast<size_t>(type));

    }

}

S3RateLimitType string_to_s3_rate_limit_type(std::string_view value) {

    if (value == "get") {

        return S3RateLimitType::GET;

    } else if (value == "put") {

        return S3RateLimitType::PUT;

    }

    return S3RateLimitType::UNKNOWN;

}

std::function<void(int64_t)> s3_rate_limiter_metric_func(S3RateLimitType type) {

    switch (type) {

    case S3RateLimitType::GET:

        return metric_func_factory(s3_get_rate_limit_sleep_ns, s3_get_rate_limit_sleep_count,

                                   &s3_get_rate_limit_rejected_count);

    case S3RateLimitType::PUT:

        return metric_func_factory(s3_put_rate_limit_sleep_ns, s3_put_rate_limit_sleep_count,

                                   &s3_put_rate_limit_rejected_count);

    default:

        return [](int64_t) {};

    }

}

int64_t apply_s3_rate_limit(S3RateLimitType type, S3RateLimiterHolder* rate_limiter,

                            int64_t log_interval) {

    auto result = rate_limiter->add_with_config(1);

    auto sleep_duration = result.sleep_duration;

    if (log_interval <= 0 || sleep_duration == 0) {

        return sleep_duration;

    }

    auto is_get = type == S3RateLimitType::GET;

    auto* sleep_log_count =

            is_get ? &s3_get_rate_limit_sleep_log_count : &s3_put_rate_limit_sleep_log_count;

    auto* rejected_log_count =

            is_get ? &s3_get_rate_limit_rejected_log_count : &s3_put_rate_limit_rejected_log_count;

    if (sleep_duration > 0) {

        int64_t count = sleep_log_count->fetch_add(1, std::memory_order_relaxed) + 1;

        if (count == 1 || count % log_interval == 0) {

            LOG(INFO) << "S3 " << to_string(type) << " request is throttled by local rate limiter"

                      << ", sleep_ms=" << sleep_duration / 1000000 << ", sleep_count=" << count

                      << ", token_per_second=" << result.max_speed

                      << ", bucket_tokens=" << result.max_burst << ", token_limit=" << result.limit;

        }

    } else {

        int64_t count = rejected_log_count->fetch_add(1, std::memory_order_relaxed) + 1;

        if (count == 1 || count % log_interval == 0) {

            LOG(WARNING) << "S3 " << to_string(type) << " request is rejected by local rate limiter"

                         << ", rejected_count=" << count

                         << ", token_per_second=" << result.max_speed

                         << ", bucket_tokens=" << result.max_burst

                         << ", token_limit=" << result.limit;

        }

    }

    return sleep_duration;

}

} // namespace doris