#include "cpu_limiter.h"

using namespace NSrvKernel;

void TRejectedConnsCleaner::OnInit(TContExecutor* executor) {
    if (HoldCount_ > 0) {
        Cont_ = TCoroutine(
            ECoroType::Service,
            "cpu_limiter_cleaner",
            executor,
            &TRejectedConnsCleaner::RejectedConnsCleanerCont,
            this,
            executor
        );
    }
}

void TRejectedConnsCleaner::Push(TSocketHolder&& socket) {
    if (HoldCount_ > 0) {
        if (RejectedConns_.size() >= HoldCount_) {
            PopConn();
        } else if (RejectedConns_.empty()) {
            NoRejectedConnsCV_.notify();
        }

        TRejectedConn rejectedConn {
            .Deadline = Now() + TDuration::MicroSeconds(RandomNumber<double>() * HoldDuration_.MicroSeconds()),
            .S = std::move(socket)
        };
        PushConn(std::move(rejectedConn));
    }
}

void TRejectedConnsCleaner::RejectedConnsCleanerCont(TContExecutor* executor) {
    for (;;) {
        if (!RejectedConns_.empty()) {
            TInstant deadline = RejectedConns_.front().Deadline;
            if (NoRejectedConnsCV_.wait_until(executor, deadline) == ECANCELED) {
                break;
            }
            while (!RejectedConns_.empty() && RejectedConns_.front().Deadline <= deadline) {
                PopConn();
            }
        } else {
            if (NoRejectedConnsCV_.wait(executor) == ECANCELED) {
                break;
            }
        }
    }
}

TCpuLimiter::TCpuLimiter(const TCpuLimiterConfig& config, TCpuLimiterStat& stat, size_t workerId)
    : Config_(config)
    , Stat_(stat, workerId)
    , Cache_(config.CheckerAddressCacheSize)
    , RejectedConnsCleaner_(std::move(Stat_.ConnHoldCounter), config.ConnHoldCount, config.ConnHoldDuration) {}

void TCpuLimiter::OnInit(TContExecutor* executor, TSharedFiles* sharedFiles) {
    if (IsAnythingEnabled() && Config_.DisableFile) {
        DisableFileChecker_ = sharedFiles->FileChecker(Config_.DisableFile, TDuration::Seconds(1));
    }

    if (Config_.DisableHTTP2File) {
        DisableHTTP2FileChecker_ = sharedFiles->FileChecker(Config_.DisableHTTP2File, TDuration::Seconds(1));
    }

    RejectedConnsCleaner_.OnInit(executor);
}
