/*
 * 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 <buildboxcasd_localcasinstance.h>
#include <buildboxcasd_localcasproxyinstance.h>
#include <buildboxcasd_metricnames.h>
#include <buildboxcasd_requestcontextmanager.h>

#include <buildboxcommon_casclient.h>
#include <buildboxcommon_exception.h>
#include <buildboxcommon_fileutils.h>
#include <buildboxcommon_grpcerror.h>
#include <buildboxcommon_logging.h>
#include <buildboxcommonmetrics_countingmetricutil.h>
#include <buildboxcommonmetrics_distributionmetricutil.h>
#include <buildboxcommonmetrics_durationmetrictimer.h>
#include <buildboxcommonmetrics_metricguard.h>

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <fcntl.h>
#include <iterator>
#include <memory>
#include <optional>
#include <sstream>
#include <sys/stat.h>
#include <sys/types.h>
#include <unordered_set>
#include <utility>
#include <vector>

using namespace buildboxcasd;
using namespace buildboxcommon;

LocalCasProxyInstance::LocalCasProxyInstance(
    const std::shared_ptr<LocalCas> &storage,
    const std::unordered_map<buildboxcommon::StageTreeRequest::StagingMode,
                             std::shared_ptr<buildboxcommon::FileStager>>
        &fileStagers,
    const std::shared_ptr<ThreadPool> &digestThreadPool,
    const std::shared_ptr<ThreadPool> &ioThreadPool,
    const std::shared_ptr<std::unordered_set<std::string>>
        &allowedPreUnstageCommands,
    const buildboxcommon::ConnectionOptions &cas_endpoint,
    const std::string &instance_name, const bool read_only_remote,
    const int findmissingblobs_cache_ttl)
    : LocalCasInstance(storage, fileStagers, digestThreadPool,
                       allowedPreUnstageCommands, instance_name),
      d_grpc_client(std::make_shared<GrpcClient>()),
      d_cas_client(std::make_shared<CASClient>(d_grpc_client)),
      d_read_only_remote(read_only_remote),
      d_find_missing_blobs_cache(
          findmissingblobs_cache_ttl > 0
              ? std::make_shared<DigestCache>(findmissingblobs_cache_ttl)
              : nullptr),
      d_ioThreadPool(ioThreadPool),
      d_fetch_tree_operation_queue(std::make_shared<DigestOperationQueue>())
{
    d_grpc_client->init(cas_endpoint);
    d_cas_client->init();

    RequestContextManager::configureGrpcClient(d_grpc_client.get());

    d_find_missing_blobs_client = std::make_shared<FindMissingBlobsClient>(
        *d_cas_client, d_find_missing_blobs_cache.get());

    if (d_find_missing_blobs_cache) {
        BUILDBOX_LOG_INFO("Caching FindMissingBlobs() results for "
                          << findmissingblobs_cache_ttl << " seconds");
    }
}

LocalCasProxyInstance::LocalCasProxyInstance(
    const LocalCasProxyInstance &other, const std::string &instanceName,
    const std::string &root)
    : LocalCasInstance(other, instanceName, root),
      d_grpc_client(other.d_grpc_client), d_cas_client(other.d_cas_client),
      d_read_only_remote(other.d_read_only_remote),
      d_find_missing_blobs_cache(other.d_find_missing_blobs_cache),
      d_find_missing_blobs_client(other.d_find_missing_blobs_client),
      d_ioThreadPool(other.d_ioThreadPool),
      d_fetch_tree_operation_queue(other.d_fetch_tree_operation_queue)
{
}

std::shared_ptr<CasInstance>
LocalCasProxyInstance::clone(const std::string &instanceName,
                             const std::string &root)
{
    return std::make_shared<LocalCasProxyInstance>(*this, instanceName, root);
}

bool LocalCasProxyInstance::hasRemote() const { return true; }

grpc::Status
LocalCasProxyInstance::FindMissingBlobs(const FindMissingBlobsRequest &request,
                                        FindMissingBlobsResponse *response)
{
    BUILDBOX_LOG_INFO(
        "LocalCasProxyInstance::FindMissingBlobs request for instance name \""
        << request.instance_name() << "\" for "
        << request.blob_digests().size() << " digest(s)");

    if (d_read_only_remote) {
        FindMissingBlobsResponse local_response;
        grpc::Status localStatus =
            LocalCasInstance::FindMissingBlobs(request, &local_response);
        if (!localStatus.ok()) {
            BUILDBOX_LOG_ERROR("Local FindMissingBlobs FAILED "
                               << localStatus.error_code() << " "
                               << localStatus.error_message());
            return localStatus;
        }
        const std::vector<Digest> missing_from_local(
            local_response.missing_blob_digests().cbegin(),
            local_response.missing_blob_digests().cend());
        std::vector<Digest> digests_missing_in_remote;
        auto find_missing_blobs_status = findMissingBlobsInRemoteCas(
            missing_from_local, &digests_missing_in_remote);
        if (!find_missing_blobs_status.ok()) {
            return find_missing_blobs_status;
        }
        for (const Digest &d : digests_missing_in_remote) {
            response->add_missing_blob_digests()->CopyFrom(d);
        }
    }
    else {
        const std::vector<Digest> requested_digests(
            request.blob_digests().cbegin(), request.blob_digests().cend());

        std::vector<Digest> digests_missing_in_remote;
        auto find_missing_blobs_status = findMissingBlobsInRemoteCas(
            requested_digests, &digests_missing_in_remote);
        if (!find_missing_blobs_status.ok()) {
            return find_missing_blobs_status;
        }

        std::vector<buildboxcommon::CASClient::UploadRequest> upload_requests;
        CounterType bytes_uploaded = 0;
        CounterType blobs_uploaded = 0;
        for (const Digest &digest : digests_missing_in_remote) {
            if (d_storage->hasBlob(digest)) {
                // If we have the blob stored locally, we implicitly update the
                // remote.
                std::string path = d_storage->pathUnchecked(digest);
                upload_requests.emplace_back(
                    buildboxcommon::CASClient::UploadRequest::from_path(digest,
                                                                        path));
                bytes_uploaded += digest.size_bytes();
                blobs_uploaded++;
            }
            else {
                Digest *entry = response->add_missing_blob_digests();
                entry->CopyFrom(digest);
            }
        }

        const std::vector<CASClient::UploadResult> not_uploaded_blobs =
            d_cas_client->uploadBlobs(upload_requests);

        for (const buildboxcommon::CASClient::UploadResult &upload_result :
             not_uploaded_blobs) {
            BUILDBOX_LOG_ERROR("Error uploading "
                               << toString(upload_result.digest) << ": "
                               << upload_result.status.error_code());

            Digest *entry = response->add_missing_blob_digests();
            entry->CopyFrom(upload_result.digest);

            // Unable to upload this blob for some reason, subtract it from
            // the uploaded bytes count
            bytes_uploaded -= upload_result.digest.size_bytes();
            blobs_uploaded--;
        }

        buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
            recordCounterMetric(MetricNames::COUNTER_NUM_BLOBS_READ_FROM_LOCAL,
                                static_cast<long>(upload_requests.size()));
        buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
            recordCounterMetric(MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
                                bytes_uploaded);
        buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
            recordCounterMetric(
                MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
                blobs_uploaded);
    }
    return grpc::Status::OK;
}

grpc::Status
LocalCasProxyInstance::BatchUpdateBlobs(const BatchUpdateBlobsRequest &request,
                                        BatchUpdateBlobsResponse *response)
{
    BUILDBOX_LOG_INFO(
        "LocalCasProxyInstance::BatchUpdateBlobs request for instance name \""
        << request.instance_name() << "\" for " << request.requests().size()
        << " blob(s)");

    if (d_read_only_remote) {
        return LocalCasInstance::BatchUpdateBlobs(request, response);
    }
    else {
        for (const auto &blob : request.requests()) {
            writeToLocalStorage(blob.digest(), blob.data());
        }

        std::vector<buildboxcommon::CASClient::UploadRequest> upload_requests;
        upload_requests.reserve(static_cast<size_t>(request.requests_size()));
        for (const auto &blob : request.requests()) {
            upload_requests.emplace_back(blob.digest(), blob.data());
        }

        // We assume that the client called `FindMissingBlobs()` previous to
        // this call, so we will upload all the blobs in the request.
        const std::unordered_map<Digest, grpc::Status> digests_not_uploaded =
            batchUpdateRemoteCas(upload_requests);

        for (const auto &blob : request.requests()) {
            google::rpc::Status status;

            const auto failed_upload_it =
                digests_not_uploaded.find(blob.digest());
            if (failed_upload_it != digests_not_uploaded.cend()) {
                status.set_code(failed_upload_it->second.error_code());
                status.set_message(failed_upload_it->second.error_message());
            }
            else {
                status.set_code(grpc::StatusCode::OK);

                // We assume that the remote will keep the blob for a while, so
                // we'll avoid querying for it in subsequent
                // `FindMissingBlobs()`:
                addToFindMissingBlobsCache(blob.digest());
            }

            auto entry = response->add_responses();
            entry->mutable_digest()->CopyFrom(blob.digest());
            entry->mutable_status()->CopyFrom(status);
        }

        return grpc::Status::OK;
    }
}

grpc::Status
LocalCasProxyInstance::BatchReadBlobs(const BatchReadBlobsRequest &request,
                                      BatchReadBlobsResponse *response)
{
    BUILDBOX_LOG_INFO(
        "LocalCasProxyInstance::BatchReadBlobs request for instance name \""
        << request.instance_name() << "\" for " << request.digests().size()
        << " digest(s)");

    // Checking the local storage first:
    std::vector<Digest> digests_missing_locally;
    for (const Digest &digest : request.digests()) {
        std::string data;
        const google::rpc::Status read_status =
            readFromLocalStorage(digest, &data);

        if (read_status.code() == grpc::StatusCode::OK) {
            auto entry = response->add_responses();
            entry->mutable_digest()->CopyFrom(digest);
            entry->mutable_status()->CopyFrom(read_status);
            entry->set_data(data);
        }
        else {
            digests_missing_locally.push_back(digest);
        }
    }

    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(MetricNames::COUNTER_NUM_BLOBS_BATCH_READ,
                            request.digests().size());

    if (digests_missing_locally.empty()) {
        BUILDBOX_LOG_INFO("retrieved all digest(s) for instance name \""
                          << request.instance_name()
                          << "\" from local storage");
        return grpc::Status::OK;
    }

    // Making a request for the blobs that we couldn't find locally:
    const buildboxcommon::CASClient::DownloadBlobsResult downloaded_data =
        d_cas_client->downloadBlobs(digests_missing_locally, nullptr);

    size_t count = 0;
    CounterType bytes_downloaded = 0;
    for (const auto &digest : digests_missing_locally) {
        const auto &download_entry = downloaded_data.at(digest.hash());
        const auto &download_status = download_entry.first;

        auto entry = response->add_responses();
        entry->mutable_digest()->CopyFrom(digest);
        entry->mutable_status()->CopyFrom(download_status);

        if (download_status.code() == grpc::StatusCode::OK) {
            const auto &data = download_entry.second;
            entry->set_data(data);

            buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
                MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE, 1);

            writeToLocalStorage(digest, data);
            bytes_downloaded += digest.size_bytes();
            ++count;
        }
    }

    if (bytes_downloaded > 0) {
        buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
            recordCounterMetric(MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
                                bytes_downloaded);
    }
    BUILDBOX_LOG_INFO("downloaded "
                      << count << " out of " << digests_missing_locally.size()
                      << " digest(s) from the remote cas server");

    return grpc::Status::OK;
}

grpc::Status LocalCasProxyInstance::Write(WriteRequest *request_message,
                                          ServerReader<WriteRequest> &request,
                                          WriteResponse *response,
                                          Digest *requested_digest)
{
    if (d_read_only_remote) {
        return LocalCasInstance::Write(request_message, request, response,
                                       requested_digest);
    }
    response->set_committed_size(0);

    auto buffer_file = d_storage->createTemporaryFile();

    Digest digest_to_write;
    auto request_status = CasInstance::processWriteRequest(
        request_message, request, &digest_to_write, buffer_file.name());

    if (request_status.error_code() != grpc::INVALID_ARGUMENT) {
        *requested_digest = digest_to_write;
    }

    if (!request_status.ok()) {
        if (request_status.error_code() == grpc::StatusCode::ALREADY_EXISTS) {
            // The blob is already present in the CAS. In that case, according
            // to the REAPI spec:
            // " [...] if another client has already completed the upload
            // [...], the request will terminate immediately with a response
            // whose `committed_size` is the full size of the uploaded file
            // (regardless of how much data was transmitted by the client)"
            response->set_committed_size(digest_to_write.size_bytes());
            return grpc::Status::OK;
        }

        return request_status;
    }

    // Trying to upload the blob first:
    const auto upload_status = uploadBlob(digest_to_write, buffer_file.name());

    // Upload was successful. We can move the file directly to the CAS,
    // avoiding copies:
    auto move_status = moveTemporaryFileToLocalStorage(digest_to_write,
                                                       std::move(buffer_file));

    // We assume that the remote will keep the blob for a while, so
    // we'll avoid querying for it in subsequent `FindMissingBlobs()`:
    addToFindMissingBlobsCache(digest_to_write);

    if (move_status.ok()) {
        response->set_committed_size(digest_to_write.size_bytes());
    }

    return move_status;
}

bool LocalCasProxyInstance::hasBlob(const Digest &digest)
{
    return d_storage->hasBlob(digest);
}

google::rpc::Status LocalCasProxyInstance::readBlob(const Digest &digest,
                                                    std::string *data,
                                                    size_t read_offset,
                                                    size_t read_limit)
{
    google::rpc::Status read_status =
        readFromLocalStorage(digest, data, read_offset, read_limit);
    if (read_status.code() == grpc::StatusCode::OK) {
        return read_status;
    }

    google::rpc::Status status;
    try {
        const auto diskQuota = d_storage->getDiskQuota();
        if (diskQuota.has_value() && digest.size_bytes() > diskQuota.value()) {
            // Completely fall through as the blob cannot be cached
            data->reserve(read_limit);
            const auto readLambda = [&data](const ReadResponse &response,
                                            size_t) {
                *data += response.data();
                return grpc::Status::OK;
            };
            d_cas_client->fetchStream(digest, readLambda,
                                      static_cast<int64_t>(read_offset),
                                      static_cast<int64_t>(read_limit));
            buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
                recordCounterMetric(
                    MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
                    static_cast<int64_t>(data->size()));
        }
        else {
            auto tmp = d_storage->createTemporaryFile();
            d_cas_client->download(tmp.fd(), digest);

            // While less bytes might be served to clients due to the
            // read_limit, the entire blob is downloaded so we use the full
            // size here
            buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
                recordCounterMetric(
                    MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
                    digest.size_bytes());

            *data = FileUtils::getFileContents(
                tmp.fd(), static_cast<int64_t>(read_offset),
                read_limit == 0 ? std::nullopt
                                : std::make_optional(read_limit));

            const auto moveStatus =
                moveTemporaryFileToLocalStorage(digest, std::move(tmp));
            if (!moveStatus.ok()) {
                BUILDBOX_LOG_WARNING(
                    "Error moving temporary file to local storage, digest="
                    << digest << " code=" << moveStatus.error_code()
                    << " message=" << moveStatus.error_message());
            }
        }

        // For each blob we successfully downloaded from the remote,
        // increment the count.
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE, 1);
    }
    catch (const GrpcError &e) {
        BUILDBOX_LOG_ERROR("Error downloading "
                           << digest.hash()
                           << " code=" << e.status.error_code()
                           << " message=" << e.status.error_message());
        status.set_code(e.status.error_code());
        status.set_message(e.status.error_message());
    }
    catch (const std::runtime_error &) {
        status.set_code(grpc::StatusCode::NOT_FOUND);
        status.set_message("Blob not found in local nor remote CAS");
    }

    return status;
}

google::rpc::Status LocalCasProxyInstance::writeBlob(const Digest &digest,
                                                     const std::string &data)
{
    BUILDBOX_LOG_INFO("LocalCasProxyInstance::writeBlob: writing digest of "
                      << digest.size_bytes() << " bytes");

    google::rpc::Status status;

    // Trying to upload the blob first:
    try {
        d_cas_client->upload(data, digest);
    }
    catch (const std::logic_error &) { // Digest does not match the data.
        status.set_code(grpc::StatusCode::INVALID_ARGUMENT);
        return status;
    }
    catch (const std::runtime_error &) {
        status.set_code(grpc::StatusCode::UNAVAILABLE);
        return status;
    }
    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
                            digest.size_bytes());
    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
                            1);

    // Upload was successful. We now write to the local storage
    // (if needed, LocalCas will skip the write if the file exists.):
    google::rpc::Status write_status = writeToLocalStorage(digest, data);
    status.set_code(write_status.code());

    return status;
}

grpc::Status LocalCasProxyInstance::uploadBlob(const Digest &digest,
                                               const std::string &path) const
{
    const FileDescriptor fd(open(path.c_str(), O_RDONLY | O_CLOEXEC));
    if (fd.get() == -1) {
        std::ostringstream error_message;
        error_message << "Could not read intermediate file at \"" << path
                      << "\":" << strerror(errno);
        BUILDBOX_LOG_ERROR(error_message.str());
        return grpc::Status(grpc::StatusCode::INTERNAL, error_message.str());
    }

    grpc::Status upload_status;
    try {
        d_cas_client->upload(fd.get(), digest);
        upload_status = grpc::Status::OK;

        buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
            recordCounterMetric(MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
                                digest.size_bytes());
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE, 1);
    }
    catch (const std::logic_error &e) { // Digest does not match the data.
        upload_status =
            grpc::Status(grpc::StatusCode::INVALID_ARGUMENT, e.what());
    }
    catch (const std::runtime_error &e) {
        upload_status = grpc::Status(grpc::StatusCode::UNAVAILABLE, e.what());
    }

    return upload_status;
}

void LocalCasProxyInstance::addToFindMissingBlobsCache(const Digest &digest)
{
    if (d_find_missing_blobs_cache) {
        d_find_missing_blobs_cache->addDigest(digest);
    }
}

std::unordered_map<Digest, grpc::Status>
LocalCasProxyInstance::batchUpdateRemoteCas(
    std::vector<buildboxcommon::CASClient::UploadRequest> &upload_requests)
{
    CounterType bytes_uploaded = 0;
    for (const auto &blob : upload_requests) {
        bytes_uploaded += blob.digest.size_bytes();
    }

    const std::vector<buildboxcommon::CASClient::UploadResult>
        blobs_not_uploaded = d_cas_client->uploadBlobs(upload_requests);

    std::unordered_map<Digest, grpc::Status> res;
    for (const buildboxcommon::CASClient::UploadResult &upload_result :
         blobs_not_uploaded) {
        res.emplace(upload_result.digest, upload_result.status);
        bytes_uploaded -= upload_result.digest.size_bytes();
    }

    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
                            bytes_uploaded);

    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
                            static_cast<long>(upload_requests.size() -
                                              blobs_not_uploaded.size()));

    return res;
}

grpc::Status LocalCasProxyInstance::findMissingBlobsInRemoteCas(
    const std::vector<Digest> &digests, std::vector<Digest> *missing_digests)
{
    return d_find_missing_blobs_client->findMissingBlobs(digests,
                                                         missing_digests);
}

std::unordered_map<Digest, google::rpc::Status>
LocalCasProxyInstance::fetchBlobs(const std::vector<Digest> &digests) const
{
    CounterType bytes_downloaded = 0;
    std::unordered_map<Digest, google::rpc::Status> res;

    if (digests.empty()) {
        return res;
    }

    auto temp_dir = d_storage->createTemporaryDirectory();

    // Making a request for the blobs that we couldn't find locally:
    auto downloaded_data = d_cas_client->downloadBlobsToDirectory(
        digests, temp_dir.fd(), d_ioThreadPool.get());

    for (const Digest &digest : digests) {
        const auto &digest_path_it = downloaded_data.at(digest.hash());
        const google::rpc::Status download_status = digest_path_it.first;

        // Download failed. We add that to the reponse, which only contains
        // digests that failed:
        if (download_status.code() != grpc::StatusCode::OK) {
            res.emplace(digest, download_status);
            continue;
        }

        // Download suceeded. Add the blob to CAS:
        bytes_downloaded += digest.size_bytes();

        try {
            if (!this->d_storage->hasBlob(digest)) {
                d_storage->moveBlobFromTemporaryFile(digest, temp_dir,
                                                     digest_path_it.second);
            }
            // For each blob we successfully downloaded from the remote,
            // increment the count.
            buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
                MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE, 1);
        }
        catch (const std::runtime_error &e) {
            // We did not add the blob to the CAS. Report that as an error in
            // the response:
            google::rpc::Status status;
            status.set_code(grpc::StatusCode::INTERNAL);
            status.set_message(
                "Internal error while writing blob to local CAS: " +
                std::string(e.what()));
            res.emplace(digest, status);
        }
    }

    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
                            bytes_downloaded);

    return res;
}

Status LocalCasProxyInstance::FetchMissingBlobs(
    const FetchMissingBlobsRequest &request,
    FetchMissingBlobsResponse *response)
{
    std::unordered_set<Digest> digests_missing_locally;
    // (Repeated digests in the request will be treated as a single entry.)
    for (const Digest &digest : request.blob_digests()) {
        // Checking the local storage first:
        if (!d_storage->hasBlob(digest)) {
            digests_missing_locally.insert(digest);
        }
    }

    std::unordered_map<Digest, google::rpc::Status> fetch_result;
    try {
        fetch_result = fetchBlobs(std::vector<Digest>(
            digests_missing_locally.cbegin(), digests_missing_locally.cend()));
    }
    catch (const std::runtime_error &e) {
        const std::string error_message =
            "Error downloading blobs from remote: " + std::string(e.what());
        BUILDBOX_LOG_ERROR(error_message);
        return grpc::Status(grpc::StatusCode::INTERNAL, error_message);
    }

    // fetch_result contains failed downloads, add them to the response.
    for (const auto &digest_status_it : fetch_result) {
        const auto digest = digest_status_it.first;
        const google::rpc::Status download_status = digest_status_it.second;

        auto entry = response->add_responses();
        entry->mutable_digest()->CopyFrom(digest);
        entry->mutable_status()->CopyFrom(download_status);
    }

    return grpc::Status::OK;
}

Status LocalCasProxyInstance::UploadMissingBlobs(
    const UploadMissingBlobsRequest &request,
    UploadMissingBlobsResponse *response)
{
    if (d_read_only_remote) {
        return LocalCasInstance::UploadMissingBlobs(request, response);
    }
    // Construct std::vector from protobuf digest list
    auto pb_digests = request.blob_digests();
    std::vector<Digest> digests(pb_digests.begin(), pb_digests.end());

    std::vector<Digest> missing_digests;
    auto find_missing_blobs_status =
        findMissingBlobsInRemoteCas(digests, &missing_digests);
    if (!find_missing_blobs_status.ok()) {
        return find_missing_blobs_status;
    }

    std::vector<buildboxcommon::CASClient::UploadRequest> upload_requests;
    for (const Digest &digest : missing_digests) {
        google::rpc::Status status;
        status.set_code(grpc::StatusCode::OK);

        try {
            std::string path = d_storage->path(digest);
            upload_requests.emplace_back(
                buildboxcommon::CASClient::UploadRequest::from_path(digest,
                                                                    path));
        }
        catch (const BlobNotFoundException &) {
            status.set_code(grpc::StatusCode::NOT_FOUND);
            status.set_message("Blob not found in local cache");
        }
        catch (const std::runtime_error &e) {
            status.set_code(grpc::StatusCode::INTERNAL);
            status.set_message(e.what());
        }

        if (status.code() != grpc::StatusCode::OK) {
            auto entry = response->add_responses();
            entry->mutable_digest()->CopyFrom(digest);
            entry->mutable_status()->CopyFrom(status);
        }
    }

    const std::unordered_map<Digest, grpc::Status> digests_not_uploaded =
        batchUpdateRemoteCas(upload_requests);

    for (const auto &blob : upload_requests) {
        google::rpc::Status status;

        const auto failed_upload_it = digests_not_uploaded.find(blob.digest);
        if (failed_upload_it != digests_not_uploaded.cend()) {
            status.set_code(failed_upload_it->second.error_code());
            status.set_message(failed_upload_it->second.error_message());
        }
        else {
            status.set_code(grpc::StatusCode::OK);

            // We assume that the remote will keep the blob for a while, so
            // we'll avoid querying for it in subsequent
            // `FindMissingBlobs()`:
            addToFindMissingBlobsCache(blob.digest);
        }

        auto entry = response->add_responses();
        entry->mutable_digest()->CopyFrom(blob.digest);
        entry->mutable_status()->CopyFrom(status);
    }

    return grpc::Status::OK;
}

Status
LocalCasProxyInstance::prepareTreeForStaging(const Digest &root_digest) const
{
    // Proxy mode. (We can try and fetch the blobs that might be missing.)
    grpc::Status status;
    try {
        buildboxcommon::buildboxcommonmetrics::MetricGuard<
            buildboxcommon::buildboxcommonmetrics::DurationMetricTimer>
            mt(MetricNames::TIMER_NAME_LOCAL_CAS_STAGE_TREE_PREPARE);

        status = fetchTreeMissingBlobs(root_digest);
    }
    catch (const GrpcError &e) {
        status = e.status;
    }
    catch (const std::runtime_error &e) {
        status = Status(grpc::StatusCode::INTERNAL, e.what());
    }

    if (status.error_code() == grpc::StatusCode::NOT_FOUND) {
        const std::string errorMessage =
            "At least one blob is missing locally and remotely for "
            "staging directory with root digest \"" +
            toString(root_digest) + "\": " + status.error_message();
        BUILDBOX_LOG_ERROR(errorMessage);
        status = Status(grpc::StatusCode::FAILED_PRECONDITION, errorMessage);
    }
    else if (!status.ok()) {
        const std::string errorMessage = "Error while fetching blobs that "
                                         "are missing locally for staging "
                                         "directory with root digest \"" +
                                         toString(root_digest) +
                                         "\": " + status.error_message();
        BUILDBOX_LOG_ERROR(errorMessage);
        status = Status(status.error_code(), errorMessage);
    }
    return status;
}

Status LocalCasProxyInstance::fetchTreeMissingBlobs(const Digest &root_digest,
                                                    bool file_blobs) const
{
    auto tree_cache = this->getTreeCache();
    int64_t tree_cache_hits = 0;
    int64_t tree_cache_misses = 0;

    // Digests of blobs that are pending to be fetched
    std::unordered_set<Digest> digests_missing_locally;
    // Digests of directories that are in the local cache but not yet processed
    std::deque<Digest> directory_queue;
    // Digests of directories that need to be fetched before processing
    std::deque<Digest> deferred_directory_queue;
    // Digests of directories that have been fetched or checked to already be
    // in the local cache
    std::vector<Digest> completed_directories;

    directory_queue.push_back(root_digest);

    while (!directory_queue.empty()) {
        const auto directory_digest = directory_queue[0];
        directory_queue.pop_front();

        if (tree_cache->hasRootDigest(directory_digest, file_blobs)) {
            tree_cache_hits++;
        }
        else {
            const Directory directory = fetchDirectory(directory_digest);
            // (This `fetchDirectory()` throws if it fails)

            const std::vector<Digest> missing_directory_digests =
                digestsMissingFromDirectory(directory, file_blobs);
            digests_missing_locally.insert(missing_directory_digests.cbegin(),
                                           missing_directory_digests.cend());

            // Add subdirectories to processing queue
            for (const DirectoryNode &subdir : directory.directories()) {
                if (digests_missing_locally.count(subdir.digest()) == 0) {
                    // Subdirectory is already in local cache
                    directory_queue.push_back(subdir.digest());
                }
                else {
                    // Subdirectory is missing locally
                    deferred_directory_queue.push_back(subdir.digest());
                }
            }

            // Don't add it directly to the directory cache
            // as subdirectories may not be cached yet.
            completed_directories.push_back(directory_digest);
            tree_cache_misses++;
        }

        if (directory_queue.empty()) {
            // No directory left that can be processed without fetching.
            // Fetch the missing blobs.
            const auto fetch_result = fetchBlobs(
                std::vector<Digest>(digests_missing_locally.cbegin(),
                                    digests_missing_locally.cend()));

            for (const auto &digest_status_it : fetch_result) {
                const auto digest = digest_status_it.first;
                const google::rpc::Status &status = digest_status_it.second;

                if (status.code() != grpc::StatusCode::OK) {
                    // If a single blob is missing we can't stage the
                    // directory, aborting:
                    return Status(
                        static_cast<grpc::StatusCode>(status.code()),
                        "Could not fetch missing blob with digest \"" +
                            toString(digest) + "\": " + ": " +
                            status.message());
                }
            }

            // The batch has been fetched successfully and thus,
            // these blobs are no longer missing locally.
            digests_missing_locally.clear();

            // Swap the contents of the two queues.
            directory_queue.swap(deferred_directory_queue);
        }
    }

    for (const auto &directory_digest : completed_directories) {
        tree_cache->addRootDigest(directory_digest, file_blobs);
    }

    recordTreeCacheMetrics(tree_cache_hits, tree_cache_misses);

    return grpc::Status::OK;
}

Directory
LocalCasProxyInstance::fetchDirectory(const Digest &root_digest) const
{
    Directory directory;
    if (this->d_storage->hasBlob(root_digest) &&
        directory.ParseFromString(*d_storage->readBlob(root_digest))) {
        return directory;
    }

    const std::string directory_blob = d_cas_client->fetchString(root_digest);
    if (!directory.ParseFromString(directory_blob)) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error, "Error parsing fetched Directory with digest "
                                    << toString(root_digest));
    }
    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(MetricNames::COUNTER_NAME_REMOTE_BYTES_READ,
                            root_digest.size_bytes());

    d_storage->writeBlob(root_digest, directory.SerializeAsString());
    // For each blob we successfully downloaded from the remote,
    // increment the count.
    buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
        MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE, 1);

    return directory;
}

grpc::Status
LocalCasProxyInstance::GetTree(const GetTreeRequest &request,
                               ServerWriter<GetTreeResponse> *writer)
{
    const Digest &rootDigest = request.root_digest();

    // GetTree may be forwarded to the remote or handled locally.
    // Define a prefix for forwarded page tokens such that subsequent requests
    // are handled by the same CAS implementation.
    const std::string REMOTE_TOKEN_PREFIX = "#";

    if (!request.page_token().empty() &&
        !request.page_token().starts_with(REMOTE_TOKEN_PREFIX)) {
        // Continuation of previous request that was handled locally.
        return CasInstance::GetTree(request, writer);
    }

    if (request.page_token().empty() && this->d_storage->hasBlob(rootDigest)) {
        // First request and root Directory is in local cache, use local
        // implementation.
        return CasInstance::GetTree(request, writer);
    }

    std::string pageToken;
    if (request.page_token().starts_with(REMOTE_TOKEN_PREFIX)) {
        // Continuation of previous request that was forwarded to the remote,
        // strip prefix.
        pageToken = request.page_token().substr(REMOTE_TOKEN_PREFIX.size());
    }

    grpc::ClientContext context;
    auto reader = d_cas_client->getTree(&context, request.root_digest(),
                                        request.page_size(), pageToken);

    GetTreeResponse response;
    while (reader->Read(&response)) {
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NUM_BLOBS_READ_FROM_REMOTE,
            response.directories_size());

        // Cache Directory objects locally
        for (const auto &directory : response.directories()) {
            const std::string serialized = directory.SerializeAsString();
            const auto digest =
                buildboxcommon::DigestGenerator::hash(serialized);
            this->d_storage->writeBlob(digest, serialized);
        }

        // Forward response page to client
        if (!response.next_page_token().empty()) {
            response.set_next_page_token(REMOTE_TOKEN_PREFIX +
                                         response.next_page_token());
        }
        if (!writer->Write(response)) {
            // The stream has been closed
            return grpc::Status::CANCELLED;
        }
    }

    return reader->Finish();
}

Digest LocalCasProxyInstance::UploadAndStore(
    int dirfd, const buildboxcommon::digest_string_map &digest_blob_map,
    const buildboxcommon::digest_string_map &digest_path_map,
    const std::optional<Tree> &t, const bool bypass_local_cache,
    const bool move_files_hint)
{
    std::vector<CASClient::UploadRequest> upload_requests;
    // Reserving one extra space for the `Tree`:
    upload_requests.reserve(digest_blob_map.size() + digest_path_map.size() +
                            1);

    CounterType bytes_uploaded = 0;
    for (const auto &it : digest_blob_map) {
        const Digest &digest = it.first;
        const std::string &serialized_directory = it.second;
        if (!bypass_local_cache) {
            d_storage->writeBlob(digest, serialized_directory);
        }

        upload_requests.emplace_back(digest, serialized_directory);
    }
    for (const auto &it : digest_path_map) {
        const Digest &digest = it.first;
        std::string path = it.second;

        if (!bypass_local_cache) {
            addCapturedFileToLocalStorage(digest, -1, dirfd, path,
                                          move_files_hint);
            // The file could have been moved, use its new path:
            path = d_storage->pathUnchecked(digest);
        }

        upload_requests.emplace_back(
            CASClient::UploadRequest::from_path(digest, dirfd, path));
        bytes_uploaded += digest.size_bytes();
    }

    // If set, upload and store the `Tree` message as well:
    Digest tree_digest;
    if (t.has_value()) {
        tree_digest = buildboxcommon::DigestGenerator::hash(*t);
        const std::string serialized_tree = t->SerializeAsString();

        if (!bypass_local_cache) {
            d_storage->writeBlob(tree_digest, serialized_tree);
        }

        upload_requests.emplace_back(tree_digest, serialized_tree);
    }

    // Filter requests by FindMissingBlobs
    std::vector<Digest> allDigests, missingDigests;
    allDigests.reserve(upload_requests.size());
    std::transform(
        upload_requests.cbegin(), upload_requests.cend(),
        std::back_inserter(allDigests),
        [](const CASClient::UploadRequest &req) { return req.digest; });
    const auto status =
        findMissingBlobsInRemoteCas(allDigests, &missingDigests);
    if (!status.ok()) {
        BUILDBOXCOMMON_THROW_EXCEPTION(
            std::runtime_error,
            "Unable to call FindMissingBlobs before uploading "
                << allDigests.size()
                << " blobs to remote CAS, code=" << status.error_code()
                << " message=" << status.error_message());
    }
    std::unordered_set<Digest> missingDigestSet(
        std::make_move_iterator(missingDigests.begin()),
        std::make_move_iterator(missingDigests.end()));
    upload_requests.erase(
        std::remove_if(
            upload_requests.begin(), upload_requests.end(),
            [&missingDigestSet](const CASClient::UploadRequest &req) {
                return missingDigestSet.find(req.digest) ==
                       missingDigestSet.end();
            }),
        upload_requests.end());

    // Upload to remote cas:
    const auto blobs_not_uploaded =
        d_cas_client->uploadBlobs(upload_requests, d_ioThreadPool.get());
    for (const auto &upload_result : blobs_not_uploaded) {
        bytes_uploaded -= upload_result.digest.size_bytes();
    }

    // Update metrics:
    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
                            bytes_uploaded);
    buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
        MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
        static_cast<long>(upload_requests.size() - blobs_not_uploaded.size()));

    if (blobs_not_uploaded.size() > 0) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Unable to upload "
                                           << blobs_not_uploaded.size()
                                           << " blobs to remote CAS");
    }

    return tree_digest;
}

void LocalCasProxyInstance::captureFileData(const Digest &digest, int fd,
                                            const std::string &path,
                                            const bool bypass_local_cache,
                                            const bool move_file)
{
    // fd is required, path may be empty
    if (fd < 0) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::invalid_argument,
                                       "captureFileData(): Invalid fd " << fd);
    }

    if (!bypass_local_cache) {
        addCapturedFileToLocalStorage(digest, fd, AT_FDCWD, path, move_file);
    }

    // Upload to remote CAS:
    const auto blobs_not_uploaded = d_cas_client->uploadBlobs(
        {CASClient::UploadRequest::from_fd(digest, fd)});

    // Update metrics:
    buildboxcommon::buildboxcommonmetrics::CountingMetricUtil::
        recordCounterMetric(MetricNames::COUNTER_NAME_REMOTE_BYTES_WRITE,
                            blobs_not_uploaded.empty() ? digest.size_bytes()
                                                       : 0);

    buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
        MetricNames::COUNTER_NUM_BLOBS_WRITTEN_TO_REMOTE,
        static_cast<long>(1 - blobs_not_uploaded.size()));

    if (blobs_not_uploaded.size() > 0) {
        BUILDBOXCOMMON_THROW_EXCEPTION(std::runtime_error,
                                       "Unable to upload "
                                           << blobs_not_uploaded.size()
                                           << " blobs to remote CAS");
    }
}

Status LocalCasProxyInstance::FetchTree(const FetchTreeRequest &request,
                                        FetchTreeResponse *)
{
    buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
        MetricNames::COUNTER_NAME_LOCAL_CAS_NUM_FETCH_TREE_REQUESTS, 1);

    auto tree_cache = this->getTreeCache();

    // We might have seen this request recently. If so, don't check the
    // storage:
    if (tree_cache->hasRootDigest(request.root_digest(),
                                  request.fetch_file_blobs())) {
        buildboxcommonmetrics::CountingMetricUtil::recordCounterMetric(
            MetricNames::COUNTER_NAME_LOCAL_CAS_NUM_FETCH_TREE_CACHE_HITS, 1);
        return grpc::Status::OK;
    }

    grpc::Status status;
    const auto fetchTreeLambda = [this, &request, &status](const Digest &) {
        try {
            const buildboxcommon::Digest &root_digest = request.root_digest();
            const bool fetch_files = request.fetch_file_blobs();

            status = fetchTreeMissingBlobs(root_digest, fetch_files);
        }
        catch (const GrpcError &e) {
            status = e.status;
        }
        catch (const std::runtime_error &e) {
            status = grpc::Status(grpc::StatusCode::INTERNAL, e.what());
        }
        if (!status.ok()) {
            const std::string errorMessage = "Error while fetching blobs for "
                                             "directory with root digest \"" +
                                             toString(request.root_digest()) +
                                             "\": " + status.error_message();
            BUILDBOX_LOG_ERROR(errorMessage);
            status = Status(status.error_code(), errorMessage);
        }
    };

    // Making sure that there is never more than one call to
    // `fetchTreeMissingBlobs(d)` in-flight for a same digest.
    d_fetch_tree_operation_queue->runExclusively(request.root_digest(),
                                                 fetchTreeLambda);

    return status;
}

Status LocalCasProxyInstance::UploadTree(const UploadTreeRequest &request,
                                         UploadTreeResponse *response)
{
    if (d_read_only_remote) {
        return LocalCasInstance::UploadTree(request, response);
    }

    // Collect digests of the directory tree
    std::deque<Digest> directoryQueue;
    std::vector<Digest> digests;
    directoryQueue.push_back(request.root_digest());
    while (!directoryQueue.empty()) {
        const auto directoryDigest = directoryQueue[0];
        directoryQueue.pop_front();

        if (!this->d_storage->hasBlob(directoryDigest)) {
            return grpc::Status(grpc::StatusCode::NOT_FOUND,
                                "Directory " + toString(directoryDigest) +
                                    " not found in local cache");
        }

        digests.push_back(directoryDigest);

        Directory directory;
        if (!directory.ParseFromString(
                *d_storage->readBlob(directoryDigest))) {
            BUILDBOXCOMMON_THROW_EXCEPTION(
                std::runtime_error, "Error parsing Directory with digest "
                                        << toString(directoryDigest));
        }

        for (const FileNode &file : directory.files()) {
            digests.push_back(file.digest());
        }

        // Add subdirectories to the processing queue
        for (const DirectoryNode &subdir : directory.directories()) {
            directoryQueue.push_back(subdir.digest());
        }
    }

    // Determine which blobs are missing in the remote CAS
    std::vector<Digest> missingDigests;
    auto findMissingBlobsStatus =
        findMissingBlobsInRemoteCas(digests, &missingDigests);
    if (!findMissingBlobsStatus.ok()) {
        return findMissingBlobsStatus;
    }

    // Upload missing blobs from local cache to the remote CAS
    std::vector<buildboxcommon::CASClient::UploadRequest> uploadRequests;
    for (const Digest &digest : missingDigests) {
        try {
            std::string path = d_storage->path(digest);
            uploadRequests.emplace_back(
                buildboxcommon::CASClient::UploadRequest::from_path(digest,
                                                                    path));
        }
        catch (const BlobNotFoundException &) {
            return grpc::Status(grpc::StatusCode::NOT_FOUND,
                                "Blob " + toString(digest) +
                                    " not found in local cache");
        }
        catch (const std::runtime_error &e) {
            return grpc::Status(grpc::StatusCode::INTERNAL, e.what());
        }
    }

    const std::unordered_map<Digest, grpc::Status> digestsNotUploaded =
        batchUpdateRemoteCas(uploadRequests);

    // If upload failed, return the first error
    if (!digestsNotUploaded.empty()) {
        return digestsNotUploaded.begin()->second;
    }

    return grpc::Status::OK;
}

Status LocalCasProxyInstance::ExportFiles(const ExportFilesRequest &request,
                                          ExportFilesResponse *response)
{
    std::unordered_set<Digest> digestsMissingLocally;
    for (const auto &file : request.output_files()) {
        if (!d_storage->hasBlob(file.digest())) {
            digestsMissingLocally.insert(file.digest());
        }
    }

    try {
        fetchBlobs(std::vector<Digest>(digestsMissingLocally.cbegin(),
                                       digestsMissingLocally.cend()));
    }
    catch (const std::runtime_error &e) {
        const std::string error_message =
            "Error downloading blobs from remote: " + std::string(e.what());
        BUILDBOX_LOG_ERROR(error_message);
        return grpc::Status(grpc::StatusCode::INTERNAL, error_message);
    }

    return LocalCasInstance::ExportFiles(request, response);
}
