/*
 * Copyright 2019 Bloomberg Finance LP
 *
 * Licensed 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 <buildboxcommon_lrulocalcas.h>

#include <buildboxcommon_exception.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_logging.h>

#include <queue>
#include <stdexcept>

#include <buildboxcommon_logging.h>
#include <buildboxcommon_metricnames.h>
#include <buildboxcommonmetrics_countingmetricutil.h>
#include <buildboxcommonmetrics_durationmetrictimer.h>
#include <buildboxcommonmetrics_gaugemetricutil.h>
#include <buildboxcommonmetrics_metricguard.h>

using namespace buildboxcommon;
using namespace std::chrono_literals;

using buildboxcommon::buildboxcommonmetrics::CountingMetricUtil;
using buildboxcommon::buildboxcommonmetrics::GaugeMetricUtil;

LruLocalCas::LruLocalCas(std::unique_ptr<LocalCas> baseCas,
                         double quotaLowRatio,
                         std::optional<int64_t> quotaHigh,
                         int64_t reservedSpace)
    : d_base_cas(std::move(baseCas)), d_quotaLowRatio(quotaLowRatio),
      d_configuredQuotaHigh(quotaHigh), d_reserved_space(reservedSpace),
      d_pending_write_space(0), d_cancel_cleanup(false)
{
    if (quotaHigh.has_value()) {
        BUILDBOX_LOG_INFO(
            "Initializing LRU with quota: " << quotaHigh.value());
        GaugeMetricUtil::setGauge(
            CommonMetricNames::GAUGE_NAME_LRU_CONFIGURED_HIGH_WATERMARK,
            quotaHigh.value());
    }
    else {
        BUILDBOX_LOG_INFO("Initializing LRU without quota");
    }
    BUILDBOX_LOG_INFO("Low watermark: " << (quotaLowRatio * 100) << "%");
    BUILDBOX_LOG_INFO("Reserved disk space: " << reservedSpace);
    const int PERCENT = 100;
    GaugeMetricUtil::setGauge(
        CommonMetricNames::GAUGE_NAME_LRU_LOW_WATERMARK_PERCENTAGE,
        static_cast<long>(quotaLowRatio * PERCENT));

    // Set initial effective disk quota based on available disk space.
    updateDiskQuota();

    // Start thread to monitor available disk space.
    d_disk_monitor_thread = std::thread(&LruLocalCas::diskMonitorThread, this);
}

LruLocalCas::~LruLocalCas()
{
    // Indicate potentially running cleanup to stop as soon as possible.
    d_cancel_cleanup = true;

    // Stop and wait for disk monitor thread.
    {
        std::lock_guard<std::mutex> lock(d_stop_mutex);
        d_stop = true;
    }
    d_stop_cv.notify_all();
    d_disk_monitor_thread.join();
}

void LruLocalCas::protectActiveBlobs(blob_time_type start_time)
{
    d_protect_active_blobs = true;
    d_protect_active_time = start_time;
}

bool LruLocalCas::hasBlob(const Digest &digest)
{
    return d_base_cas->hasBlob(digest);
}

bool LruLocalCas::writeBlob(const Digest &digest, const std::string &data)
{
    const auto write_function = [&] {
        return d_base_cas->writeBlob(digest, data);
    };

    return tryToAddBlob(digest.size_bytes(), write_function);
}

bool LruLocalCas::writeBlob(int fd, Digest *digest)
{
    const int64_t file_size = buildboxcommon::FileUtils::getFileSize(fd);

    const auto write_function = [&] {
        return d_base_cas->writeBlob(fd, digest);
    };

    return tryToAddBlob(file_size, write_function);
}

bool LruLocalCas::moveBlobFromTemporaryFile(
    const Digest &digest, buildboxcommon::TemporaryFile tempFile)
{
    const auto write_function = [&] {
        return d_base_cas->moveBlobFromTemporaryFile(digest,
                                                     std::move(tempFile));
    };

    return tryToAddBlob(digest.size_bytes(), write_function);
}

bool LruLocalCas::moveBlobFromTemporaryFile(
    const Digest &digest, const buildboxcommon::TemporaryDirectory &directory,
    const std::string &filename)
{
    const auto write_function = [&] {
        return d_base_cas->moveBlobFromTemporaryFile(digest, directory,
                                                     filename);
    };

    return tryToAddBlob(digest.size_bytes(), write_function);
}

std::unique_ptr<std::string>
LruLocalCas::readBlob(const Digest &digest, int64_t offset, size_t length)
{
    return d_base_cas->readBlob(digest, offset, length);
}

std::unique_ptr<std::string> LruLocalCas::readBlob(const Digest &digest)
{
    return d_base_cas->readBlob(digest);
}

bool LruLocalCas::copyBlob(const Digest &digest, int dirfd, const char *path,
                           mode_t mode, struct stat *destStat)
{
    return d_base_cas->copyBlob(digest, dirfd, path, mode, destStat);
}

bool LruLocalCas::deleteBlob(const Digest &digest)
{
    return d_base_cas->deleteBlob(digest);
}

bool LruLocalCas::deleteBlob(const Digest &digest,
                             const blob_time_type &if_unmodified_since)
{
    return d_base_cas->deleteBlob(digest, if_unmodified_since);
}

void LruLocalCas::listBlobs(digest_time_callback_type callback)
{
    d_base_cas->listBlobs(callback);
}

std::string LruLocalCas::path(const Digest &digest)
{
    return d_base_cas->path(digest);
}

std::string LruLocalCas::pathUnchecked(const Digest &digest)
{
    return d_base_cas->pathUnchecked(digest);
}

FdNamePair LruLocalCas::pathAt(const Digest &digest)
{
    return d_base_cas->pathAt(digest);
}

FdNamePair LruLocalCas::pathAtUnchecked(const Digest &digest)
{
    return d_base_cas->pathAtUnchecked(digest);
}

std::string const &LruLocalCas::path() { return d_base_cas->path(); }

buildboxcommon::TemporaryDirectory LruLocalCas::createStagingDirectory()
{
    return d_base_cas->createStagingDirectory();
}

buildboxcommon::TemporaryFile LruLocalCas::createTemporaryFile()
{
    return d_base_cas->createTemporaryFile();
}

buildboxcommon::TemporaryDirectory LruLocalCas::createTemporaryDirectory()
{
    return d_base_cas->createTemporaryDirectory();
}

int64_t LruLocalCas::getDiskUsage(bool force_scan)
{
    return d_base_cas->getDiskUsage(force_scan) + d_pending_write_space;
}

std::optional<int64_t> LruLocalCas::getDiskQuota()
{
    return d_effective_quota_high;
}

void LruLocalCas::verifyDiskUsage()
{
    auto actual_disk_usage = getDiskUsage(true);

    if (getDiskUsage() != actual_disk_usage) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Cached disk usage "
                                    << getDiskUsage()
                                    << " does not match calculated disk usage "
                                    << actual_disk_usage);
    }
}

bool LruLocalCas::externalFileMovesAllowed() const
{
    return d_base_cas->externalFileMovesAllowed();
}

bool LruLocalCas::moveBlobFromExternalFile(const Digest &digest, int dirfd,
                                           const std::string &path)
{
    const auto write_function = [&] {
        return d_base_cas->moveBlobFromExternalFile(digest, dirfd, path);
    };

    return tryToAddBlob(digest.size_bytes(), write_function);
}

int64_t LruLocalCas::getAvailableDiskSpace()
{
    return d_base_cas->getAvailableDiskSpace();
}

void LruLocalCas::updateDiskQuota()
{
    int64_t available_disk_space = getAvailableDiskSpace();

    int64_t max_quota = std::max<int64_t>(
        getDiskUsage() + available_disk_space - d_reserved_space, 0);

    if (d_configuredQuotaHigh.has_value()) {
        d_effective_quota_high =
            std::min<int64_t>(d_configuredQuotaHigh.value(), max_quota);
    }
    else {
        d_effective_quota_high = max_quota;
    }

    GaugeMetricUtil::setGauge(
        CommonMetricNames::GAUGE_NAME_LRU_EFFECTIVE_HIGH_WATERMARK,
        d_effective_quota_high);

    // Effective low watermark cannot be calculated here if high watermark is
    // not configured as it's unknown what amount of disk space is used by
    // files that aren't CAS blobs.
    if (d_configuredQuotaHigh.has_value()) {
        GaugeMetricUtil::setGauge(
            CommonMetricNames::GAUGE_NAME_LRU_EFFECTIVE_LOW_WATERMARK,
            static_cast<int64_t>(static_cast<double>(d_effective_quota_high) *
                                 d_quotaLowRatio));
    }
}

void LruLocalCas::diskMonitorThread()
{
    // Periodically check available disk space
    const auto interval = 5s;

    std::unique_lock<std::mutex> lock(d_stop_mutex);
    while (!d_stop_cv.wait_for(lock, interval, [&] { return d_stop; })) {
        updateDiskQuota();

        if (getDiskUsage() >= d_effective_quota_high) {
            // High watermark reached. Trigger expiry and wait for it to
            // complete.
            try {
                cleanup();
            }
            catch (const OutOfSpaceException &) {
                // Cleanup failed because no more blobs can be deleted due to
                // `d_protect_active_blobs` being set.
                // We ignore the error here; the user will be notified of this
                // when the next write throws.
                ;
            }
        }
    }
}

void LruLocalCas::cleanup()
{
    std::lock_guard<std::mutex> lock(d_cleanup_mutex);

    int64_t quota_high = d_effective_quota_high;

    if (getDiskUsage() < quota_high) {
        // Another thread beat us to it, don't trigger another cleanup
        return;
    }

    BUILDBOX_LOG_INFO("Starting LRU cleanup. Disk usage="
                      << getDiskUsage() << " >= quota=" << quota_high);

    using time_digest_pair = std::pair<blob_time_type, Digest>;

    auto pq_compare = [](time_digest_pair &a, time_digest_pair &b) {
        return a.first > b.first;
    };

    std::priority_queue<time_digest_pair, std::vector<time_digest_pair>,
                        decltype(pq_compare)>
        pq(pq_compare);

    {
        buildboxcommon::buildboxcommonmetrics::MetricGuard<
            buildboxcommon::buildboxcommonmetrics::DurationMetricTimer>
            mt(CommonMetricNames::TIMER_NAME_LRU_CLEANUP);

        int64_t blobDiskUsage = d_pending_write_space;

        listBlobs([&pq, &blobDiskUsage](const buildboxcommon::Digest &digest,
                                        buildboxcommon::blob_time_type time) {
            pq.emplace(time, digest);
            blobDiskUsage += digest.size_bytes();
        });

        // Disk usage of files that aren't cache blobs
        // (should be zero if an absolute quota is set and thus, live blob disk
        // usage tracking is enabled)
        int64_t otherDiskUsage = getDiskUsage() - blobDiskUsage;

        if (!d_configuredQuotaHigh.has_value()) {
            BUILDBOX_LOG_INFO("Disk usage of blobs: " << blobDiskUsage
                                                      << ", other disk usage: "
                                                      << otherDiskUsage);
        }

        int64_t quotaLow =
            otherDiskUsage +
            static_cast<int64_t>(
                static_cast<double>(quota_high - otherDiskUsage) *
                d_quotaLowRatio);

        BUILDBOX_LOG_INFO("Low watermark: " << quotaLow << " ("
                                            << (d_quotaLowRatio * 100)
                                            << "%)");
        GaugeMetricUtil::setGauge(
            CommonMetricNames::GAUGE_NAME_LRU_EFFECTIVE_LOW_WATERMARK,
            quotaLow);

        while (getDiskUsage() > quotaLow && !pq.empty()) {
            const auto blob_time = pq.top().first;

            if (d_protect_active_blobs) {
                if (blob_time >= d_protect_active_time) {
                    // Blob was created or used after the start point of active
                    // blob protection. Do not delete blob.

                    if (getDiskUsage() >= quota_high) {
                        // Disk usage still above high watermark. This means
                        // that the working set is larger than the quota and we
                        // can't continue.

                        BUILDBOX_LOG_ERROR(
                            "Cleanup operation failed: disk usage ("
                            << getDiskUsage() << ") is above maximum quota ("
                            << quota_high
                            << ") and no inactive blobs are "
                               "available for deletion");

                        BUILDBOXCOMMON_THROW_EXCEPTION(
                            OutOfSpaceException, "Insufficient storage quota");
                    }

                    break;
                }
            }

            // On shutdown, cleanup only as much as needed to complete
            // pending operations.
            if (d_cancel_cleanup && getDiskUsage() < quota_high) {
                BUILDBOX_LOG_INFO(
                    "Cleanup operation cancelled due to shutdown.");
                break;
            }

            auto digest = pq.top().second;
            pq.pop();

            try {
                deleteBlob(digest, blob_time);
            }
            catch (const std::runtime_error &e) {
                BUILDBOX_LOG_ERROR("Error deleting blob "
                                   << digest.hash() << "/"
                                   << digest.size_bytes() << ": " << e.what());
            }
        }
    }

    BUILDBOX_LOG_INFO("Finished LRU cleanup. Disk usage: " << getDiskUsage());
}

LocalCas *LruLocalCas::getBase() { return d_base_cas.get(); }

bool LruLocalCas::tryToAddBlob(google::protobuf::int64 blob_size,
                               const std::function<bool()> &add_blob_function)
{
    if (blob_size >= d_effective_quota_high) {
        // Blob is larger than quota. Do not attempt cleanup and bail out.
        BUILDBOXCOMMON_THROW_EXCEPTION(
            OutOfSpaceException,
            "Blob size exceeds total storage allocation ("
                << blob_size << " >= " << d_effective_quota_high << " bytes)");
    }

    // Reserve disk usage before actual write to ensure expiry is triggered
    // early enough in case of parallel write requests.
    d_pending_write_space += blob_size;

    try {
        if (getDiskUsage() >= d_effective_quota_high) {
            // High watermark reached. Trigger expiry and wait for it to
            // complete.
            cleanup();
        }
        const bool blob_was_written = add_blob_function();
        d_pending_write_space -= blob_size;
        return blob_was_written;
    }
    catch (...) {
        d_pending_write_space -= blob_size;
        throw;
    }
}

bool LruLocalCas::objectFilesWritable() const
{
    return d_base_cas->objectFilesWritable();
};
