/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::borrow::Cow;
use std::fmt;
use std::hash::Hash;
use std::hash::Hasher;

use anyhow::Context;
use buck2_core::buck_path::path::BuckPathRef;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::fs::buck_out_path::BuckOutPath;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use dupe::Dupe;
use either::Either;
use gazebo::cell::ARef;
use gazebo::eq_chain;

#[derive(Debug)]
pub struct ArtifactPath<'a> {
    pub base_path: Either<ARef<'a, BuckOutPath>, BuckPathRef<'a>>,
    pub projected_path: Option<&'a ForwardRelativePath>,
    /// The number of components at the prefix of that path that are internal details to the rule,
    /// not returned by `.short_path`. Omitted from Eq and Hash comparisons.
    pub hidden_components_count: usize,
}

impl<'a> ArtifactPath<'a> {
    pub fn with_filename<F, T>(&self, f: F) -> T
    where
        for<'b> F: FnOnce(anyhow::Result<&'b FileName>) -> T,
    {
        let file_name = match self.projected_path.as_ref() {
            Some(projected_path) => projected_path,
            None => match self.base_path.as_ref() {
                Either::Left(buck_out) => buck_out.path(),
                Either::Right(buck) => buck.path().as_ref(),
            },
        }
        .file_name()
        .with_context(|| format!("Artifact has no file name: `{}`", self));

        f(file_name)
    }

    pub fn with_short_path<F, T>(&self, f: F) -> T
    where
        for<'b> F: FnOnce(&'b ForwardRelativePath) -> T,
    {
        let base_short_path = match self.base_path.as_ref() {
            Either::Left(buck_out) => buck_out.path(),
            Either::Right(buck) => buck.path().as_ref(),
        };

        let path = match self.projected_path.as_ref() {
            Some(projected_path) => Cow::Owned(base_short_path.join(projected_path)),
            None => Cow::Borrowed(base_short_path),
        };

        let path = match path.strip_prefix_components(self.hidden_components_count) {
            Some(p) => p,
            None => ForwardRelativePath::empty(),
        };

        f(path)
    }

    pub fn with_full_path<F, T>(&self, f: F) -> T
    where
        for<'b> F: FnOnce(&'b ForwardRelativePath) -> T,
    {
        let base_path = match self.base_path.as_ref() {
            Either::Left(buck_out) => Cow::Borrowed(buck_out.path()),
            Either::Right(buck) => Cow::Owned(
                buck.package()
                    .cell_relative_path()
                    .as_forward_relative_path()
                    .join(buck.path()),
            ),
        };

        let path = match self.projected_path.as_ref() {
            Some(projected_path) => Cow::Owned(base_path.join(projected_path)),
            None => base_path,
        };

        f(&path)
    }

    pub fn resolve(&self, artifact_fs: &ArtifactFs) -> anyhow::Result<ProjectRelativePathBuf> {
        let ArtifactPath {
            base_path,
            projected_path,
            hidden_components_count: _,
        } = self;

        let base_path = match base_path {
            Either::Left(build) => artifact_fs.buck_out_path_resolver().resolve_gen(build),
            Either::Right(source) => artifact_fs
                .buck_path_resolver()
                .resolve_buck_path(source.dupe())?,
        };

        Ok(match projected_path {
            Some(projected_path) => base_path.join(projected_path),
            None => base_path,
        })
    }
}

impl Hash for ArtifactPath<'_> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.base_path.hash(state);
        self.projected_path.as_ref().hash(state);
    }
}

impl PartialEq for ArtifactPath<'_> {
    fn eq(&self, other: &Self) -> bool {
        eq_chain! {
            self.base_path == other.base_path,
            self.projected_path.as_ref() == other.projected_path.as_ref()
        }
    }
}

impl Eq for ArtifactPath<'_> {}

impl fmt::Display for ArtifactPath<'_> {
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        // NOTE: This produces a representation we tend to use in Starlark for those, which isn't
        // really consistent with what we use when *not* in Starlark.
        self.with_short_path(|p| write!(fmt, "{}", p))
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

pub mod artifact_path;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use buck2_events::dispatch::get_dispatcher;
use buck2_wrapper_common::invocation_id::TraceId;

use crate::execute::request::CommandExecutionPaths;
use crate::execute::target::CommandExecutionTarget;

pub struct ReActionIdentity<'a> {
    /// This is currently unused, but historically it has been useful to add logging in the RE
    /// client, so it's worth keeping around.
    _target: &'a dyn CommandExecutionTarget,

    /// Actions with the same action key share e.g. memory requirements learnt by RE.
    pub action_key: String,

    /// Actions with the same affinity key get scheduled on similar hosts.
    pub affinity_key: String,

    /// Details about the action collected while uploading
    pub paths: &'a CommandExecutionPaths,

    //// Trace ID which started the execution of this action, to be added on the RE side
    pub trace_id: TraceId,
}

impl<'a> ReActionIdentity<'a> {
    pub fn new(
        target: &'a dyn CommandExecutionTarget,
        executor_action_key: Option<&str>,
        paths: &'a CommandExecutionPaths,
    ) -> Self {
        let mut action_key = target.re_action_key();
        if let Some(executor_action_key) = executor_action_key {
            action_key = format!("{} {}", executor_action_key, action_key);
        }

        let trace_id = get_dispatcher().trace_id().to_owned();

        Self {
            _target: target,
            action_key,
            affinity_key: target.re_affinity_key(),
            paths,
            trace_id,
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::sync::Arc;
use std::time::Duration;

use allocative::Allocative;
use anyhow::Context;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_re_configuration::RemoteExecutionStaticMetadata;
use buck2_re_configuration::RemoteExecutionStaticMetadataImpl;
use chrono::DateTime;
use chrono::Utc;
use dupe::Dupe;
use either::Either;
use fbinit::FacebookInit;
use futures::stream::BoxStream;
use futures::FutureExt;
use futures::StreamExt;
use futures::TryFutureExt;
use gazebo::prelude::*;
use itertools::Itertools;
use prost::Message;
use remote_execution as RE;
use remote_execution::ActionHistoryInfo;
use remote_execution::ActionResultRequest;
use remote_execution::ActionResultResponse;
use remote_execution::BuckInfo;
use remote_execution::DownloadRequest;
use remote_execution::ExecuteRequest;
use remote_execution::ExecuteResponse;
use remote_execution::ExecuteWithProgressResponse;
use remote_execution::GetDigestsTtlRequest;
use remote_execution::HostResourceRequirements;
use remote_execution::InlinedBlobWithDigest;
use remote_execution::NamedDigest;
use remote_execution::NamedDigestWithPermissions;
use remote_execution::REClient;
use remote_execution::REClientBuilder;
use remote_execution::REClientError;
use remote_execution::RemoteExecutionMetadata;
use remote_execution::Stage;
use remote_execution::TActionResult2;
use remote_execution::TCode;
use remote_execution::TDigest;
use remote_execution::TExecutionPolicy;
use remote_execution::UploadRequest;
use remote_execution::WriteActionResultRequest;
use tokio::sync::Semaphore;

use crate::digest::CasDigestToReExt;
use crate::digest_config::DigestConfig;
use crate::directory::ActionImmutableDirectory;
use crate::execute::action_digest::ActionDigest;
use crate::execute::blobs::ActionBlobs;
use crate::execute::executor_stage_async;
use crate::execute::manager::CommandExecutionManager;
use crate::knobs::ExecutorGlobalKnobs;
use crate::materialize::materializer::Materializer;
use crate::re::action_identity::ReActionIdentity;
use crate::re::metadata::RemoteExecutionMetadataExt;
use crate::re::stats::OpStats;
use crate::re::stats::RemoteExecutionClientOpStats;
use crate::re::stats::RemoteExecutionClientStats;
use crate::re::uploader::UploadStats;
use crate::re::uploader::Uploader;

#[derive(Clone, Dupe, Allocative)]
pub struct RemoteExecutionClient {
    data: Arc<RemoteExecutionClientData>,
}

// The large one is the actual default case
#[allow(clippy::large_enum_variant)]
pub enum ExecuteResponseOrCancelled {
    Response(ExecuteResponse),
    Cancelled,
}

#[derive(Allocative)]
struct RemoteExecutionClientData {
    client: RemoteExecutionClientImpl,
    uploads: OpStats,
    downloads: OpStats,
    action_cache: OpStats,
    executes: OpStats,
    materializes: OpStats,
    write_action_results: OpStats,
    get_digest_expirations: OpStats,
}

impl RemoteExecutionClient {
    pub async fn new(
        fb: FacebookInit,
        skip_remote_cache: bool,
        static_metadata: Arc<RemoteExecutionStaticMetadata>,
        logs_dir_path: Option<&AbsNormPath>,
        buck_out_path: &AbsNormPath,
    ) -> anyhow::Result<Self> {
        let client = RemoteExecutionClientImpl::new(
            fb,
            skip_remote_cache,
            static_metadata,
            logs_dir_path,
            buck_out_path,
        )
        .await?;

        Ok(Self {
            data: Arc::new(RemoteExecutionClientData {
                client,
                uploads: OpStats::default(),
                downloads: OpStats::default(),
                action_cache: OpStats::default(),
                executes: OpStats::default(),
                materializes: OpStats::default(),
                write_action_results: OpStats::default(),
                get_digest_expirations: OpStats::default(),
            }),
        })
    }

    pub async fn new_retry(
        fb: FacebookInit,
        skip_remote_cache: bool,
        times: usize, // 0 is treated as 1
        static_metadata: Arc<RemoteExecutionStaticMetadata>,
        logs_dir_path: Option<&AbsNormPath>,
        buck_out_path: &AbsNormPath,
    ) -> anyhow::Result<Self> {
        // Loop happens times-1 times at most
        for i in 1..times {
            match Self::new(
                fb,
                skip_remote_cache,
                static_metadata.dupe(),
                logs_dir_path,
                buck_out_path,
            )
            .await
            {
                Ok(v) => return Ok(v),
                Err(e) => {
                    tracing::warn!(
                        "Failed to connect to RE, retrying after sleeping {} seconds: {:#?}",
                        i,
                        e
                    );
                    tokio::time::sleep(Duration::from_secs(i as u64)).await;
                }
            }
        }
        Self::new(
            fb,
            skip_remote_cache,
            static_metadata,
            logs_dir_path,
            buck_out_path,
        )
        .await
    }

    fn decorate_error(&self, source: anyhow::Error) -> anyhow::Error {
        source.context(format!(
            "Remote Execution Error ({})",
            self.get_session_id()
        ))
    }

    pub async fn action_cache(
        &self,
        action_digest: ActionDigest,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Option<ActionResultResponse>> {
        self.data
            .action_cache
            .op(self.data.client.action_cache(action_digest, use_case))
            .await
    }

    pub async fn upload(
        &self,
        fs: &ProjectRoot,
        materializer: &Arc<dyn Materializer>,
        blobs: &ActionBlobs,
        dir_path: &ProjectRelativePath,
        input_dir: &ActionImmutableDirectory,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> anyhow::Result<UploadStats> {
        self.data
            .uploads
            .op(self
                .data
                .client
                .upload(
                    fs,
                    materializer,
                    blobs,
                    dir_path,
                    input_dir,
                    use_case,
                    digest_config,
                )
                .map_err(|e| self.decorate_error(e)))
            .await
    }

    pub async fn upload_files_and_directories(
        &self,
        files_with_digest: Vec<NamedDigest>,
        directories: Vec<remote_execution::Path>,
        inlined_blobs_with_digest: Vec<InlinedBlobWithDigest>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<()> {
        self.data
            .uploads
            .op(self
                .data
                .client
                .upload_files_and_directories(
                    files_with_digest,
                    directories,
                    inlined_blobs_with_digest,
                    use_case,
                )
                .map_err(|e| self.decorate_error(e)))
            .await
    }

    pub async fn execute(
        &self,
        action_digest: ActionDigest,
        platform: &RE::Platform,
        use_case: RemoteExecutorUseCase,
        identity: &ReActionIdentity<'_>,
        manager: &mut CommandExecutionManager,
        skip_cache_read: bool,
        skip_cache_write: bool,
        re_max_queue_time: Option<Duration>,
        knobs: &ExecutorGlobalKnobs,
    ) -> anyhow::Result<ExecuteResponseOrCancelled> {
        self.data
            .executes
            .op(self
                .data
                .client
                .execute(
                    action_digest,
                    platform,
                    use_case,
                    identity,
                    manager,
                    skip_cache_read,
                    skip_cache_write,
                    re_max_queue_time,
                    knobs,
                )
                .map_err(|e| self.decorate_error(e)))
            .await
    }

    pub async fn materialize_files(
        &self,
        files: Vec<NamedDigestWithPermissions>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<()> {
        self.data
            .materializes
            .op(self.data.client.materialize_files(files, use_case))
            .await
    }

    pub async fn download_typed_blobs<T: Message + Default>(
        &self,
        digests: Vec<TDigest>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Vec<T>> {
        self.data
            .downloads
            .op(self
                .data
                .client
                .download_typed_blobs(digests, use_case)
                .map_err(|e| self.decorate_error(e)))
            .await
    }

    pub async fn download_blob(
        &self,
        digest: &TDigest,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Vec<u8>> {
        self.data
            .downloads
            .op(self
                .data
                .client
                .download_blob(digest, use_case)
                .map_err(|e| self.decorate_error(e)))
            .await
    }

    pub async fn upload_blob(
        &self,
        blob: Vec<u8>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<TDigest> {
        self.data
            .uploads
            .op(self
                .data
                .client
                .upload_blob(blob, use_case)
                .map_err(|e| self.decorate_error(e)))
            .await
    }

    pub async fn get_digest_expirations(
        &self,
        digests: Vec<TDigest>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Vec<(TDigest, DateTime<Utc>)>> {
        self.data
            .get_digest_expirations
            .op(self
                .data
                .client
                .get_digest_expirations(digests, use_case)
                .map_err(|e| self.decorate_error(e)))
            .await
    }

    pub async fn write_action_result(
        &self,
        digest: TDigest,
        result: TActionResult2,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<()> {
        self.data
            .write_action_results
            .op(self
                .data
                .client
                .write_action_result(digest, result, use_case)
                .map_err(|e| self.decorate_error(e)))
            .await
    }

    pub fn get_session_id(&self) -> &str {
        self.data.client.client().get_session_id()
    }

    pub fn get_experiment_name(&self) -> anyhow::Result<Option<String>> {
        self.data.client.client().get_experiment_name()
    }

    pub fn get_network_stats(&self) -> anyhow::Result<RemoteExecutionClientStats> {
        let updated = self
            .data
            .client
            .client()
            .get_network_stats()
            .context("Error getting updated network stats")?;

        let uploaded = updated.uploaded as u64;
        let downloaded = updated.downloaded as u64;

        Ok(RemoteExecutionClientStats {
            uploaded,
            downloaded,
            uploads: RemoteExecutionClientOpStats::from(&self.data.uploads),
            downloads: RemoteExecutionClientOpStats::from(&self.data.downloads),
            executes: RemoteExecutionClientOpStats::from(&self.data.executes),
            action_cache: RemoteExecutionClientOpStats::from(&self.data.action_cache),
            write_action_results: RemoteExecutionClientOpStats::from(
                &self.data.write_action_results,
            ),
            materializes: RemoteExecutionClientOpStats::from(&self.data.materializes),
            get_digest_expirations: RemoteExecutionClientOpStats::from(
                &self.data.get_digest_expirations,
            ),
        })
    }
}

#[derive(Allocative)]
struct RemoteExecutionClientImpl {
    #[allocative(skip)]
    client: Option<REClient>,
    skip_remote_cache: bool,
    /// How many simultaneous requests to RE
    #[allocative(skip)]
    cas_semaphore: Arc<Semaphore>,
    /// How many files we can be downloading concurrently.
    #[allocative(skip)]
    download_files_semapore: Arc<Semaphore>,
    /// How many files to kick off downloading concurrently for one request. This should be smaller
    /// than the files semaphore to ensure we can actually *acquire* that semaphore.
    download_chunk_size: usize,
}

fn re_platform(x: &RE::Platform) -> remote_execution::TPlatform {
    remote_execution::TPlatform {
        properties: x.properties.map(|x| remote_execution::TProperty {
            name: x.name.clone(),
            value: x.value.clone(),
            ..Default::default()
        }),
        ..Default::default()
    }
}

impl RemoteExecutionClientImpl {
    async fn new(
        fb: FacebookInit,
        skip_remote_cache: bool,
        static_metadata: Arc<RemoteExecutionStaticMetadata>,
        maybe_logs_dir_path: Option<&AbsNormPath>,
        buck_out_path: &AbsNormPath,
    ) -> anyhow::Result<Self> {
        tracing::info!("Creating a new RE client");

        let res: anyhow::Result<Self> = try {
            static DOWNLOAD_CONCURRENCY: EnvHelper<usize> =
                EnvHelper::new("BUCK2_RE_DOWNLOAD_CONCURRENCY");

            let download_concurrency = DOWNLOAD_CONCURRENCY.get_copied()?.unwrap_or(256);

            // Split things up into smaller chunks.
            let download_chunk_size = std::cmp::max(download_concurrency / 8, 1);

            #[cfg(fbcode_build)]
            let client = {
                use buck2_core::fs::fs_util;
                use remote_execution::create_default_config;
                use remote_execution::CASDaemonClientCfg;
                use remote_execution::CopyPolicy;
                use remote_execution::CurlReactorConfig;
                use remote_execution::EmbeddedCASDaemonClientCfg;
                use remote_execution::RichClientMode;

                let mut re_client_config = create_default_config();
                re_client_config.action_cache_client_config.connection_count =
                    static_metadata.action_cache_connection_count;
                re_client_config.action_cache_client_config.address =
                    static_metadata.action_cache_address.clone();
                re_client_config.execution_client_config.connection_count =
                    static_metadata.engine_connection_count;
                re_client_config.execution_client_config.address =
                    static_metadata.engine_address.clone();

                let mut embedded_cas_daemon_config = EmbeddedCASDaemonClientCfg {
                    connection_count: static_metadata.cas_connection_count,
                    address: static_metadata.cas_address.clone(),
                    name: "buck2".to_owned(),
                    ..Default::default()
                };

                // buck2 makes find_missing calls for the same blobs
                // so having a 50Mb cache to amortize that
                embedded_cas_daemon_config
                    .cache_config
                    .find_missing_cache_size_byte = 50 << 20;
                // a small cache maps inodes to digests
                // useful for both uploads and downloads
                embedded_cas_daemon_config.cache_config.digest_cache_size = 100000;

                // prevents downloading the same trees (dirs)
                embedded_cas_daemon_config
                    .rich_client_config
                    .get_tree_cache_size = 50 << 20; // 50 Mb

                // disabling zippy rich client until we have limits in place
                embedded_cas_daemon_config
                    .rich_client_config
                    .zdb_client_mode = if static_metadata.use_zippy_rich_client {
                    RichClientMode::HYBRID
                } else {
                    RichClientMode::DISABLED
                };
                embedded_cas_daemon_config.rich_client_config.disable_p2p =
                    !static_metadata.use_p2p;
                embedded_cas_daemon_config
                    .rich_client_config
                    .enable_rich_client = static_metadata.use_manifold_rich_client;

                if let Some(channels) = static_metadata.rich_client_channels_per_blob {
                    embedded_cas_daemon_config
                        .rich_client_config
                        .number_of_parallel_channels = channels;
                }

                if let Some(attempt_timeout) = static_metadata.rich_client_attempt_timeout_ms {
                    embedded_cas_daemon_config
                        .rich_client_config
                        .attempt_timeout_ms = attempt_timeout;
                }

                if let Some(retries_count) = static_metadata.rich_client_retries_count {
                    embedded_cas_daemon_config
                        .rich_client_config
                        .number_of_retries = retries_count;
                }

                embedded_cas_daemon_config.force_enable_deduplicate_find_missing =
                    static_metadata.force_enable_deduplicate_find_missing;

                // Will either choose the SOFT_COPY (on some linux fs like btrfs/extfs etc, on Mac if using APFS) or FULL_COPY otherwise
                embedded_cas_daemon_config.copy_policy = CopyPolicy::BEST_AVAILABLE;
                embedded_cas_daemon_config.meterialization_mount_path = Some(
                    buck_out_path
                        .to_str()
                        .context("invalid meterialization_mount_path")?
                        .to_owned(),
                );

                embedded_cas_daemon_config.thread_count = static_metadata.cas_thread_count;

                // make sure that outputs are writable
                // otherwise actions that are modifying outputs will fail due to a permission error
                embedded_cas_daemon_config.writable_outputs = true;

                re_client_config.cas_client_config =
                    CASDaemonClientCfg::embedded_config(embedded_cas_daemon_config);
                if let Some(logs_dir_path) = maybe_logs_dir_path {
                    // make sure that the log dir exists as glog is expecting that :(
                    fs_util::create_dir_all(logs_dir_path)?;
                    re_client_config.log_file_location = Some(
                        logs_dir_path
                            .to_str()
                            .context("Invalid log_file_location")?
                            .to_owned(),
                    );
                    // keep last 10 sessions (similar to a number of buck builds)
                    re_client_config.log_rollup_window_size = 10;
                }

                re_client_config.check_file_existence_when_no_cache_enabled = false;
                re_client_config.check_parent_directories_exist = false;
                if static_metadata.curl_reactor_max_number_of_retries.is_some()
                    || static_metadata.curl_reactor_connection_timeout_ms.is_some()
                    || static_metadata.curl_reactor_request_timeout_ms.is_some()
                {
                    let mut curl_config = CurlReactorConfig {
                        ..Default::default()
                    };

                    if let Some(max_number_of_retries) =
                        static_metadata.curl_reactor_max_number_of_retries
                    {
                        curl_config.max_number_of_retries = max_number_of_retries;
                    }

                    if let Some(request_timeout_ms) =
                        static_metadata.curl_reactor_request_timeout_ms
                    {
                        curl_config.request_timeout_ms = request_timeout_ms;
                    }

                    if let Some(connection_timeout_ms) =
                        static_metadata.curl_reactor_connection_timeout_ms
                    {
                        curl_config.connection_timeout_ms = connection_timeout_ms;
                    }

                    re_client_config.curl_reactor_config = Some(curl_config);
                }

                re_client_config.features_config_path = static_metadata
                    .features_config_path
                    .as_deref()
                    .unwrap_or("remote_execution/features/client_buck2")
                    .to_owned();

                // TODO(ndmitchell): For now, we just drop RE log messages, but ideally we'd put them in our log stream.
                let logger = slog::Logger::root(slog::Discard, slog::o!());
                REClientBuilder::new(fb)
                    .with_config(re_client_config)
                    .with_logger(logger)
                    .build_and_connect()
                    .await?
            };

            #[cfg(not(fbcode_build))]
            let client = {
                let _unused = (fb, maybe_logs_dir_path, buck_out_path);

                REClientBuilder::build_and_connect(&static_metadata.0).await?
            };

            Self {
                client: Some(client),
                skip_remote_cache,
                cas_semaphore: Arc::new(Semaphore::new(static_metadata.cas_semaphore_size())),
                download_files_semapore: Arc::new(Semaphore::new(download_concurrency)),
                download_chunk_size,
            }
        };

        res.context("RE: creating client")
    }

    fn client(&self) -> &REClient {
        self.client
            .as_ref()
            .expect("REClient is always present unless dropped")
    }

    async fn action_cache(
        &self,
        action_digest: ActionDigest,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Option<ActionResultResponse>> {
        let res = self
            .client()
            .get_action_cache_client()
            .get_action_result(
                use_case.metadata(),
                ActionResultRequest {
                    digest: action_digest.to_re(),
                    ..Default::default()
                },
            )
            .await;

        match res {
            Ok(r) => Ok(Some(r)),
            Err(e) => {
                if e.downcast_ref::<REClientError>()
                    .map(|e| e.code == TCode::NOT_FOUND)
                    == Some(true)
                {
                    Ok(None)
                } else {
                    Err(e)
                }
            }
        }
    }

    async fn upload(
        &self,
        fs: &ProjectRoot,
        materializer: &Arc<dyn Materializer>,
        blobs: &ActionBlobs,
        dir_path: &ProjectRelativePath,
        input_dir: &ActionImmutableDirectory,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> anyhow::Result<UploadStats> {
        // Actually upload to CAS
        let _cas = self.cas_semaphore.acquire().await;
        Uploader::upload(
            fs,
            self.client().get_cas_client(),
            materializer,
            dir_path,
            input_dir,
            blobs,
            use_case,
            digest_config,
        )
        .await
    }

    async fn upload_files_and_directories(
        &self,
        files_with_digest: Vec<NamedDigest>,
        directories: Vec<remote_execution::Path>,
        inlined_blobs_with_digest: Vec<InlinedBlobWithDigest>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<()> {
        self.client()
            .get_cas_client()
            .upload(
                use_case.metadata(),
                UploadRequest {
                    files_with_digest: Some(files_with_digest),
                    inlined_blobs_with_digest: Some(inlined_blobs_with_digest),
                    directories: Some(directories),
                    upload_only_missing: true,
                    ..Default::default()
                },
            )
            .await?;
        Ok(())
    }

    async fn execute_impl(
        &self,
        metadata: RemoteExecutionMetadata,
        request: ExecuteRequest,
        action_digest: &ActionDigest,
        manager: &mut CommandExecutionManager,
        re_max_queue_time: Option<Duration>,
        platform: &remote_execution::Platform,
        knobs: &ExecutorGlobalKnobs,
    ) -> anyhow::Result<ExecuteResponseOrCancelled> {
        use buck2_data::re_stage;
        use buck2_data::ReExecute;
        use buck2_data::ReQueue;
        use buck2_data::ReUnknown;
        use buck2_data::ReWorkerDownload;
        use buck2_data::ReWorkerUpload;

        let action_key = if knobs.log_action_keys {
            metadata
                .action_history_info
                .as_ref()
                .map(|h| h.action_key.clone())
        } else {
            None
        };

        #[allow(clippy::large_enum_variant)]
        enum ResponseOrStateChange {
            Present(ExecuteWithProgressResponse),
            Cancelled,
        }

        /// Wait for either the ExecuteResponse to show up, or a stage change, within a span
        /// on the CommandExecutionManager.
        async fn wait_for_response_or_stage_change(
            receiver: &mut BoxStream<'static, anyhow::Result<ExecuteWithProgressResponse>>,
            previous_stage: Stage,
            report_stage: re_stage::Stage,
            manager: &mut CommandExecutionManager,
            re_max_queue_time: Option<Duration>,
        ) -> anyhow::Result<ResponseOrStateChange> {
            executor_stage_async(
                buck2_data::ReStage {
                    stage: Some(report_stage),
                },
                async move {
                    loop {
                        let next = futures::future::select(
                            manager.liveliness_observer.while_alive(),
                            receiver.next(),
                        );

                        let event = match next.await {
                            futures::future::Either::Left((_dead, _)) => {
                                return Ok(ResponseOrStateChange::Cancelled);
                            }
                            futures::future::Either::Right((event, _)) => match event {
                                Some(event) => event,
                                None => {
                                    return Err(anyhow::anyhow!(
                                        "RE execution did not yield a ExecuteResponse"
                                    ));
                                }
                            },
                        };

                        let event = event.context("Error was returned on the stream by RE")?;

                        if event.execute_response.is_some() || event.stage != previous_stage {
                            return Ok(ResponseOrStateChange::Present(event));
                        }

                        // TODO: This should be one block up?
                        if let Some(re_max_queue_time) = re_max_queue_time {
                            if let Some(info) = event.metadata.task_info {
                                let est = u64::try_from(info.estimated_queue_time_ms)
                                    .context("estimated_queue_time_ms from RE is negative")?;
                                let queue_time = Duration::from_millis(est);

                                if queue_time > re_max_queue_time {
                                    manager.on_result_delayed();
                                }
                            }
                        }
                    }
                },
            )
            .await
        }

        fn re_stage_from_exe_stage(
            stage: Stage,
            action_digest: String,
            platform: &remote_execution::Platform,
            action_key: &Option<String>,
        ) -> re_stage::Stage {
            match stage {
                Stage::QUEUED => re_stage::Stage::Queue(ReQueue { action_digest }),
                Stage::MATERIALIZING_INPUT => {
                    re_stage::Stage::WorkerDownload(ReWorkerDownload { action_digest })
                }
                Stage::EXECUTING => re_stage::Stage::Execute(ReExecute {
                    action_digest,
                    platform: Some(transform_platform(platform)),
                    action_key: action_key.clone(),
                }),
                Stage::UPLOADING_OUTPUT => {
                    re_stage::Stage::WorkerUpload(ReWorkerUpload { action_digest })
                }
                _ => {
                    tracing::debug!(
                        "Received unexpected RE stage {:#?} for action: {}",
                        stage,
                        action_digest
                    );
                    re_stage::Stage::Unknown(ReUnknown { action_digest })
                }
            }
        }

        // There are four possible outcomes, three of them failures:
        // 1. We get Err, which means we failed to communicate to RE.
        // 2. We get Ok(error.code != TCode::Ok), which means we have an RE-level failure
        // 3. We get Ok(action_result.exit_code != 0), which means the command failed
        // 4. We get Ok(0 everywhere), which means it worked

        // Obtain a stream of events from RE. If this fails then that is case #1 above so we
        // bail.
        let mut receiver = self
            .client()
            .get_execution_client()
            .execute_with_progress(metadata, request)
            // boxed() to segment the future
            .boxed()
            .await
            .map_err(anyhow::Error::from)
            .context("Failed to start remote execution")?;

        // Now we wait until the ExecuteResponse shows up, and produce events accordingly. If
        // this doesn't give us an ExecuteResponse then this is case #1 again so we also fail.
        let action_digest_str = action_digest.to_string();
        let mut exe_stage = Stage::QUEUED;

        loop {
            let progress_response = wait_for_response_or_stage_change(
                &mut receiver,
                exe_stage,
                re_stage_from_exe_stage(
                    exe_stage,
                    action_digest_str.clone(),
                    platform,
                    &action_key,
                ),
                manager,
                re_max_queue_time,
            )
            .await?;

            let progress_response = match progress_response {
                ResponseOrStateChange::Present(r) => r,
                ResponseOrStateChange::Cancelled => {
                    return Ok(ExecuteResponseOrCancelled::Cancelled);
                }
            };

            // Return the result if we're done
            if let Some(execute_response) = progress_response.execute_response {
                return Ok(ExecuteResponseOrCancelled::Response(execute_response));
            }

            // Change the stage
            exe_stage = progress_response.stage;
        }
    }

    pub async fn execute(
        &self,
        action_digest: ActionDigest,
        platform: &RE::Platform,
        use_case: RemoteExecutorUseCase,
        identity: &ReActionIdentity<'_>,
        manager: &mut CommandExecutionManager,
        skip_cache_read: bool,
        skip_cache_write: bool,
        re_max_queue_time: Option<Duration>,
        knobs: &ExecutorGlobalKnobs,
    ) -> anyhow::Result<ExecuteResponseOrCancelled> {
        let metadata = RemoteExecutionMetadata {
            action_history_info: Some(ActionHistoryInfo {
                action_key: identity.action_key.clone(),
                disable_retry_on_oom: false,
                ..Default::default()
            }),
            host_resource_requirements: Some(HostResourceRequirements {
                affinity_keys: vec![identity.affinity_key.clone()],
                input_files_bytes: identity.paths.input_files_bytes() as i64,
                ..Default::default()
            }),
            platform: Some(re_platform(platform)),
            do_not_cache: skip_cache_write,
            buck_info: Some(BuckInfo {
                build_id: identity.trace_id.to_string(),
                ..Default::default()
            }),
            ..use_case.metadata()
        };
        let request = ExecuteRequest {
            skip_cache_lookup: self.skip_remote_cache || skip_cache_read,
            execution_policy: Some(TExecutionPolicy::default()),
            action_digest: action_digest.to_re(),
            ..Default::default()
        };
        self.execute_impl(
            metadata,
            request,
            &action_digest,
            manager,
            re_max_queue_time,
            platform,
            knobs,
        )
        .await
        .with_context(|| format!("RE: execution with digest {}", &action_digest))
    }

    /// Fetches a list of digests from the CAS and casts them to Tree objects.
    /// If fetching or decoding fails for one or more digests, returns an Err.
    async fn download_typed_blobs<T: Message + Default>(
        &self,
        digests: Vec<TDigest>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Vec<T>> {
        if digests.is_empty() {
            return Ok(Vec::new());
        }
        let expected_blobs = digests.len();
        let response = self
            .client()
            .get_cas_client()
            .download(
                use_case.metadata(),
                DownloadRequest {
                    inlined_digests: Some(digests),
                    ..Default::default()
                },
            )
            .await?;

        let mut blobs: Vec<T> = Vec::with_capacity(expected_blobs);
        if let Some(ds) = response.inlined_blobs {
            for d in ds {
                blobs.push(Message::decode(d.blob.as_slice()).with_context(|| {
                    format!("Failed to Protobuf decode tree at `{}`", d.digest)
                })?);
            }
        }

        // This shouldn't happen, but we can't just assume the CAS won't ever break
        if blobs.len() != expected_blobs {
            return Err(anyhow::anyhow!(
                "CAS client returned fewer blobs than expected."
            ));
        }

        Ok(blobs)
    }

    pub async fn download_blob(
        &self,
        digest: &TDigest,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Vec<u8>> {
        let response = self
            .client()
            .get_cas_client()
            .download(
                use_case.metadata(),
                DownloadRequest {
                    inlined_digests: Some(vec![digest.clone()]),
                    ..Default::default()
                },
            )
            // boxed() to segment the future
            .boxed()
            .await
            .with_context(|| format!("Download request failed for digest {}", digest))?;

        response
            .inlined_blobs
            .into_iter()
            .flat_map(|blobs| blobs.into_iter())
            .next()
            .map(|blob| blob.blob)
            .with_context(|| format!("No digest was returned in request for {}", digest))
    }

    pub async fn upload_blob(
        &self,
        blob: Vec<u8>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<TDigest> {
        self.client().upload_blob(blob, use_case.metadata()).await
    }

    async fn materialize_files(
        &self,
        files: Vec<NamedDigestWithPermissions>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<()> {
        let use_case = &use_case;

        let futs = chunks(files, self.download_chunk_size).map(|chunk| async move {
            let _permit = self
                .download_files_semapore
                .acquire_many(chunk.len().try_into().context("chunk is too large")?)
                .await
                .context("Failed to acquire download_files_semapore")?;

            self.client()
                .get_cas_client()
                .download(
                    use_case.metadata(),
                    DownloadRequest {
                        file_digests: Some(chunk),
                        ..Default::default()
                    },
                )
                .await?;

            anyhow::Ok(())
        });

        futures::future::try_join_all(futs).await?;

        Ok(())
    }

    async fn get_digest_expirations(
        &self,
        digests: Vec<TDigest>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Vec<(TDigest, DateTime<Utc>)>> {
        let now = Utc::now();

        let ttls = self
            .client()
            .get_cas_client()
            .get_digests_ttl(
                use_case.metadata(),
                GetDigestsTtlRequest {
                    digests,
                    ..Default::default()
                },
            )
            .await?
            .digests_with_ttl;

        Ok(ttls
            .into_iter()
            .map(|t| (t.digest, now + chrono::Duration::seconds(t.ttl)))
            .collect())
    }

    async fn write_action_result(
        &self,
        digest: TDigest,
        result: TActionResult2,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<()> {
        self.client()
            .get_action_cache_client()
            .write_action_result(
                use_case.metadata(),
                WriteActionResultRequest {
                    action_digest: digest,
                    action_result: result,
                    ..Default::default()
                },
            )
            .await?;

        Ok(())
    }
}

fn transform_platform(platform: &remote_execution::Platform) -> buck2_data::RePlatform {
    buck2_data::RePlatform {
        properties: platform
            .properties
            .iter()
            .map(|property| buck2_data::re_platform::Property {
                name: property.name.clone(),
                value: property.value.clone(),
            })
            .collect(),
    }
}

/// Drop the REClient on a blocking thread. The REClient destructor does a blocking wait on async
/// calls (it tells the server to cancel its calls, but it waits for an ack), so we shouldn't drop
/// it on a runtime thread.
impl Drop for RemoteExecutionClientImpl {
    fn drop(&mut self) {
        tracing::info!("Dropping RE client");
        let client = self.client.take();

        if let Ok(handle) = tokio::runtime::Handle::try_current() {
            handle.spawn_blocking(move || {
                drop(client);
                tracing::info!("Dropped RE client");
            });
        }
    }
}

#[allow(clippy::needless_collect)] // chunks() is not Send.
fn chunks<T>(v: Vec<T>, chunk_size: usize) -> impl Iterator<Item = Vec<T>> {
    if !v.is_empty() && v.len() <= chunk_size {
        return Either::Left(std::iter::once(v));
    }

    let chunks = v
        .into_iter()
        .chunks(chunk_size)
        .into_iter()
        .map(|c| c.collect::<Vec<_>>())
        .collect::<Vec<_>>();

    Either::Right(chunks.into_iter())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_chunks_skips() {
        assert_eq!(chunks(Vec::<usize>::new(), 1).next(), None);
    }

    #[test]
    fn test_chunks_reuses() {
        let v = vec![1, 2, 3];
        let addr = v.as_ptr() as usize;
        let mut it = chunks(v, 3);
        assert_eq!(it.next().unwrap().as_ptr() as usize, addr);
        assert_eq!(it.next(), None);
    }

    #[test]
    fn test_chunks_splits() {
        let v = vec![1, 2, 3];
        let mut it = chunks(v, 2);
        assert_eq!(it.next(), Some(vec![1, 2]));
        assert_eq!(it.next(), Some(vec![3]));
        assert_eq!(it.next(), None);
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

// This triggers on Arc<Arc<...>>, but we do that here for lifetime/ownership reasons
#![allow(clippy::redundant_allocation)]

use std::sync::Arc;
use std::sync::Mutex;
use std::sync::RwLock;
use std::sync::Weak;
use std::time::Duration;

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_common::result::SharedResult;
use buck2_core::async_once_cell::AsyncOnceCell;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_re_configuration::RemoteExecutionStaticMetadata;
use chrono::DateTime;
use chrono::Utc;
use dupe::Dupe;
use fbinit::FacebookInit;
use futures::FutureExt;
use prost::Message;
use remote_execution as RE;
use remote_execution::ActionResultResponse;
use remote_execution::InlinedBlobWithDigest;
use remote_execution::NamedDigest;
use remote_execution::NamedDigestWithPermissions;
use remote_execution::TActionResult2;
use remote_execution::TDigest;

use crate::digest_config::DigestConfig;
use crate::directory::ActionImmutableDirectory;
use crate::execute::action_digest::ActionDigest;
use crate::execute::blobs::ActionBlobs;
use crate::execute::manager::CommandExecutionManager;
use crate::knobs::ExecutorGlobalKnobs;
use crate::materialize::materializer::Materializer;
use crate::re::action_identity::ReActionIdentity;
use crate::re::client::ExecuteResponseOrCancelled;
use crate::re::client::RemoteExecutionClient;
use crate::re::re_get_session_id::ReGetSessionId;
use crate::re::stats::RemoteExecutionClientStats;
use crate::re::uploader::UploadStats;

/// Lifetime management of the Remote Execution connection (i.e. the RemoteExecutionClient).
///
/// The [ReConnectionManager] ensures that there is only ever a single remote execution connection. It
/// vends out a [ReConnectionHandle] on calls to [ReConnectionManager::get_re_connection()].
///
/// There will only be at most a single RE connection during the lifetime of an [ReConnectionHandle], the
/// connection will not be closed as long as any handle still exists. The [ReConnectionHandle] then vends
/// out a [ManagedRemoteExecutionClient].
///
/// The [ManagedRemoteExecutionClient] has all the RE functions on it. The first time one of these is
/// called will be when the RE connection is actually made, if no calls to RE are made, we don't create
/// a connection.
///
/// When a [ManagedRemoteExecutionClient] is used, if it's source [ReConnectionHandle] still exists then
/// the calls will go to the underlying RE client (and making the connection if this is the first call). If,
/// however, the [ReConnectionHandle] has been dropped the calls will fail (even if the connection is still
/// alive from other handles).
///
/// When all existing [ReConnectionHandle] (for a particular [ReConnectionManager]) are dropped, the
/// underlying connection will be closed.
///
/// This approach allows us to ensure that all RE interactions for a particular buck command use the
/// same RE session. Concurrent commands will share an RE session.

#[derive(Clone, Allocative)]
struct RemoteExecutionConfig {
    #[allocative(skip)] // TODO(nga): implement in `allocative`.
    fb: FacebookInit,
    /// whether to skip the cache when performing RE
    skip_remote_cache: bool,
    /// number of retries when attempting the initial RE connection
    connection_retries: usize,
    static_metadata: Arc<RemoteExecutionStaticMetadata>,
    logs_dir_path: Option<AbsNormPathBuf>,
    buck_out_path: AbsNormPathBuf,
}

impl RemoteExecutionConfig {
    async fn connect_now(&self) -> anyhow::Result<RemoteExecutionClient> {
        RemoteExecutionClient::new_retry(
            self.fb,
            self.skip_remote_cache,
            self.connection_retries,
            self.static_metadata.dupe(),
            self.logs_dir_path.as_deref(),
            &self.buck_out_path,
        )
        .await
    }
}

pub trait ReConnectionObserver: Allocative + 'static + Send + Sync {
    fn session_created(&self, client: &RemoteExecutionClient);
}

#[derive(Allocative)]
struct LazyRemoteExecutionClient {
    client: AsyncOnceCell<SharedResult<RemoteExecutionClient>>,
    observers: Mutex<Vec<Weak<dyn ReConnectionObserver>>>,
    config: RemoteExecutionConfig,
}

impl LazyRemoteExecutionClient {
    fn new(config: RemoteExecutionConfig) -> LazyRemoteExecutionClient {
        Self {
            client: AsyncOnceCell::new(),
            observers: Mutex::new(Vec::new()),
            config,
        }
    }

    /// Apply F to the client contained in this LazyRemoteExecutionClient if and only if there is a
    /// valid client.
    fn with_client<F, T>(&self, f: F) -> Option<T>
    where
        F: FnOnce(&RemoteExecutionClient) -> T,
    {
        match self.client.get().as_ref() {
            Some(Ok(client)) => Some(f(client)),
            _ => None,
        }
    }

    async fn get(&self) -> anyhow::Result<&RemoteExecutionClient> {
        // The future from self.init() is large and very rarely required. We want
        // to ensure that (1) it doesn't contribute to the size of the get() future
        // (which is used for basically every call) and (2) isn't even allocated on
        // the fast path.
        // To do this, we .boxed() it and wrap it in an async block that only needs to
        // hold the &self (until it's first polled).
        let init_fut = async move { self.init().boxed().await };
        match self.client.get_or_init(init_fut).await {
            Ok(v) => Ok(v),
            Err(e) => Err(e.dupe().into()),
        }
    }

    async fn init(&self) -> SharedResult<RemoteExecutionClient> {
        let client = self.config.connect_now().await?;

        let mut observers = self.observers.lock().unwrap();
        let owned = std::mem::take(&mut *observers);
        // session_created is the only thing on observer right now, so we can just throw them all out at this point.
        for o in owned.into_iter() {
            if let Some(o) = o.upgrade() {
                o.session_created(&client);
            }
        }
        Ok(client)
    }

    // Adds an observer. This observer will immediately get a session_created() call if we are already connected.
    fn observe(&self, observer: &Arc<dyn ReConnectionObserver>) {
        // lock observers so we don't race with the notifications in init()
        let mut observers = self.observers.lock().unwrap();
        match self.client.get() {
            Some(Ok(client)) => {
                observer.session_created(client);
            }
            Some(_) => {
                // ignored
            }
            _ => {
                observers.push(Arc::downgrade(observer));
            }
        }
    }
}

/// The main manager for the RE connections
#[derive(Allocative)]
pub struct ReConnectionManager {
    // We hold a single Weak to the lazy client. ReConnectionHandle will hold a strong Arc to the same. Once the
    // last ReConnectionHandle is dropped, the client we point to will be dropped and we'll create a new one for
    // the next ReConnectionHandle.
    data: RwLock<Weak<LazyRemoteExecutionClient>>,
    config: RemoteExecutionConfig,
}

impl ReConnectionManager {
    pub fn new(
        fb: FacebookInit,
        skip_remote_cache: bool,
        connection_retries: usize,
        static_metadata: Arc<RemoteExecutionStaticMetadata>,
        logs_dir_path: Option<AbsNormPathBuf>,
        buck_out_path: AbsNormPathBuf,
    ) -> Self {
        Self {
            data: RwLock::new(Weak::new()),
            config: RemoteExecutionConfig {
                fb,
                skip_remote_cache,
                connection_retries,
                static_metadata,
                logs_dir_path,
                buck_out_path,
            },
        }
    }

    /// Gets a new guard that holds a RE connection open
    pub fn get_re_connection(&self) -> ReConnectionHandle {
        ReConnectionHandle::new(self.get_client_handle())
    }

    fn get_client_handle(&self) -> Arc<LazyRemoteExecutionClient> {
        if let Some(conn) = self.data.read().unwrap().upgrade() {
            return conn;
        }

        let mut conn = self.data.write().unwrap();
        match conn.upgrade() {
            None => {
                let new_connection = Arc::new(LazyRemoteExecutionClient::new(self.config.clone()));
                *conn = Arc::downgrade(&new_connection);
                new_connection
            }
            Some(conn) => conn,
        }
    }

    pub fn get_network_stats(&self) -> anyhow::Result<Option<RemoteExecutionClientStats>> {
        let conn = self.data.read().unwrap().upgrade();
        conn.as_ref()
            .and_then(|lazy_client| lazy_client.with_client(|client| client.get_network_stats()))
            .transpose()
    }
}

#[async_trait]
impl ReGetSessionId for ReConnectionManager {
    async fn get_session_id(&self) -> anyhow::Result<String> {
        self.get_re_connection().get_client().get_session_id().await
    }
}

/// The guard that holds the RE connection open. This can outlive the main manager
/// The client obtained from this handle will only be valid while this handle is held.
pub struct ReConnectionHandle {
    /// here we hold an `Arc<Arc>` so that we can hand out a `Weak<Arc>` so that upon the handle
    /// being dropped, the clients will lose the ability to access connection
    // TODO(cjhopman): While we vend out Weak<Arc> with this, due to the way this is stored/used on
    // the dice graph we are guaranteed that there's an Arc that outlives all the Weak from it. That
    // kind of defeats the purpose of this.
    #[cfg_attr(feature = "gazebo_lint", allow(gazebo_lint_arc_on_dupe))]
    connection: Arc<Arc<LazyRemoteExecutionClient>>,
    // We use a similar Arc/Weak to ensure that the observer only lives as long as the connection
    // handle. Otherwise it would be possible to observer a session create attached to one command
    // after that command ended. An alternative would be to register/deregister the connection
    // handle itself as an observer on the lazy client, but that doesn't seem any simpler.
    observer: Option<Arc<dyn ReConnectionObserver>>,
}

impl ReConnectionHandle {
    fn new(connection: Arc<LazyRemoteExecutionClient>) -> Self {
        Self {
            connection: Arc::new(connection),
            observer: None,
        }
    }

    /// Sets the connection observer. This will drop the previous observer if there is one.
    pub fn set_observer(&mut self, observer: Arc<dyn ReConnectionObserver>) {
        // We store it just to give it this handle's lifetime.
        self.connection.observe(&observer);
        self.observer = Some(observer);
    }

    /// gets a client that is tied to the scope of this guard
    pub fn get_client(&self) -> ManagedRemoteExecutionClient {
        ManagedRemoteExecutionClient {
            data: Arc::downgrade(&self.connection),
        }
    }
}

#[derive(Clone, Dupe)]
pub struct ManagedRemoteExecutionClient {
    data: Weak<Arc<LazyRemoteExecutionClient>>,
}

impl ManagedRemoteExecutionClient {
    #[cfg_attr(feature = "gazebo_lint", allow(gazebo_lint_arc_on_dupe))]
    fn lock(&self) -> anyhow::Result<Arc<Arc<LazyRemoteExecutionClient>>> {
        self.data
            .upgrade()
            .context("Internal error: the underlying RE connection has terminated because the corresponding guard has been dropped.")
    }

    pub async fn action_cache(
        &self,
        action_digest: ActionDigest,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Option<ActionResultResponse>> {
        Ok(self
            .lock()?
            .get()
            .await?
            .action_cache(action_digest, use_case)
            .await
            .ok()
            .flatten())
    }

    pub async fn upload(
        &self,
        fs: &ProjectRoot,
        materializer: &Arc<dyn Materializer>,
        blobs: &ActionBlobs,
        dir_path: &ProjectRelativePath,
        input_dir: &ActionImmutableDirectory,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> anyhow::Result<UploadStats> {
        self.lock()?
            .get()
            .await?
            .upload(
                fs,
                materializer,
                blobs,
                dir_path,
                input_dir,
                use_case,
                digest_config,
            )
            .await
    }

    pub async fn upload_files_and_directories(
        &self,
        files_with_digest: Vec<NamedDigest>,
        directories: Vec<remote_execution::Path>,
        inlined_blobs_with_digest: Vec<InlinedBlobWithDigest>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<()> {
        self.lock()?
            .get()
            .await?
            .upload_files_and_directories(
                files_with_digest,
                directories,
                inlined_blobs_with_digest,
                use_case,
            )
            .await
    }

    pub async fn execute(
        &self,
        action_digest: ActionDigest,
        platform: &RE::Platform,
        use_case: RemoteExecutorUseCase,
        identity: &ReActionIdentity<'_>,
        manager: &mut CommandExecutionManager,
        skip_cache_read: bool,
        skip_cache_write: bool,
        re_max_queue_time: Option<Duration>,
        knobs: &ExecutorGlobalKnobs,
    ) -> anyhow::Result<ExecuteResponseOrCancelled> {
        self.lock()?
            .get()
            .await?
            .execute(
                action_digest,
                platform,
                use_case,
                identity,
                manager,
                skip_cache_read,
                skip_cache_write,
                re_max_queue_time,
                knobs,
            )
            .await
    }

    pub async fn materialize_files(
        &self,
        files: Vec<NamedDigestWithPermissions>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<()> {
        self.lock()?
            .get()
            .await?
            .materialize_files(files, use_case)
            .await
    }

    pub async fn download_typed_blobs<T: Message + Default>(
        &self,
        digests: Vec<TDigest>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Vec<T>> {
        self.lock()?
            .get()
            .await?
            .download_typed_blobs(digests, use_case)
            .await
    }

    pub async fn download_blob(
        &self,
        digest: &TDigest,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Vec<u8>> {
        self.lock()?
            .get()
            .await?
            .download_blob(digest, use_case)
            .await
    }

    pub async fn upload_blob(
        &self,
        blob: Vec<u8>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<TDigest> {
        self.lock()?.get().await?.upload_blob(blob, use_case).await
    }

    pub async fn get_digest_expirations(
        &self,
        digests: Vec<TDigest>,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<Vec<(TDigest, DateTime<Utc>)>> {
        self.lock()?
            .get()
            .await?
            .get_digest_expirations(digests, use_case)
            .await
    }

    pub async fn write_action_result(
        &self,
        digest: TDigest,
        result: TActionResult2,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<()> {
        self.lock()?
            .get()
            .await?
            .write_action_result(digest, result, use_case)
            .await
    }

    pub async fn get_session_id(&self) -> anyhow::Result<String> {
        let session_id = self.lock()?.get().await?.get_session_id().to_owned();
        Ok(session_id)
    }

    /// Construct a dummy ManagedRemoteExecutionClient that won't actually work. This is only
    /// remotely useful in tests.
    pub fn testing_new_dummy() -> Self {
        Self { data: Weak::new() }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use buck2_common::executor_config::RemoteExecutorUseCase;
use remote_execution::RemoteExecutionMetadata;

pub trait RemoteExecutionMetadataExt {
    fn metadata(&self) -> RemoteExecutionMetadata;
}

impl RemoteExecutionMetadataExt for RemoteExecutorUseCase {
    fn metadata(&self) -> RemoteExecutionMetadata {
        RemoteExecutionMetadata {
            use_case_id: self.as_str().to_owned(),
            ..Default::default()
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

pub mod action_identity;
pub mod client;
pub mod manager;
pub mod metadata;
pub mod re_get_session_id;
pub mod remote_action_result;
mod stats;
pub mod streams;
pub mod uploader;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::time::Duration;
use std::time::SystemTime;

use anyhow::Context as _;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_miniperf_proto::MiniperfCounter;
use remote_execution::ActionResultResponse;
use remote_execution::ExecuteResponse;
use remote_execution::TDirectory2;
use remote_execution::TExecutedActionMetadata;
use remote_execution::TFile;
use remote_execution::TPerfCount;
use remote_execution::TSubsysPerfCount;
use remote_execution::TTimestamp;

use crate::digest_config::DigestConfig;
use crate::execute::action_digest::ActionDigest;
use crate::execute::kind::CommandExecutionKind;
use crate::execute::result::CommandExecutionMetadata;
use crate::re::manager::ManagedRemoteExecutionClient;
use crate::re::streams::RemoteCommandStdStreams;

pub trait RemoteActionResult: Send + Sync {
    fn output_files(&self) -> &[TFile];
    fn output_directories(&self) -> &[TDirectory2];

    fn execution_kind(&self, digest: ActionDigest) -> CommandExecutionKind;

    fn timing(&self) -> CommandExecutionMetadata;

    fn std_streams(
        &self,
        client: &ManagedRemoteExecutionClient,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> RemoteCommandStdStreams;

    /// The TTL given by RE for the outputs for this action.
    fn ttl(&self) -> i64;
}

impl RemoteActionResult for ExecuteResponse {
    fn output_files(&self) -> &[TFile] {
        &self.action_result.output_files
    }

    fn output_directories(&self) -> &[TDirectory2] {
        &self.action_result.output_directories
    }

    fn execution_kind(&self, digest: ActionDigest) -> CommandExecutionKind {
        CommandExecutionKind::Remote { digest }
    }

    fn timing(&self) -> CommandExecutionMetadata {
        timing_from_re_metadata(&self.action_result.execution_metadata)
    }

    fn std_streams(
        &self,
        client: &ManagedRemoteExecutionClient,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> RemoteCommandStdStreams {
        RemoteCommandStdStreams::new(&self.action_result, client, use_case, digest_config)
    }

    fn ttl(&self) -> i64 {
        self.action_result_ttl
    }
}

impl RemoteActionResult for ActionResultResponse {
    fn output_files(&self) -> &[TFile] {
        &self.action_result.output_files
    }

    fn output_directories(&self) -> &[TDirectory2] {
        &self.action_result.output_directories
    }

    fn execution_kind(&self, digest: ActionDigest) -> CommandExecutionKind {
        CommandExecutionKind::ActionCache { digest }
    }

    fn timing(&self) -> CommandExecutionMetadata {
        let mut timing = timing_from_re_metadata(&self.action_result.execution_metadata);
        timing.wall_time = Duration::ZERO; // This was a cache hit so we didn't wait.
        timing.input_materialization_duration = Duration::ZERO; // This was a cache hit so we didn't wait.
        timing
    }

    fn std_streams(
        &self,
        client: &ManagedRemoteExecutionClient,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> RemoteCommandStdStreams {
        RemoteCommandStdStreams::new(&self.action_result, client, use_case, digest_config)
    }

    fn ttl(&self) -> i64 {
        self.ttl
    }
}

fn timing_from_re_metadata(meta: &TExecutedActionMetadata) -> CommandExecutionMetadata {
    let execution_time = meta
        .execution_completed_timestamp
        .saturating_duration_since(&meta.execution_start_timestamp);

    let re_queue_time = meta
        .last_queued_timestamp
        .saturating_duration_since(&meta.queued_timestamp);

    let start_time = SystemTime::UNIX_EPOCH
        + meta
            .execution_start_timestamp
            .saturating_duration_since(&TTimestamp::unix_epoch());

    let execution_stats = match convert_perf_counts(&meta.instruction_counts) {
        Ok(v) => Some(v),
        Err(e) => {
            tracing::warn!("Invalid instruction counts received from RE: {:#}", e);
            None
        }
    };

    let fetch_input_time = meta
        .input_fetch_completed_timestamp
        .saturating_duration_since(&meta.input_fetch_start_timestamp);

    CommandExecutionMetadata {
        wall_time: execution_time,
        re_queue_time: Some(re_queue_time),
        execution_time,
        start_time,
        execution_stats,
        input_materialization_duration: fetch_input_time,
    }
}

fn convert_perf_counts(
    perf_counts: &TPerfCount,
) -> anyhow::Result<buck2_data::CommandExecutionStats> {
    Ok(buck2_data::CommandExecutionStats {
        cpu_instructions_user: convert_perf_count(&perf_counts.userspace_events)?
            .map(|p| p.adjusted_count()),
        cpu_instructions_kernel: convert_perf_count(&perf_counts.kernel_events)?
            .map(|p| p.adjusted_count()),
    })
}

fn convert_perf_count(perf_count: &TSubsysPerfCount) -> anyhow::Result<Option<MiniperfCounter>> {
    if perf_count.time_running == 0 {
        return Ok(None);
    }

    Ok(Some(MiniperfCounter {
        count: perf_count.count.try_into().context("Invalid count")?,
        time_enabled: perf_count
            .time_enabled
            .try_into()
            .context("Invalid time_enabled")?,
        time_running: perf_count
            .time_running
            .try_into()
            .context("Invalid time_running")?,
    }))
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use allocative::Allocative;
use async_trait::async_trait;

#[async_trait]
pub trait ReGetSessionId: Allocative + Send + Sync + 'static {
    async fn get_session_id(&self) -> anyhow::Result<String>;
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::future::Future;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::Ordering;

use allocative::Allocative;
use futures::FutureExt;

pub struct RemoteExecutionClientOpStats {
    pub started: u32,
    pub finished_successfully: u32,
    pub finished_with_error: u32,
}

impl From<&'_ OpStats> for RemoteExecutionClientOpStats {
    fn from(stats: &OpStats) -> RemoteExecutionClientOpStats {
        RemoteExecutionClientOpStats {
            started: stats.started.load(Ordering::Relaxed),
            finished_successfully: stats.finished_successfully.load(Ordering::Relaxed),
            finished_with_error: stats.finished_with_error.load(Ordering::Relaxed),
        }
    }
}

pub struct RemoteExecutionClientStats {
    /// In bytes.
    pub uploaded: u64,
    /// In bytes.
    pub downloaded: u64,
    pub uploads: RemoteExecutionClientOpStats,
    pub downloads: RemoteExecutionClientOpStats,
    pub action_cache: RemoteExecutionClientOpStats,
    pub executes: RemoteExecutionClientOpStats,
    pub materializes: RemoteExecutionClientOpStats,
    pub write_action_results: RemoteExecutionClientOpStats,
    pub get_digest_expirations: RemoteExecutionClientOpStats,
}

#[derive(Default, Allocative)]
pub(super) struct OpStats {
    started: AtomicU32,
    finished_successfully: AtomicU32,
    finished_with_error: AtomicU32,
}

impl OpStats {
    pub(super) fn op<'a, R, F>(&'a self, f: F) -> impl Future<Output = anyhow::Result<R>> + 'a
    where
        F: Future<Output = anyhow::Result<R>> + 'a,
    {
        // We avoid using `async fn` or `async move` here to avoid doubling the
        // future size. See https://github.com/rust-lang/rust/issues/62958
        self.started.fetch_add(1, Ordering::Relaxed);
        f.map(|result| {
            (if result.is_ok() {
                &self.finished_successfully
            } else {
                &self.finished_with_error
            })
            .fetch_add(1, Ordering::Relaxed);
            result
        })
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use buck2_common::executor_config::RemoteExecutorUseCase;
use derivative::Derivative;
use dupe::Dupe;
use futures::future;
use remote_execution::TActionResult2;

use crate::digest_config::DigestConfig;
use crate::execute::output::ReStdStream;
use crate::re::manager::ManagedRemoteExecutionClient;

#[derive(Derivative, Clone)]
#[derivative(Debug)]
pub struct RemoteCommandStdStreams {
    #[derivative(Debug = "ignore")]
    client: ManagedRemoteExecutionClient,
    #[derivative(Debug = "ignore")]
    use_case: RemoteExecutorUseCase,
    #[derivative(Debug = "ignore")]
    digest_config: DigestConfig,
    stdout: ReStdStream,
    stderr: ReStdStream,
}

impl RemoteCommandStdStreams {
    pub fn new(
        action_result: &TActionResult2,
        client: &ManagedRemoteExecutionClient,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> Self {
        let stdout = ReStdStream::new(
            action_result.stdout_raw.clone(),
            action_result.stdout_digest.clone(),
        );
        let stderr = ReStdStream::new(
            action_result.stderr_raw.clone(),
            action_result.stderr_digest.clone(),
        );

        Self {
            client: client.dupe(),
            use_case,
            digest_config,
            stdout,
            stderr,
        }
    }

    pub async fn prefetch_lossy_stderr(mut self) -> Self {
        self.stderr
            .prefetch_lossy(&self.client, self.use_case, self.digest_config)
            .await;
        self
    }

    pub(crate) async fn to_lossy_stdout(&self) -> String {
        self.stdout
            .to_lossy(&self.client, self.use_case, self.digest_config)
            .await
    }

    pub(crate) async fn to_lossy_stderr(&self) -> String {
        self.stderr
            .to_lossy(&self.client, self.use_case, self.digest_config)
            .await
    }

    pub(crate) async fn into_stdout_stderr_bytes(self) -> anyhow::Result<(Vec<u8>, Vec<u8>)> {
        future::try_join(
            self.stdout
                .into_bytes(&self.client, self.use_case, self.digest_config),
            self.stderr
                .into_bytes(&self.client, self.use_case, self.digest_config),
        )
        .await
    }

    pub(crate) fn use_case(&self) -> RemoteExecutorUseCase {
        self.use_case
    }

    pub(crate) fn into_stdout_stderr(self) -> (ReStdStream, ReStdStream) {
        (self.stdout, self.stderr)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashSet;
use std::str::FromStr;
use std::sync::Arc;

use anyhow::Context;
use buck2_common::cas_digest::TrackedCasDigest;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_common::file_ops::FileDigest;
use buck2_common::file_ops::FileDigestKind;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_core::directory::DirectoryEntry;
use buck2_core::directory::DirectoryIterator;
use buck2_core::directory::FingerprintedDirectory;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::soft_error;
use chrono::Duration;
use chrono::Utc;
use futures::FutureExt;
use gazebo::prelude::*;
use remote_execution::GetDigestsTtlRequest;
use remote_execution::InlinedBlobWithDigest;
use remote_execution::NamedDigest;
use remote_execution::REClient;
use remote_execution::REClientError;
use remote_execution::TCode;
use remote_execution::TDigest;
use remote_execution::UploadRequest;

use crate::digest::CasDigestFromReExt;
use crate::digest::CasDigestToReExt;
use crate::digest_config::DigestConfig;
use crate::directory::ActionDirectoryMember;
use crate::directory::ActionFingerprintedDirectory;
use crate::directory::ActionImmutableDirectory;
use crate::directory::ReDirectorySerializer;
use crate::execute::blobs::ActionBlobs;
use crate::materialize::materializer::ArtifactNotMaterializedReason;
use crate::materialize::materializer::CasDownloadInfo;
use crate::materialize::materializer::Materializer;
use crate::re::metadata::RemoteExecutionMetadataExt;

#[derive(Clone, Debug, Default)]
pub struct UploadStats {
    pub bytes_uploaded: u64,
    pub digests_uploaded: u64,
}

pub struct Uploader {}

impl Uploader {
    async fn find_missing<'a>(
        client: &REClient,
        input_dir: &'a ActionImmutableDirectory,
        blobs: &'a ActionBlobs,
        use_case: &RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> anyhow::Result<(
        Vec<InlinedBlobWithDigest>,
        HashSet<&'a TrackedCasDigest<FileDigestKind>>,
    )> {
        // RE mentions they usually take 5-10 minutes of leeway so we mirror this here.
        let now = Utc::now();
        let ttl_wanted = if buck2_core::is_open_source() {
            1
        } else {
            600i64
        };
        let ttl_deadline = now + Duration::seconds(ttl_wanted);

        // See if anything needs uploading
        let mut input_digests = blobs.keys().collect::<HashSet<_>>();
        let digest_ttls = {
            // Collect the digests we need to upload
            for entry in input_dir.fingerprinted_unordered_walk().without_paths() {
                let digest = match entry {
                    DirectoryEntry::Dir(d) => d.fingerprint(),
                    DirectoryEntry::Leaf(ActionDirectoryMember::File(f)) => &f.digest,
                    DirectoryEntry::Leaf(..) => continue,
                };

                if digest.expires() <= ttl_deadline {
                    input_digests.insert(digest);
                }
            }

            let root_dir_digest = input_dir.fingerprint();
            if root_dir_digest.expires() <= ttl_deadline {
                input_digests.insert(root_dir_digest);
            }

            // Find out which ones are missing
            let request = GetDigestsTtlRequest {
                digests: input_digests.iter().map(|d| d.to_re()).collect(),
                ..Default::default()
            };
            client
                .get_digests_ttl(use_case.metadata(), request)
                .boxed()
                .await?
                .digests_with_ttl
        };

        let mut upload_blobs = Vec::new();
        let mut missing_digests = HashSet::new();
        add_injected_missing_digests(&input_digests, &mut missing_digests)?;

        let mut input_digests = input_digests.into_iter().collect::<Vec<_>>();
        input_digests.sort();

        let mut digest_ttls = digest_ttls.into_try_map(|d| {
            anyhow::Ok((FileDigest::from_re(&d.digest, digest_config)?, d.ttl))
        })?;
        digest_ttls.sort();

        if input_digests.len() != digest_ttls.len() {
            return Err(anyhow::anyhow!(
                "Invalid response from get_digests_ttl: expected {}, got {} digests",
                input_digests.len(),
                digest_ttls.len()
            ));
        }

        // Find the blobs that need to be uploaded

        for (digest, (matching_digest, digest_ttl)) in
            input_digests.into_iter().zip(digest_ttls.into_iter())
        {
            if *digest.data() != matching_digest {
                return Err(anyhow::anyhow!("Invalid response from get_digests_ttl"));
            }

            if digest_ttl <= ttl_wanted {
                tracing::debug!(digest=%digest, ttl=digest_ttl, "Mark for upload");

                match blobs.get(digest) {
                    Some(blob) => {
                        upload_blobs.push(InlinedBlobWithDigest {
                            blob: blob.clone(),
                            digest: digest.to_re(),
                            ..Default::default()
                        });
                    }
                    None => {
                        missing_digests.insert(digest);
                    }
                }
            } else {
                tracing::debug!(digest=%digest, ttl=digest_ttl, "Not uploading");
                let ttl = Duration::seconds(digest_ttl);
                digest.update_expires(now + ttl);
            }
        }

        Ok((upload_blobs, missing_digests))
    }

    pub async fn upload(
        fs: &ProjectRoot,
        client: &REClient,
        materializer: &Arc<dyn Materializer>,
        dir_path: &ProjectRelativePath,
        input_dir: &ActionImmutableDirectory,
        blobs: &ActionBlobs,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> anyhow::Result<UploadStats> {
        let (mut upload_blobs, mut missing_digests) =
            Self::find_missing(client, input_dir, blobs, &use_case, digest_config).await?;

        if upload_blobs.is_empty() && missing_digests.is_empty() {
            return Ok(UploadStats::default());
        }

        // Find the file paths and directory blobs that need to be uploaded
        let mut upload_files = Vec::new();

        // Track what files should be materialized before we upload.
        let mut paths_to_materialize = Vec::new();

        if !missing_digests.is_empty() {
            let mut upload_file_paths = Vec::new();
            let mut upload_file_digests = Vec::new();

            {
                let mut walk = input_dir.fingerprinted_unordered_walk();
                while let Some((path, entry)) = walk.next() {
                    let digest = match entry {
                        DirectoryEntry::Dir(d) => d.fingerprint(),
                        DirectoryEntry::Leaf(ActionDirectoryMember::File(f)) => &f.digest,
                        DirectoryEntry::Leaf(..) => continue,
                    };

                    if !missing_digests.remove(digest) {
                        continue;
                    }

                    match entry {
                        DirectoryEntry::Dir(d) => {
                            upload_blobs.push(directory_to_blob(d));
                        }
                        DirectoryEntry::Leaf(ActionDirectoryMember::File(..)) => {
                            upload_file_paths.push(dir_path.join(path.get()));
                            upload_file_digests.push(digest.to_re());
                        }
                        DirectoryEntry::Leaf(..) => unreachable!(), // TODO: Better representation of this.
                    };
                }
            }

            if missing_digests.remove(input_dir.fingerprint()) {
                upload_blobs.push(directory_to_blob(input_dir));
            }

            assert!(
                missing_digests.is_empty(),
                "Expected a path to be found for every digest, traversal code is inconsistent. Left with {:?}.",
                missing_digests
            );

            // Get the real path of the files we are going to upload.
            // This needs to be done because we could have A copied to B, and
            // we are asked to upload B. But since we defer local copies until
            // it's actually needed for a local run, B might not have been
            // copied yet (or ever), so we should upload A directly instead.
            let upload_file_paths = materializer
                .get_materialized_file_paths(upload_file_paths)
                .await?;

            for (name, digest) in upload_file_paths
                .into_iter()
                .zip(upload_file_digests.into_iter())
            {
                match name {
                    Ok(name) => {
                        upload_files.push(NamedDigest {
                            name: fs.resolve(&name).as_maybe_relativized_str()?.to_owned(),
                            digest,
                            ..Default::default()
                        });
                    }
                    Err(
                        ref err @ ArtifactNotMaterializedReason::RequiresCasDownload {
                            ref entry,
                            ref info,
                            ..
                        },
                    ) => {
                        if let DirectoryEntry::Leaf(ActionDirectoryMember::File(ref file)) =
                            entry.as_ref()
                        {
                            // NOTE: find_missing has negative caching, so when we query to know if an
                            // artifact was uploaded, if it was the result of an action we just ran, it
                            // won't be here. On the flip side, if a digest has been in the CAS for
                            // a very long time, it might have expired.
                            if file.digest.to_re() == digest {
                                if should_error_for_missing_digest(info) {
                                    soft_error!(
                                        "cas_missing_fatal",
                                        anyhow::anyhow!(
                                            "{} missing (origin: {})",
                                            file.digest,
                                            info.origin.as_display_for_not_found(),
                                        ),
                                        daemon_in_memory_state_is_corrupted: true,
                                        action_cache_is_corrupted: info.origin.guaranteed_by_action_cache()
                                    )?;

                                    return Err(anyhow::anyhow!(
                                        "Your build requires an artifact that has expired in the RE CAS \
                                        and Buck does not have it. This likely happened because your Buck daemon \
                                        has been online for a long time. This error is currently unrecoverable. \
                                        To proceed, you should restart Buck using `buck2 killall`. \
                                        Debug information: {:#}",
                                        err
                                    ));
                                }

                                soft_error!(
                                    "cas_missing",
                                    anyhow::anyhow!(
                                        "{} (expires = {}) is missing in the CAS but expected to exist as per: {:#}",
                                        file.digest,
                                        file.digest.expires(),
                                        err
                                    )
                                )?;

                                continue;
                            }
                        }

                        return Err(error_for_missing_file(&digest, err));
                    }
                    Err(ArtifactNotMaterializedReason::RequiresMaterialization { path }) => {
                        upload_files.push(NamedDigest {
                            name: fs.resolve(&path).as_maybe_relativized_str()?.to_owned(),
                            digest,
                            ..Default::default()
                        });
                        paths_to_materialize.push(path);
                    }
                    Err(
                        ref err @ ArtifactNotMaterializedReason::DeferredMaterializerCorruption {
                            ..
                        },
                    ) => {
                        return Err(error_for_missing_file(&digest, err));
                    }
                };
            }
        }

        if !paths_to_materialize.is_empty() {
            materializer
                .ensure_materialized(paths_to_materialize)
                .await
                .context("Error materializing paths for upload")?;
        }

        // Compute stats of digests we're about to upload so we can report them
        // to the span end event of this stage of execution.
        let stats = {
            let named_digest_byte_count: u64 = upload_files
                .iter()
                .map(|nd| {
                    let byte_count: u64 = nd.digest.size_in_bytes.try_into().unwrap_or_default();
                    byte_count
                })
                .sum();
            let blob_byte_count: u64 = upload_blobs
                .iter()
                .map(|blob| {
                    let byte_count: u64 = blob.digest.size_in_bytes.try_into().unwrap_or_default();
                    byte_count
                })
                .sum();

            UploadStats {
                digests_uploaded: (upload_files.len() + upload_blobs.len()) as u64,
                bytes_uploaded: named_digest_byte_count + blob_byte_count,
            }
        };

        // Upload
        let upload_res = if !upload_files.is_empty() || !upload_blobs.is_empty() {
            client
                .upload(
                    use_case.metadata(),
                    UploadRequest {
                        files_with_digest: Some(upload_files),
                        inlined_blobs_with_digest: Some(upload_blobs),
                        // all find missing checks are done previously
                        // and we can skip them and upload all digests
                        upload_only_missing: false,
                        ..Default::default()
                    },
                )
                .boxed()
                .await
                .map(|_| ())
        } else {
            Ok(())
        };

        if let Err(e) = upload_res.as_ref() {
            if let Some(re_client_error) = e.downcast_ref::<REClientError>() {
                if re_client_error.code == TCode::INVALID_ARGUMENT {
                    return Err(anyhow::anyhow!(
                        "RE Upload failed. It looks like you might have modified files while the build \
                        was in progress. Retry your build to proceed. Debug information: {:#}",
                        e
                    ));
                }
            }
        }

        upload_res.context("RE: upload")?;

        Ok(stats)
    }
}

fn should_error_for_missing_digest(info: &CasDownloadInfo) -> bool {
    // RE sometimes reports things that exist as missing. We don't fully understand why at this
    // time and this is being investigated, but we know that RE normally ensures that anything it
    // returns to us will last for another 6 hours at least. So, we silence all errors when the
    // download info is less than 5 hours, because we know those are most likely bogus. This
    // basically means that after 5 hours we might tell the user to restart even though they don't
    // need to. However, the alternative is confused users who see errors after a couple days (we
    // had 3 reports of this in a week), so for now we do accept some false positives (when RE
    // tells us a digest doesn't exist even though it does) in order to provide better UX when we
    // hit a true positive.
    if let Some(age) = info.action_age() {
        age >= Duration::seconds(3600 * 5)
    } else {
        true
    }
}

fn directory_to_blob<D>(d: &D) -> InlinedBlobWithDigest
where
    D: ActionFingerprintedDirectory + ?Sized,
{
    InlinedBlobWithDigest {
        digest: d.fingerprint().to_re(),
        blob: ReDirectorySerializer::serialize_entries(d.fingerprinted_entries()),
        ..Default::default()
    }
}

fn error_for_missing_file(
    digest: &TDigest,
    cause: &ArtifactNotMaterializedReason,
) -> anyhow::Error {
    anyhow::anyhow!(
        "Action execution requires artifact `{}` but the materializer did not return a matching \
        file for this path. This error is unrecoverable and you should restart Buck using \
        `buck2 killall`. We would appreciate a bug report. Debug information: {:#}",
        digest,
        cause,
    )
}

/// This is used for tests. We allow an environment variable to be set to report that some digests
/// are _always_ missing if they are required. This lets us test our upload paths more easily.
fn add_injected_missing_digests<'a>(
    input_digests: &HashSet<&'a TrackedFileDigest>,
    missing_digests: &mut HashSet<&'a TrackedFileDigest>,
) -> anyhow::Result<()> {
    fn convert_digests(val: &str) -> anyhow::Result<Vec<FileDigest>> {
        val.split(' ')
            .map(|digest| {
                let digest = TDigest::from_str(digest)
                    .with_context(|| format!("Invalid digest: `{}`", digest))?;
                // This code does not run in a test but it is only used for testing.
                let digest = FileDigest::from_re(&digest, DigestConfig::testing_default())?;
                anyhow::Ok(digest)
            })
            .collect()
    }

    static INJECTED_DIGESTS: EnvHelper<Vec<FileDigest>> =
        EnvHelper::with_converter("BUCK2_TEST_INJECTED_MISSING_DIGESTS", convert_digests);

    if let Some(digests) = INJECTED_DIGESTS.get()? {
        for d in digests {
            if let Some(i) = input_digests.get(d) {
                missing_digests.insert(i);
            }
        }
    }

    Ok(())
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

#![feature(try_blocks)]
#![feature(box_patterns)]
#![feature(try_trait_v2)]

pub mod executors;
pub mod low_pass_filter;
pub mod materializers;
pub mod re;