/*
 * 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.
 */

// Eden's Thrift API does sometime want &Vec<...>.
#![allow(clippy::useless_vec)]

use std::collections::BTreeMap;
use std::fmt::Display;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use allocative::Allocative;
use anyhow::Context as _;
use buck2_common::result::SharedResult;
use buck2_core;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::paths::abs_path::AbsPath;
use dupe::Dupe;
use edenfs::client::EdenService;
use edenfs::errors::eden_service::ListMountsError;
use edenfs::types::BinaryHash;
use edenfs::types::EdenErrorType;
use edenfs::types::FileAttributeData;
use edenfs::types::FileAttributeDataOrErrorV2;
use edenfs::types::FileAttributeDataV2;
use edenfs::types::MountState;
use edenfs::types::PathString;
use edenfs::types::SourceControlType;
use fb303_core::client::BaseService;
use fbinit::FacebookInit;
use futures::future::BoxFuture;
use futures::future::Future;
use futures::future::FutureExt;
use futures::future::Shared;
use parking_lot::Mutex;
use serde::Deserialize;
use sorted_vector_map::SortedVectorMap;
use thiserror::Error;
use tokio::sync::Semaphore;

#[derive(Allocative)]
pub struct EdenConnectionManager {
    connector: EdenConnector,
    connection: Mutex<EdenConnection>,
    /// Eden has limits on concurrency and will return server overloaded (or timeout) errors if we
    /// send too many. Experimentally, even for large builds (see details in D36136516), we don't
    /// much performance improvement beyond 2K concurrent requests, regardless of whether Eden has
    /// a fast or slow connectin to source control, a warm cache or not, and a lot of CPU available
    /// to run or not.
    #[allocative(skip)]
    semaphore: Semaphore,
}

#[derive(Deserialize, Debug)]
struct Config {
    root: String,
    socket: String,
}

#[derive(Deserialize, Debug)]
struct EdenConfig {
    #[serde(rename = "Config")]
    config: Config,
}

impl EdenConnectionManager {
    pub fn new(
        fb: FacebookInit,
        root: &AbsNormPath,
        semaphore: Semaphore,
    ) -> anyhow::Result<Option<Self>> {
        let eden_root = root.as_abs_path().join(".eden");
        if !eden_root.exists() {
            return Ok(None);
        }
        let connector = Self::get_eden_connector(fb, &eden_root)?;

        let connection = Mutex::new(EdenConnection {
            epoch: 0,
            client: connector.connect(),
        });

        Ok(Some(Self {
            connector,
            connection,
            semaphore,
        }))
    }

    fn get_eden_connector(fb: FacebookInit, eden_root: &AbsPath) -> anyhow::Result<EdenConnector> {
        // Based off of how watchman picks up the config: fbcode/watchman/watcher/eden.cpp:138
        if cfg!(windows) {
            let config_path = eden_root.join("config");
            let config_contents = fs_util::read_to_string(config_path)?;
            let config: EdenConfig = toml::from_str(&config_contents)?;
            let root = Arc::new(config.config.root);
            let socket = PathBuf::from(config.config.socket);
            Ok(EdenConnector { fb, root, socket })
        } else {
            let root = fs_util::read_link(eden_root.join("root"))?
                .to_str()
                .context("Eden root is not UTF-8")?
                .to_owned();
            let root = Arc::new(root);
            let socket = fs_util::read_link(eden_root.join("socket"))?;
            Ok(EdenConnector { fb, root, socket })
        }
    }

    pub fn get_root(&self) -> &str {
        &self.connector.root
    }

    pub fn get_mount_point(&self) -> Vec<u8> {
        self.connector.root.as_bytes().to_vec()
    }

    /// Returns a string like "20220102-030405", assuming this is a release version. This is
    /// pattern-matched off of what the Eden CLI does.
    pub async fn get_eden_version(&self) -> anyhow::Result<Option<String>> {
        let fb303 = self.connector.connect_fb303()?;
        let values = fb303.getRegexExportedValues("^build_.*").await?;

        fn join_version(values: &BTreeMap<String, String>) -> Option<String> {
            let version = values.get("build_package_version")?;
            let release = values.get("build_package_release")?;
            if version.is_empty() || release.is_empty() {
                return None;
            }
            Some(format!("{}-{}", version, release))
        }

        Ok(join_version(&values))
    }

    pub async fn with_eden<F, Fut, T, E>(&self, f: F) -> Result<T, ConnectAndRequestError<E>>
    where
        F: Fn(&(dyn EdenService + Send + Sync)) -> Fut,
        Fut: Future<Output = Result<T, E>>,
        E: HasErrorHandlingStrategy + Display,
    {
        const MAX_ATTEMPTS: usize = 3;

        let mut connection = (*self.connection.lock()).clone();
        let mut attempts = 0;

        let _permit = self
            .semaphore
            .acquire()
            .await
            .expect("Eden I/O semaphore is never closed");

        loop {
            attempts += 1;

            let res = async {
                let client = connection
                    .client
                    .clone()
                    .await
                    .map_err(|e| ConnectAndRequestError::ConnectionError(e.into()))?;

                f(client.as_ref())
                    .await
                    .map_err(|e| ConnectAndRequestError::RequestError(e))
            }
            .await;

            let err = match res {
                Ok(res) => break Ok(res),
                Err(e) => e,
            };

            match err.error_handling_strategy() {
                ErrorHandlingStrategy::Reconnect => {
                    // Our connection to Eden broke. This typically means Eden restarted. Just
                    // reconnect.
                    tracing::info!("Reconnecting to Eden after: {:#}", err);
                    let mut guard = self.connection.lock();
                    if guard.epoch == connection.epoch {
                        guard.client = self.connector.connect();
                        guard.epoch += 1;
                    }
                    connection = (*guard).clone();
                }
                ErrorHandlingStrategy::Retry => {
                    // Our request failed but needs retrying.
                    tracing::info!("Retrying Eden request after: {:#}", err);
                }
                ErrorHandlingStrategy::Abort => {
                    break Err(err);
                }
            };

            if attempts > MAX_ATTEMPTS {
                break Err(err);
            }
        }
    }
}

/// A (potentially pending) Eden client.
type EdenClientFuture =
    Shared<BoxFuture<'static, SharedResult<Arc<dyn EdenService + Send + Sync>>>>;

/// An Eden client and an epoch to keep track of reconnections.
#[derive(Clone, Allocative)]
struct EdenConnection {
    /// This stats at zero and increments every time we reconnect. We use this to keep track of
    /// whether another client already recycled the connection when we need to reconnect.
    epoch: usize,
    #[allocative(skip)]
    client: EdenClientFuture,
}

/// A factory for Eden clients.
#[derive(Allocative)]
struct EdenConnector {
    #[allocative(skip)]
    fb: FacebookInit,
    root: Arc<String>,
    socket: PathBuf,
}

impl EdenConnector {
    fn connect(&self) -> EdenClientFuture {
        let socket = self.socket.clone();
        let fb = self.fb;
        let root = self.root.dupe();

        tokio::task::spawn(async move {
            tracing::info!("Creating a new Eden connection via `{}`", socket.display());
            let eden: anyhow::Result<Arc<dyn EdenService + Send + Sync>>;

            #[cfg(fbcode_build)]
            {
                eden = fbcode::thrift_builder(fb, socket)?
                    .build_client(::edenfs::client::make_EdenService);
            }

            #[cfg(not(fbcode_build))]
            {
                let _ignored = fb;
                let _ignored = socket;
                eden = Err(anyhow::anyhow!("Eden I/O is not available in Cargo builds"))
            }

            let eden = eden.context("Error constructing Eden client")?;

            wait_until_mount_is_ready(eden.as_ref(), &root).await?;

            Ok(eden)
        })
        .map(|r| match r {
            Ok(r) => r,
            Err(e) => Err(e.into()), // Turn the JoinError into a SharedError.
        })
        .boxed()
        .shared()
    }

    #[cfg(fbcode_build)]
    fn connect_fb303(&self) -> anyhow::Result<Arc<dyn BaseService + Send + Sync>> {
        fbcode::thrift_builder(self.fb, &self.socket)?
            .build_client(::fb303_core::client::make_BaseService)
    }

    #[cfg(not(fbcode_build))]
    fn connect_fb303(&self) -> anyhow::Result<Arc<dyn BaseService + Send + Sync>> {
        Err(anyhow::anyhow!("Eden I/O is not available in Cargo builds"))
    }
}

#[derive(Error, Debug)]
#[error("Mount never became ready: `{}`", self.mount)]
struct MountNeverBecameReady {
    mount: Arc<String>,
}

/// Delay until a mount becomes ready (up to 10 seconds).
async fn wait_until_mount_is_ready(
    eden: &(dyn EdenService + Send + Sync),
    root: &Arc<String>,
) -> anyhow::Result<()> {
    let mut interval = tokio::time::interval(Duration::from_secs(1));
    interval.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip);

    for _ in 0..10 {
        interval.tick().await;
        match is_mount_ready(eden, root).await {
            Ok(true) => return Ok(()),
            Ok(false) => {
                // Fallthrough to keep going
            }
            Err(e) if e.error_handling_strategy() == ErrorHandlingStrategy::Retry => {
                // Fallthrough to keep going
            }
            Err(e) => return Err(e.into()),
        }
    }

    Err(MountNeverBecameReady { mount: root.dupe() }.into())
}

#[derive(Error, Debug)]
pub enum IsMountReadyError {
    #[error("Mount does not exist in Eden: `{}`", .mount)]
    MountDoesNotExist { mount: Arc<String> },
    #[error(transparent)]
    RequestError(ListMountsError),
}

/// Check if a given mount is ready.
async fn is_mount_ready(
    eden: &(dyn EdenService + Send + Sync),
    root: &Arc<String>,
) -> Result<bool, IsMountReadyError> {
    let mounts = eden
        .listMounts()
        .await
        .map_err(IsMountReadyError::RequestError)?;

    for mount in mounts {
        if mount.mountPoint == root.as_bytes() {
            return Ok(mount.state == MountState::RUNNING);
        }
    }

    Err(IsMountReadyError::MountDoesNotExist { mount: root.dupe() })
}

#[derive(Error, Debug)]
pub enum ConnectAndRequestError<E> {
    #[error(transparent)]
    ConnectionError(anyhow::Error),

    #[error(transparent)]
    RequestError(E),
}

#[derive(Copy, Clone, Dupe, PartialEq, Eq)]
pub enum ErrorHandlingStrategy {
    Reconnect,
    Retry,
    Abort,
}

pub trait HasErrorHandlingStrategy {
    fn error_handling_strategy(&self) -> ErrorHandlingStrategy;
}

impl<E: HasErrorHandlingStrategy> HasErrorHandlingStrategy for ConnectAndRequestError<E> {
    fn error_handling_strategy(&self) -> ErrorHandlingStrategy {
        match self {
            Self::ConnectionError(..) => ErrorHandlingStrategy::Reconnect,
            Self::RequestError(e) => e.error_handling_strategy(),
        }
    }
}

impl HasErrorHandlingStrategy for IsMountReadyError {
    fn error_handling_strategy(&self) -> ErrorHandlingStrategy {
        match self {
            Self::MountDoesNotExist { .. } => ErrorHandlingStrategy::Abort,
            Self::RequestError(e) => e.error_handling_strategy(),
        }
    }
}

macro_rules! impl_has_error_handling_strategy {
    ($err: ident) => {
        impl HasErrorHandlingStrategy for ::edenfs::errors::eden_service::$err {
            fn error_handling_strategy(&self) -> ErrorHandlingStrategy {
                match self {
                    Self::ThriftError(..) => ErrorHandlingStrategy::Reconnect,
                    Self::ApplicationException(..) => ErrorHandlingStrategy::Retry,
                    Self::ex(..) => ErrorHandlingStrategy::Abort,
                }
            }
        }
    };
}

impl_has_error_handling_strategy!(GetAttributesFromFilesError);
impl_has_error_handling_strategy!(GetAttributesFromFilesV2Error);
impl_has_error_handling_strategy!(GlobFilesError);
impl_has_error_handling_strategy!(ListMountsError);
impl_has_error_handling_strategy!(SynchronizeWorkingCopyError);
impl_has_error_handling_strategy!(SetPathObjectIdError);
impl_has_error_handling_strategy!(RemoveRecursivelyError);
impl_has_error_handling_strategy!(EnsureMaterializedError);
impl_has_error_handling_strategy!(ReaddirError);

#[derive(Debug, Error)]
pub enum EdenError {
    #[error("Eden POSIX error (code = {}): {}", .code, .error.message)]
    PosixError {
        error: edenfs::types::EdenError,
        code: i32,
    },

    #[error("Eden service error: {}", .error.message)]
    ServiceError { error: edenfs::types::EdenError },

    #[error("Eden returned an unexpected field: {}", .field)]
    UnknownField { field: i32 },
}

impl From<edenfs::types::EdenError> for EdenError {
    fn from(error: edenfs::types::EdenError) -> Self {
        if error.errorType == EdenErrorType::POSIX_ERROR {
            if let Some(error_code) = error.errorCode {
                return Self::PosixError {
                    error,
                    code: error_code,
                };
            }
        }

        Self::ServiceError { error }
    }
}

pub trait EdenDataIntoResult {
    type Data;

    fn into_result(self) -> Result<Self::Data, EdenError>;
}

macro_rules! impl_eden_data_into_result {
    ($typ: ident, $data: ty, $ok_variant: ident) => {
        impl EdenDataIntoResult for ::edenfs::types::$typ {
            type Data = $data;

            fn into_result(self) -> Result<Self::Data, EdenError> {
                match self {
                    Self::$ok_variant(data) => Ok(data),
                    Self::error(e) => Err(e.into()),
                    Self::UnknownField(field) => Err(EdenError::UnknownField { field }),
                }
            }
        }
    };
}

impl_eden_data_into_result!(
    SourceControlTypeOrError,
    SourceControlType,
    sourceControlType
);

impl_eden_data_into_result!(FileAttributeDataOrError, FileAttributeData, data);

impl_eden_data_into_result!(
    FileAttributeDataOrErrorV2,
    FileAttributeDataV2,
    fileAttributeData
);

impl_eden_data_into_result!(SizeOrError, i64, size);

impl_eden_data_into_result!(Sha1OrError, BinaryHash, sha1);

impl_eden_data_into_result!(Blake3OrError, BinaryHash, blake3);

impl_eden_data_into_result!(
    DirListAttributeDataOrError,
    SortedVectorMap<PathString, FileAttributeDataOrErrorV2>,
    dirListAttributeData
);

#[cfg(fbcode_build)]
mod fbcode {
    use std::path::Path;

    use super::*;

    pub fn thrift_builder<P: AsRef<Path>>(
        fb: FacebookInit,
        socket: P,
    ) -> anyhow::Result<::thriftclient::ThriftChannelBuilder> {
        // NOTE: This timeout is absurdly high, but bear in mind that what we're
        // "comparing" to is a FS call that has no timeouts at all.
        const THRIFT_TIMEOUT_MS: u32 = 120_000;

        Ok(::thriftclient::ThriftChannelBuilder::from_path(fb, socket)?
            .with_conn_timeout(THRIFT_TIMEOUT_MS)
            .with_recv_timeout(THRIFT_TIMEOUT_MS)
            .with_secure(false))
    }
}


/*
 * 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.
 */

// Eden's Thrift API does sometime want &Vec<...>.
#![allow(clippy::useless_vec)]

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_common::cas_digest::CasDigestConfig;
use buck2_common::file_ops::FileDigest;
use buck2_common::file_ops::FileMetadata;
use buck2_common::file_ops::FileType;
use buck2_common::file_ops::RawDirEntry;
use buck2_common::file_ops::RawPathMetadata;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_common::io::fs::FsIoProvider;
use buck2_common::io::IoProvider;
use buck2_core;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_core::io_counters::IoCounterKey;
use buck2_core::soft_error;
use compact_str::CompactString;
use dupe::Dupe;
use edenfs::types::FileAttributes;
use edenfs::types::GetAttributesFromFilesParams;
use edenfs::types::ReaddirParams;
use edenfs::types::SourceControlType;
use edenfs::types::SyncBehavior;
use edenfs::types::SynchronizeWorkingCopyParams;
use fbinit::FacebookInit;
use libc::EINVAL;
use libc::EISDIR;
use libc::ENOENT;
use libc::ENOTDIR;
use tokio::sync::Semaphore;

use crate::connection::EdenConnectionManager;
use crate::connection::EdenDataIntoResult;
use crate::connection::EdenError;

#[derive(Allocative)]
pub struct EdenIoProvider {
    manager: EdenConnectionManager,
    fs: FsIoProvider,
    version: Version,
}

#[derive(Allocative, Copy, Clone, Dupe)]
enum Version {
    V2(Digest),
    V1,
}

#[derive(Allocative, Copy, Clone, Dupe)]
enum Digest {
    Sha1,
    Blake3Keyed,
}

impl EdenIoProvider {
    pub async fn new(
        fb: FacebookInit,
        fs: &ProjectRoot,
        cas_digest_config: CasDigestConfig,
        v2: bool,
    ) -> anyhow::Result<Option<Self>> {
        if cfg!(not(fbcode_build)) {
            tracing::warn!("Disabling Eden I/O: Cargo build detected");
            return Ok(None);
        }

        let (version, min_eden_version) = if cas_digest_config.source_files_config().allows_sha1() {
            if v2 {
                (Version::V2(Digest::Sha1), "20220905-214046")
            } else {
                (Version::V1, "20220720-094125")
            }
        } else if cas_digest_config
            .source_files_config()
            .allows_blake3_keyed()
        {
            if v2 {
                (Version::V2(Digest::Blake3Keyed), "20230612-191714")
            } else {
                return Err(anyhow::anyhow!(
                    "Enabling Blake3 for source files requires Eden I/O v2"
                ));
            }
        } else {
            tracing::warn!("Disabling Eden I/O: Digest config is not compatible with Eden version");
            return Ok(None);
        };

        static EDEN_SEMAPHORE: EnvHelper<usize> = EnvHelper::new("BUCK2_EDEN_SEMAPHORE");
        let eden_semaphore = EDEN_SEMAPHORE.get_copied()?.unwrap_or(2048);

        let manager =
            match EdenConnectionManager::new(fb, fs.root(), Semaphore::new(eden_semaphore))? {
                Some(manager) => manager,
                None => return Ok(None),
            };

        let eden_version = manager
            .get_eden_version()
            .await
            .context("Error querying Eden version")?;

        if let Some(eden_version) = &eden_version {
            if eden_version.as_str() < min_eden_version {
                tracing::warn!(
                    "Disabling Eden I/O: \
                    your Eden version ({}) is too old to support Eden I/O (minimum required: {}). \
                    Update Eden then restart Buck2 to use Eden I/O.",
                    eden_version,
                    min_eden_version
                );
                return Ok(None);
            }
        } else {
            tracing::warn!("You are using a development version of Eden, enabling Eden I/O");
        }

        Ok(Some(Self {
            manager,
            fs: FsIoProvider::new(fs.dupe(), cas_digest_config),
            version,
        }))
    }

    async fn read_path_metadata_if_exists(
        &self,
        path: ProjectRelativePathBuf,
    ) -> anyhow::Result<Option<RawPathMetadata<ProjectRelativePathBuf>>> {
        let _guard = IoCounterKey::StatEden.guard();

        let requested_attributes = i64::from(
            i32::from(FileAttributes::SHA1_HASH)
                | i32::from(FileAttributes::FILE_SIZE)
                | i32::from(FileAttributes::SOURCE_CONTROL_TYPE),
        );

        let params = GetAttributesFromFilesParams {
            mountPoint: self.manager.get_mount_point(),
            paths: vec![path.to_string().into_bytes()],
            requestedAttributes: requested_attributes,
            sync: no_sync(),
            ..Default::default()
        };

        let attrs = self
            .manager
            .with_eden(|eden| {
                tracing::trace!("getAttributesFromFiles({})", path);
                eden.getAttributesFromFiles(&params)
            })
            .await?;

        match attrs
            .res
            .into_iter()
            .next()
            .context("Eden did not return file info")?
            .into_result()
        {
            Ok(data) => {
                tracing::debug!("getAttributesFromFiles({}): ok", path,);
                let digest = FileDigest::new_sha1(
                    data.sha1
                        .context("Eden did not return a sha1")?
                        .try_into()
                        .ok()
                        .context("Eden returned an invalid sha1")?,
                    data.fileSize
                        .context("Eden did not return a fileSize")?
                        .try_into()
                        .context("Eden returned an invalid fileSize")?,
                );

                // TODO (DigestConfig): Check that the config allows SHA1 earlier here.
                let digest = TrackedFileDigest::new(digest, self.fs.cas_digest_config());

                let is_executable = data.r#type.context("Eden did not return a type")?
                    == SourceControlType::EXECUTABLE_FILE;

                let meta = FileMetadata {
                    digest,
                    is_executable,
                };

                Ok(Some(RawPathMetadata::File(meta)))
            }
            Err(EdenError::PosixError { code, .. }) if code == EISDIR => {
                tracing::debug!("getAttributesFromFiles({}): EISDIR", path);
                Ok(Some(RawPathMetadata::Directory))
            }
            Err(EdenError::PosixError { code, .. }) if code == ENOENT => {
                tracing::debug!("getAttributesFromFiles({}): ENOENT", path);
                Ok(None)
            }
            Err(EdenError::PosixError { code, .. }) if code == EINVAL || code == ENOTDIR => {
                // If we get EINVAL it means the target wasn't a file, and since we know it
                // existed and it wasn't a dir, then that means it must be a symlink. If we get
                // ENOTDIR, that means we tried to traverse a path component that was a
                // symlink. In both cases, we need to both a) handle ExternalSymlink and b)
                // look through to the target, so we do that.
                // TODO: It would be better to read the link then ask Eden for the SHA1.
                tracing::debug!("getAttributesFromFiles({}): fallthrough", path);
                self.fs.read_path_metadata_if_exists(path).await
            }
            Err(err) => Err(err.into()),
        }
    }

    async fn read_path_metadata_if_exists_v2(
        &self,
        path: ProjectRelativePathBuf,
        digest: Digest,
    ) -> anyhow::Result<Option<RawPathMetadata<ProjectRelativePathBuf>>> {
        let _guard = IoCounterKey::StatEden.guard();

        let hash_attribute = match digest {
            Digest::Sha1 => FileAttributes::SHA1_HASH,
            // The only digests we get out of Eden are Blake3-keyed.
            Digest::Blake3Keyed => FileAttributes::BLAKE3_HASH,
        };

        let requested_attributes = i64::from(
            i32::from(hash_attribute)
                | i32::from(FileAttributes::FILE_SIZE)
                | i32::from(FileAttributes::SOURCE_CONTROL_TYPE),
        );

        let params = GetAttributesFromFilesParams {
            mountPoint: self.manager.get_mount_point(),
            paths: vec![path.to_string().into_bytes()],
            requestedAttributes: requested_attributes,
            sync: no_sync(),
            ..Default::default()
        };

        let attrs = self
            .manager
            .with_eden(|eden| {
                tracing::trace!("getAttributesFromFilesV2({})", path);
                eden.getAttributesFromFilesV2(&params)
            })
            .await?;

        match attrs
            .res
            .into_iter()
            .next()
            .context("Eden did not return file info")?
            .into_result()
        {
            Ok(data) => {
                let source_control_type = data
                    .sourceControlType
                    .context("Eden did not return a type")?
                    .into_result()
                    .context("Eden returned an error for sourceControlType")?;

                if source_control_type == SourceControlType::TREE {
                    return Ok(Some(RawPathMetadata::Directory));
                };

                if source_control_type == SourceControlType::SYMLINK {
                    // TODO: Ideally we would read *the link we just got*, instead of letting
                    // read_path_metadata_if_exists do the traversal of the whole path
                    let meta = self
                        .fs
                        .read_unchecked_symlink(path.clone())
                        .await
                        .with_context(|| {
                            format!(
                                "Eden returned that `{}` was a symlink, but it was not.  \
                                This path may have changed during the build",
                                path
                            )
                        })?;

                    return Ok(Some(meta));
                };

                let size = data
                    .size
                    .context("Eden did not return a size")?
                    .into_result()
                    .context("Eden returned an error for size")?
                    .try_into()
                    .context("Eden returned an invalid size")?;

                tracing::debug!("getAttributesFromFilesV2({}): ok", path,);
                let digest = match digest {
                    Digest::Sha1 => {
                        let sha1 = data
                            .sha1
                            .context("Eden did not return a sha1")?
                            .into_result()
                            .context("Eden returned an error for sha1")?
                            .try_into()
                            .ok()
                            .context("Eden returned an invalid sha1")?;
                        FileDigest::new_sha1(sha1, size)
                    }
                    Digest::Blake3Keyed => {
                        let blake3 = data
                            .blake3
                            .context("Eden did not return a blake3")?
                            .into_result()
                            .context("Eden returned an error for blake3")?
                            .try_into()
                            .ok()
                            .context("Eden returned an invalid blake3")?;
                        FileDigest::new_blake3_keyed(blake3, size)
                    }
                };

                let digest = TrackedFileDigest::new(digest, self.fs.cas_digest_config());

                let is_executable = source_control_type == SourceControlType::EXECUTABLE_FILE;

                let meta = FileMetadata {
                    digest,
                    is_executable,
                };

                Ok(Some(RawPathMetadata::File(meta)))
            }
            Err(EdenError::PosixError { code, .. }) if code == EISDIR => {
                tracing::debug!("getAttributesFromFilesV2({}): EISDIR", path);
                soft_error!(
                    "eden_io_eisdir",
                    anyhow::anyhow!("Eden returned EISDIR for {}", path)
                )?;
                Ok(Some(RawPathMetadata::Directory))
            }
            Err(EdenError::PosixError { code, .. }) if code == ENOENT => {
                tracing::debug!("getAttributesFromFilesV2({}): ENOENT", path);
                Ok(None)
            }
            Err(EdenError::PosixError { code, .. }) if code == EINVAL || code == ENOTDIR => {
                // If we get EINVAL it means the target wasn't a file, and since we know it
                // existed and it wasn't a dir, then that means it must be a symlink. If we get
                // ENOTDIR, that means we tried to traverse a path component that was a
                // symlink. In both cases, we need to both a) handle ExternalSymlink and b)
                // look through to the target, so we do that.
                tracing::debug!("getAttributesFromFilesV2({}): fallthrough", path);
                self.fs.read_path_metadata_if_exists(path).await
            }
            Err(err) => Err(err.into()),
        }
    }
}

#[async_trait]
impl IoProvider for EdenIoProvider {
    async fn read_file_if_exists(
        &self,
        path: ProjectRelativePathBuf,
    ) -> anyhow::Result<Option<String>> {
        self.fs.read_file_if_exists(path).await
    }

    async fn read_dir(&self, path: ProjectRelativePathBuf) -> anyhow::Result<Vec<RawDirEntry>> {
        let _guard = IoCounterKey::ReadDirEden.guard();

        let requested_attributes = i64::from(i32::from(FileAttributes::SOURCE_CONTROL_TYPE));

        let params = ReaddirParams {
            mountPoint: self.manager.get_mount_point(),
            // TODO(nga): this assumes eden mount point is the project root.
            //   Which is not the case for isolated test data directories for example.
            directoryPaths: vec![path.to_string().into_bytes()],
            requestedAttributes: requested_attributes,
            sync: no_sync(),
            ..Default::default()
        };

        let res = self
            .manager
            .with_eden(|eden| {
                tracing::trace!("readdir({})", path);
                eden.readdir(&params)
            })
            .await?
            .dirLists;

        let data = res
            .into_iter()
            .next()
            .context("Eden did not return a directory result")?
            .into_result()?;

        tracing::debug!("readdir({}): {} entries", path, data.len(),);

        let entries = data
            .into_iter()
            .map(|(file_name, attrs)| {
                let file_name =
                    CompactString::from_utf8(file_name).context("Filename is not UTF-8")?;

                let source_control_type = attrs
                    .into_result()?
                    .sourceControlType
                    .context("Missing sourceControlType")?
                    .into_result()?;

                let file_type = match source_control_type {
                    SourceControlType::TREE => FileType::Directory,
                    SourceControlType::REGULAR_FILE | SourceControlType::EXECUTABLE_FILE => {
                        FileType::File
                    }
                    SourceControlType::SYMLINK => FileType::Symlink,
                    _ => FileType::Unknown,
                };

                anyhow::Ok(RawDirEntry {
                    file_name,
                    file_type,
                })
            })
            .collect::<Result<Vec<_>, _>>()?;

        Ok(entries)
    }

    async fn read_path_metadata_if_exists(
        &self,
        path: ProjectRelativePathBuf,
    ) -> anyhow::Result<Option<RawPathMetadata<ProjectRelativePathBuf>>> {
        match self.version {
            Version::V2(digest) => Self::read_path_metadata_if_exists_v2(self, path, digest).await,
            Version::V1 => Self::read_path_metadata_if_exists(self, path).await,
        }
    }

    async fn settle(&self) -> anyhow::Result<()> {
        let _guard = IoCounterKey::EdenSettle.guard();

        let root = self.manager.get_mount_point();

        self.manager
            .with_eden(move |eden| {
                eden.synchronizeWorkingCopy(
                    &root,
                    &SynchronizeWorkingCopyParams {
                        sync: SyncBehavior {
                            syncTimeoutSeconds: None,
                            ..Default::default()
                        },
                        ..Default::default()
                    },
                )
            })
            .await
            .context("Error synchronizing Eden working copy")
    }

    fn name(&self) -> &'static str {
        "eden"
    }

    async fn eden_version(&self) -> anyhow::Result<Option<String>> {
        self.manager.get_eden_version().await
    }

    fn project_root(&self) -> &ProjectRoot {
        self.fs.project_root()
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }
}

/// We don't request sync on individual calls because we make a dedicated call for this before we
/// start an operation.
fn no_sync() -> SyncBehavior {
    // NOTE/ we'd rather make this a `const` if we could, but we can't, because
    // `Default::default()` is not a const-fn.
    SyncBehavior {
        syncTimeoutSeconds: Some(0),
        ..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.
 */

// This can't be built in our OSS implementation.
#![cfg(any(fbcode_build, cargo_internal_build))]

pub mod connection;
pub mod io_provider;


/*
 * 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 once_cell::sync::Lazy;
use uuid::Uuid;

#[derive(derive_more::Display)]
#[display(fmt = "{}", uuid.hyphenated())]
pub struct DaemonId {
    uuid: Uuid,
}

pub static DAEMON_UUID: Lazy<DaemonId> = Lazy::new(|| DaemonId {
    uuid: Uuid::new_v4(),
});

#[cfg(test)]
mod tests {
    use std::str::FromStr;

    use uuid::Uuid;

    use crate::daemon_id::DaemonId;

    #[test]
    fn test_display() {
        assert_eq!(
            "5621e477-ca64-4c68-888c-c47bdba0ff77",
            DaemonId {
                uuid: Uuid::from_str("5621e477-ca64-4c68-888c-c47bdba0ff77").unwrap()
            }
            .to_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.
 */

//! Mechanism for dispatching events. Where the rubber meets the road, for producers of events.
//!
//! The [`EventDispatcher`] is a type-erased, dupe-able container for a [`crate::EventSink`]. It is intended to be
//! liberally duplicated and passed around to the depths of buck2 so that consumers can insert events into it.

use std::cell::Cell;
use std::cell::UnsafeCell;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::pin::Pin;
use std::process::Stdio;
use std::sync::Arc;
use std::task;
use std::time::Duration;
use std::time::Instant;
use std::time::SystemTime;

use buck2_core::env_helper::EnvHelper;
use buck2_data::buck_event;
use buck2_data::instant_event::Data::HgInfo;
use buck2_data::span_end_event;
use buck2_data::span_start_event;
use buck2_data::MercurialInfo;
use buck2_data::SpanEndEvent;
use buck2_data::SpanStartEvent;
use buck2_util::process::background_command;
use buck2_wrapper_common::invocation_id::TraceId;
use dupe::Dupe;
use futures::Future;
use pin_project::pin_project;
use smallvec::SmallVec;

use crate::sink::null::NullEventSink;
use crate::span::SpanId;
use crate::BuckEvent;
use crate::ControlEvent;
use crate::EventSink;
use crate::EventSinkStats;

/// A type-erased and dupe-able container for EventSinks, containing some additional metadata useful for all events
/// emitted through the dispatcher.
#[derive(Clone, Dupe, Allocative)]
pub struct EventDispatcher {
    /// The Trace ID of the current trace.
    trace_id: TraceId,
    /// The sink to log events to.
    #[allocative(skip)] // TODO(nga): do not skip.
    sink: Arc<dyn EventSink>,
}

impl EventDispatcher {
    /// Creates a new Event Dispatcher using the given TraceId and logging events to the given sink.
    pub fn new<T: EventSink + 'static>(trace_id: TraceId, sink: T) -> EventDispatcher {
        EventDispatcher {
            trace_id,
            sink: Arc::new(sink),
        }
    }

    pub fn sink(&self) -> Arc<dyn EventSink> {
        self.sink.dupe()
    }

    /// Creates a new null Event Dispatcher that accepts events but does not write them anywhere.
    pub fn null() -> EventDispatcher {
        EventDispatcher {
            trace_id: TraceId::null(),
            sink: Arc::new(NullEventSink::new()),
        }
    }

    /// Creates a new null Event Dispatcher with trace ID that accepts events but does not write them anywhere.
    pub fn null_sink_with_trace(trace_id: TraceId) -> EventDispatcher {
        EventDispatcher {
            trace_id,
            sink: Arc::new(NullEventSink::new()),
        }
    }

    /// Emits an event annotated with the current trace ID.
    pub fn event<E: Into<buck_event::Data>>(&self, data: E) {
        self.event_with_span_id(data, None, current_span());
    }

    /// Emits an InstantEvent annotated with the current trace ID
    pub fn instant_event<E: Into<buck2_data::instant_event::Data>>(&self, data: E) {
        let instant = buck2_data::InstantEvent {
            data: Some(data.into()),
        };
        self.event_with_span_id(instant, None, current_span());
    }

    pub fn console_message(&self, message: String) {
        self.instant_event(buck2_data::ConsoleMessage { message })
    }

    fn event_with_span_id<E: Into<buck_event::Data>>(
        &self,
        data: E,
        span_id: Option<SpanId>,
        parent_id: Option<SpanId>,
    ) {
        let now = SystemTime::now();

        let event = BuckEvent::new(now, self.trace_id.dupe(), span_id, parent_id, data.into());
        self.sink.send(event);
    }

    pub fn control_event<E: Into<ControlEvent>>(&self, data: E) {
        self.sink.send_control(data.into());
    }

    // Logs mercurial data
    pub async fn instant_hg(&self) {
        // TODO use tokio/tokio::process::Command instead of command (see D29824148)
        let committed = background_command("hg").arg("status").arg("-mard").output();
        let log_changes = if let Ok(status) = committed {
            !status.stdout.is_empty()
        } else {
            return;
        };

        // TODO use `hg debughiddencommit` instead of `hg id`, `hg diff`, and `pastry`
        let hash = if let Ok(commit) = background_command("hg")
            .arg("--debug")
            .arg("id")
            .arg("-i")
            .output()
        {
            String::from_utf8_lossy(&commit.stdout).to_string()
        } else {
            return;
        };
        let mut pastry = "".to_owned();
        if log_changes {
            if let Ok(diff) = background_command("hg")
                .arg("diff")
                .arg("-r")
                .arg("master")
                .stdout(Stdio::piped())
                .spawn()
            {
                if let Some(d) = diff.stdout {
                    if let Ok(paste) = background_command("pastry").stdin(d).output() {
                        pastry = String::from_utf8_lossy(&paste.stdout).to_string();
                    }
                }
            }
        }
        self.instant_event(HgInfo(MercurialInfo {
            commit: hash,
            diff: pastry,
        }));
    }

    /// Introduces a new span and immediately fires the given start event. When the given synchronous function returns,
    /// the span is closed and the end event emitted.
    pub fn span<Start, End, F, R>(&self, start: Start, func: F) -> R
    where
        Start: Into<span_start_event::Data>,
        End: Into<span_end_event::Data>,
        F: FnOnce() -> (R, End),
    {
        let mut span = Span::start(self.dupe(), start);
        let (result, end) = span.call_in_span(func);
        span.end(end);
        result
    }

    /// Introduces a new span and immediately fires the given start event. When the given future resolves,  the span is
    /// closed and the event is emitted. This span is a "suspending span"; it is intended to suspend and resume whenever
    /// the future itself is suspended and resumed, respectively.
    ///
    /// TODO(swgillespie) actually implement suspend/resume
    pub fn span_async<Start, End, Fut, R>(&self, start: Start, fut: Fut) -> impl Future<Output = R>
    where
        Start: Into<span_start_event::Data>,
        End: Into<span_end_event::Data>,
        Fut: Future<Output = (R, End)>,
    {
        Span::start(self.dupe(), start).wrap_future(fut)
    }

    /// Creates a new span but does not enter it. Use `wrap_future` or `wrap_closure` to enter it,
    /// or create children spans via `create_child` and enter those.
    pub fn create_span<Start>(&self, start: Start) -> Span
    where
        Start: Into<span_start_event::Data>,
    {
        Span::start(self.dupe(), start)
    }

    /// Returns the traceid for this event dispatcher.
    pub fn trace_id(&self) -> &TraceId {
        &self.trace_id
    }

    /// Collect stats for the underlying sink.
    pub fn stats(&self) -> Option<EventSinkStats> {
        self.sink.stats()
    }
}

// While this may be simpler to implement via futures::poll_fn or some kind of async fn/block or similar, we need
// to ensure that we don't explode the size of our futures and so we implement the future explicitly.
#[pin_project]
struct SpanAsync<End, Fut, R>
where
    End: Into<span_end_event::Data>,
    Fut: Future<Output = (R, End)>,
{
    #[pin]
    fut: Fut,
    span: Span,
}

impl<End, Fut, R> Future for SpanAsync<End, Fut, R>
where
    End: Into<span_end_event::Data>,
    Fut: Future<Output = (R, End)>,
{
    type Output = R;

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
        let this = self.project();
        let span: &mut _ = this.span;
        match span.call_in_span(|| this.fut.poll(cx)) {
            task::Poll::Ready((v, end)) => {
                span.send(end.into());
                task::Poll::Ready(v)
            }
            task::Poll::Pending => task::Poll::Pending,
        }
    }
}

impl Debug for EventDispatcher {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        write!(f, "EventDispatcher(TraceId = {})", self.trace_id().dupe())
    }
}

#[derive(Default, Allocative)]
struct SpanStats {
    max_poll_time: Duration,
    total_poll_time: Duration,
}

/// Utility to track the start & end of a span
#[must_use = "You must end a Span you started"]
#[derive(Allocative)]
pub struct Span {
    dispatcher: EventDispatcher,
    start_instant: Instant,
    span_id: SpanId,
    stats: SpanStats,
    parent_id: Option<SpanId>,
    sent: bool,
    /// Whether we recorded (in any active span recorder) that this span was entered.
    recorded: bool,
}

impl Span {
    pub fn start<D>(dispatcher: EventDispatcher, data: D) -> Self
    where
        D: Into<span_start_event::Data>,
    {
        Self::start_impl(dispatcher, data, current_span())
    }

    fn start_impl<D>(dispatcher: EventDispatcher, data: D, parent_id: Option<SpanId>) -> Self
    where
        D: Into<span_start_event::Data>,
    {
        let span_id = SpanId::new();
        let start_instant = Instant::now();

        dispatcher.event_with_span_id(
            SpanStartEvent {
                data: Some(data.into()),
            },
            Some(span_id),
            parent_id,
        );

        Self {
            dispatcher,
            start_instant,
            span_id,
            parent_id,
            stats: Default::default(),
            sent: false,
            recorded: false,
        }
    }

    pub fn end<D>(mut self, data: D)
    where
        D: Into<span_end_event::Data>,
    {
        self.send(data.into());
    }

    pub fn span_id(&self) -> SpanId {
        self.span_id
    }

    pub fn wrap_future<End, Fut, R>(self, fut: Fut) -> impl Future<Output = R>
    where
        End: Into<span_end_event::Data>,
        Fut: Future<Output = (R, End)>,
    {
        SpanAsync { fut, span: self }
    }

    pub fn wrap_closure<End, Fun, R>(mut self, fun: Fun) -> R
    where
        End: Into<span_end_event::Data>,
        Fun: FnOnce() -> (R, End),
    {
        let (r, end) = self.call_in_span(fun);
        self.send(end.into());
        r
    }

    pub fn create_child(&self, data: impl Into<span_start_event::Data>) -> Span {
        Span::start_impl(self.dispatcher.dupe(), data, Some(self.span_id))
    }

    fn call_in_span<T, F>(&mut self, f: F) -> T
    where
        F: FnOnce() -> T,
    {
        let now = Instant::now();

        if !self.recorded {
            with_thread_local_recorder(|tl_recorder| {
                if let Some(tl_recorder) = tl_recorder.as_mut() {
                    tl_recorder.push(self.span_id);
                }
            });

            self.recorded = true;
        }

        let ret = unset_thread_local_recorder(|| {
            CURRENT_SPAN.with(|tl_span| {
                let previous_span = tl_span.replace(Some(self.span_id));
                let ret = f();
                tl_span.set(previous_span);
                ret
            })
        });

        let elapsed = now.elapsed();
        self.stats.max_poll_time = std::cmp::max(elapsed, self.stats.max_poll_time);
        self.stats.total_poll_time += elapsed;

        ret
    }

    fn send(&mut self, data: buck2_data::span_end_event::Data) {
        self.sent = true;

        self.dispatcher.event_with_span_id(
            SpanEndEvent {
                duration: self.start_instant.elapsed().try_into().ok(),
                data: Some(data),
                stats: Some(buck2_data::SpanStats {
                    max_poll_time_us: self.stats.max_poll_time.as_micros() as u64,
                    total_poll_time_us: self.stats.total_poll_time.as_micros() as u64,
                }),
            },
            Some(self.span_id),
            self.parent_id,
        );
    }
}

impl Drop for Span {
    /// When dropping a Span, send a SpanEndEvent. We send SpanCancelled in Drop to ensure that
    /// even if we never `end()` a Span, we notify clients (if any exist).
    fn drop(&mut self) {
        if !self.sent {
            self.send(buck2_data::SpanCancelled {}.into())
        }
    }
}

thread_local! {
    static CURRENT_SPAN: Cell<Option<SpanId>> = Cell::new(None);
}

#[allow(clippy::declare_interior_mutable_const)]
// TODO: remove when clippy updates past https://github.com/rust-lang/rust-clippy/issues/8493
mod ignore_clippy_warning {
    use crate::dispatch::EventDispatcher;
    tokio::task_local! {
        pub static EVENTS: EventDispatcher;
    }
}
use allocative::Allocative;
pub use ignore_clippy_warning::*;

/// Invokes function func, setting the dispatcher to the task_local for the duration
/// of the call (and for any downstream events).
pub fn with_dispatcher<R, F>(dispatcher: EventDispatcher, func: F) -> R
where
    F: FnOnce() -> R,
{
    EVENTS.sync_scope(dispatcher, func)
}

// Wraps the Future fut with a TaskLocalFuture that sets the task_local dispatcher before polling fut.
pub fn with_dispatcher_async<F, R>(dispatcher: EventDispatcher, fut: F) -> impl Future<Output = R>
where
    F: Future<Output = R>,
{
    EVENTS.scope(dispatcher, fut)
}

/// Get the ambient dispatcher, if one is available (and None otherwise). In contexts that aren't
/// tied to a specific command (e.g. materializer command loop), an ambient dispatcher may not be
/// available.
pub fn get_dispatcher_opt() -> Option<EventDispatcher> {
    match EVENTS.try_with(|dispatcher| dispatcher.dupe()) {
        Ok(dispatcher) => Some(dispatcher),
        Err(..) => None,
    }
}

pub fn get_dispatcher() -> EventDispatcher {
    static ENFORCE_DISPATCHER_SET: EnvHelper<bool> = EnvHelper::new("ENFORCE_DISPATCHER_SET");

    match get_dispatcher_opt() {
        Some(dispatcher) => dispatcher,
        None => {
            let should_error = ENFORCE_DISPATCHER_SET
                .get()
                .ok()
                .flatten()
                .copied()
                .unwrap_or_default();

            if should_error {
                panic!("dispatcher is not set")
            } else {
                // TODO: This is firing millions of times, needs to fix this up before it's made a soft error.
                // let _ignored = soft_error!(anyhow::anyhow!("Task local event dispatcher not set."));
                EventDispatcher::null()
            }
        }
    }
}

pub fn current_span() -> Option<SpanId> {
    CURRENT_SPAN.with(|tl_span| tl_span.get())
}

/// Introduces a new span and immediately fires the given start event. When the given synchronous function returns,
/// the span is closed and the end event emitted.
pub fn span<Start, End, F, R>(start: Start, func: F) -> R
where
    Start: Into<span_start_event::Data>,
    End: Into<span_end_event::Data>,
    F: FnOnce() -> (R, End),
{
    get_dispatcher().span(start, func)
}

/// Emits an InstantEvent annotated with the current trace ID
pub fn instant_event<E: Into<buck2_data::instant_event::Data>>(data: E) {
    get_dispatcher().instant_event(data)
}

/// Send console message from the server.
pub fn console_message(message: String) {
    get_dispatcher().console_message(message)
}

// Logs mercurial data
pub async fn instant_hg() {
    get_dispatcher().instant_hg().await
}

/// Introduces a new span and immediately fires the given start event. When the given future resolves,  the span is
/// closed and the event is emitted. This span is a "suspending span"; it is intended to suspend and resume whenever
/// the future itself is suspended and resumed, respectively.
///
/// TODO(swgillespie) actually implement suspend/resume
pub fn span_async<Start, End, Fut, R>(start: Start, fut: Fut) -> impl Future<Output = R>
where
    Start: Into<span_start_event::Data>,
    End: Into<span_end_event::Data>,
    Fut: Future<Output = (R, End)>,
{
    get_dispatcher().span_async(start, fut)
}

/// To use when wrapping via span() and span_async is not convenient. This produces a Span guard
/// that must be ended. The span is not automatically entered.
pub fn create_span(start: impl Into<span_start_event::Data>) -> Span {
    get_dispatcher().create_span(start)
}

/// Allows the caller to record root spans being created in their current context.
#[derive(Default)]
struct RootSpansRecorder {
    spans: SmallVec<[SpanId; 1]>,
}

impl RootSpansRecorder {
    fn new() -> Self {
        Self {
            spans: SmallVec::new(),
        }
    }

    fn push(&mut self, span_id: SpanId) {
        self.spans.push(span_id);
    }

    fn call_with_recorder<T, F>(&mut self, f: F) -> T
    where
        F: FnOnce() -> T,
    {
        let this_recorder = std::mem::take(self);

        unset_thread_local_recorder(|| {
            with_thread_local_recorder(|tl_recorder| *tl_recorder = Some(this_recorder));
            let v = f();
            // Unwrap safety: only this unset_thread_local_recorder ever unsets this, and it
            // promises to put it back.
            *self = with_thread_local_recorder(|tl_recorder| tl_recorder.take().unwrap());
            v
        })
    }
}

fn unset_thread_local_recorder<F, O>(f: F) -> O
where
    F: FnOnce() -> O,
{
    struct RestoreRecorder {
        previous_recorder: Option<RootSpansRecorder>,
    }

    impl Drop for RestoreRecorder {
        fn drop(&mut self) {
            with_thread_local_recorder(|tl_recorder| {
                *tl_recorder = std::mem::take(&mut self.previous_recorder);
            })
        }
    }

    let previous_recorder = with_thread_local_recorder(|tl_recorder| std::mem::take(tl_recorder));
    let _guard = RestoreRecorder { previous_recorder };
    f()
}

fn with_thread_local_recorder<F, O>(f: F) -> O
where
    for<'a> F: FnOnce(&'a mut Option<RootSpansRecorder>) -> O,
{
    thread_local! {
        static ROOT_SPAN_RECORDER: UnsafeCell<Option<RootSpansRecorder>> = UnsafeCell::new(None);
    }

    // SAFETY: Nobody can possibly hold a reference to the contents of this cell, since the thread
    // local is defined within this function (so nobody besies this function can hold a reference
    // to *that*), and the reference we pass ends when this function ends.
    ROOT_SPAN_RECORDER.with(|tl_recorder| {
        let recorder = unsafe { &mut *tl_recorder.get() };
        f(recorder)
    })
}

pub fn record_root_spans<F, R>(f: F) -> (R, SmallVec<[SpanId; 1]>)
where
    F: FnOnce() -> R,
{
    let mut recorder = RootSpansRecorder::new();
    let v = recorder.call_with_recorder(f);
    (v, recorder.spans)
}

#[pin_project]
pub struct RootSpansRecordingFuture<Fut> {
    recorder: RootSpansRecorder,
    #[pin]
    fut: Fut,
}

impl<Fut> Future for RootSpansRecordingFuture<Fut>
where
    Fut: Future,
{
    type Output = (<Fut as Future>::Output, SmallVec<[SpanId; 1]>);

    fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
        let this = self.project();
        let v = futures::ready!(this.recorder.call_with_recorder(|| this.fut.poll(cx)));
        task::Poll::Ready((v, std::mem::take(&mut this.recorder.spans)))
    }
}

pub fn async_record_root_spans<Fut>(fut: Fut) -> RootSpansRecordingFuture<Fut> {
    RootSpansRecordingFuture {
        recorder: RootSpansRecorder::new(),
        fut,
    }
}

#[cfg(test)]
mod tests {
    use buck2_data::CommandEnd;
    use buck2_data::CommandStart;
    use buck2_data::SpanStartEvent;
    use tokio::task::JoinHandle;

    use super::*;
    use crate::sink::channel::ChannelEventSink;
    use crate::source::ChannelEventSource;
    use crate::EventSource;

    async fn next_event(source: &mut ChannelEventSource) -> BuckEvent {
        source.receive().unwrap().unpack_buck().unwrap().clone()
    }

    fn create_dispatcher() -> (EventDispatcher, ChannelEventSource, TraceId) {
        let (send, recv) = crossbeam_channel::unbounded();
        let source = ChannelEventSource::new(recv);

        let sink = ChannelEventSink::new(send);
        let trace_id = TraceId::new();
        let dispatcher = EventDispatcher::new(trace_id.dupe(), sink);

        (dispatcher, source, trace_id)
    }

    fn create_start_end_events() -> (CommandStart, CommandEnd) {
        let start = CommandStart {
            data: Default::default(),
            metadata: Default::default(),
        };
        let end = CommandEnd {
            data: Default::default(),
            metadata: Default::default(),
            is_success: true,
            error_messages: vec![],
        };

        (start, end)
    }

    #[tokio::test]
    async fn send_event_smoke() {
        let (dispatcher, mut source, trace_id) = create_dispatcher();
        let (start, _) = create_start_end_events();

        dispatcher.event(SpanStartEvent {
            data: Some(start.into()),
        });

        let event = next_event(&mut source).await;
        assert_eq!(event.trace_id().unwrap(), trace_id);
        assert!(event.parent_id.is_none());
        assert!(event.span_id.is_none());
    }

    #[tokio::test]
    async fn send_event_with_span() {
        let (dispatcher, mut source, trace_id) = create_dispatcher();
        let (start, end) = create_start_end_events();

        dispatcher.span(start, || ((), end));

        let event = next_event(&mut source).await;
        let span_id = event.span_id.unwrap();
        assert_eq!(event.trace_id().unwrap(), trace_id);
        assert!(event.parent_id.is_none());

        let end = next_event(&mut source).await;
        assert_eq!(end.span_id.unwrap(), span_id);
    }

    #[tokio::test]
    async fn send_event_with_span_async() {
        let (dispatcher, mut source, trace_id) = create_dispatcher();
        let (start, end) = create_start_end_events();

        dispatcher.span_async(start, async { ((), end) }).await;

        let event = next_event(&mut source).await;
        let span_id = event.span_id.unwrap();
        assert_eq!(event.trace_id().unwrap(), trace_id);
        assert!(event.parent_id.is_none());

        let end = next_event(&mut source).await;
        assert_eq!(end.span_id.unwrap(), span_id);
    }

    #[tokio::test]
    async fn send_event_with_nested_span() {
        let (dispatcher, mut source, _) = create_dispatcher();
        let (start, end) = create_start_end_events();

        dispatcher.span(start.clone(), || {
            dispatcher.span(start.clone(), || ((), end.clone()));
            ((), end.clone())
        });

        let e1 = next_event(&mut source).await;
        let e2 = next_event(&mut source).await;
        let e3 = next_event(&mut source).await;
        let e4 = next_event(&mut source).await;

        assert_eq!(e1.span_id, e4.span_id);
        assert_eq!(e2.span_id, e3.span_id);
        assert_eq!(e2.parent_id, e1.span_id);
    }

    #[tokio::test]
    async fn send_event_with_nested_span_async() {
        let (dispatcher, mut source, _) = create_dispatcher();
        let (start, end) = create_start_end_events();

        dispatcher
            .span_async(start.clone(), async {
                dispatcher
                    .span_async(start.clone(), async { ((), end.clone()) })
                    .await;
                ((), end.clone())
            })
            .await;

        let e1 = next_event(&mut source).await;
        let e2 = next_event(&mut source).await;
        let e3 = next_event(&mut source).await;
        let e4 = next_event(&mut source).await;

        assert_eq!(e1.span_id, e4.span_id);
        assert_eq!(e2.span_id, e3.span_id);
        assert_eq!(e2.parent_id, e1.span_id);
    }

    #[tokio::test]
    async fn send_event_with_nested_span_explicit() {
        let (dispatcher, mut source, _trace_id) = create_dispatcher();
        let (start, end) = create_start_end_events();

        let top = dispatcher.create_span(start.clone());
        let middle = top.create_child(start.clone());
        middle.wrap_closure(|| {
            (
                dispatcher.span(start.clone(), || ((), end.clone())),
                end.clone(),
            )
        });
        top.end(end);

        let enter_top = next_event(&mut source).await;
        let enter_middle = next_event(&mut source).await;
        let enter_bottom = next_event(&mut source).await;
        let exit_bottom = next_event(&mut source).await;
        let exit_middle = next_event(&mut source).await;
        let exit_top = next_event(&mut source).await;

        assert_eq!(enter_top.span_id, exit_top.span_id);
        assert_eq!(enter_middle.span_id, exit_middle.span_id);
        assert_eq!(enter_bottom.span_id, exit_bottom.span_id);

        assert_eq!(enter_middle.parent_id.unwrap(), enter_top.span_id.unwrap());
        assert_eq!(
            enter_bottom.parent_id.unwrap(),
            enter_middle.span_id.unwrap()
        );
    }

    #[test]
    fn test_span_stats() {
        let (dispatcher, _, _) = create_dispatcher();
        let (start, _) = create_start_end_events();

        let mut span = Span::start(dispatcher.dupe(), start);
        span.call_in_span(|| std::thread::sleep(Duration::from_millis(10)));
        assert_eq!(span.stats.max_poll_time, span.stats.total_poll_time);

        span.call_in_span(|| std::thread::sleep(Duration::from_millis(10)));
        assert!(span.stats.max_poll_time < span.stats.total_poll_time);
    }

    async fn dispatcher_task(
        dispatcher: EventDispatcher,
        start: CommandStart,
        end: CommandEnd,
    ) -> Result<u32, ()> {
        let mut err = false;
        let mut events = 0u32;

        async fn event(
            start: CommandStart,
            end: CommandEnd,
            events: &mut u32,
            trace_id: &TraceId,
            err: &mut bool,
        ) {
            async fn yield_and_check(trace_id: &TraceId, err: &mut bool) {
                tokio::task::yield_now().await;
                *err |= get_dispatcher().trace_id != *trace_id;
            }

            // 2*2 for Start and End events, for 2 spans
            *events += 4;

            span_async(start.clone(), async {
                for _ in 0..4 {
                    yield_and_check(trace_id, err).await;
                }
                // TODO(mufeez): re-consider adding recursive calls to event
                span_async(start.clone(), async {
                    for _ in 0..4 {
                        yield_and_check(trace_id, err).await;
                    }
                    ((), end.clone())
                })
                .await;
                ((), end.clone())
            })
            .await
        }

        with_dispatcher_async(
            dispatcher.dupe(),
            event(
                start.clone(),
                end.clone(),
                &mut events,
                &dispatcher.trace_id,
                &mut err,
            ),
        )
        .await;

        if err { Err(()) } else { Ok(events) }
    }

    // Test function used in test_concurrent_single_thread and
    // test_concurrent_multi_thread.
    async fn test_concurrent() {
        let (events1, source1, trace_id1) = create_dispatcher();
        let (events2, source2, trace_id2) = create_dispatcher();

        let mut event_sources: [ChannelEventSource; 2] = [source1, source2];
        let trace_ids: [TraceId; 2] = [trace_id1, trace_id2];

        let (start1, end1) = create_start_end_events();
        let (start2, end2) = create_start_end_events();

        let handles: [JoinHandle<Result<u32, ()>>; 2] = [
            tokio::spawn(dispatcher_task(events1, start1, end1)),
            tokio::spawn(dispatcher_task(events2, start2, end2)),
        ];

        let mut sent_counts = Vec::new();

        for handle in handles {
            match handle.await {
                Ok(result) => {
                    sent_counts.push(result.unwrap());
                }
                Err(e) => assert!(!e.is_panic()),
            }
        }

        for (i, sent) in sent_counts.iter().enumerate() {
            for _ in 0..*sent {
                let e = next_event(&mut event_sources[i]).await;
                assert!(e.trace_id().unwrap() == trace_ids[i]);
            }
        }
    }

    #[tokio::test]
    async fn test_concurrent_single_thread() {
        test_concurrent().await;
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn test_concurrent_multi_thread() {
        test_concurrent().await;
    }

    fn with_span(f: impl FnOnce()) -> SpanId {
        span(CommandStart::default(), || {
            let id = current_span().unwrap();
            f();
            (id, CommandEnd::default())
        })
    }

    #[test]
    fn test_record_root_spans() {
        let nop = || ();

        // Test a basic case.
        let (id, spans) = record_root_spans(|| with_span(nop));
        assert_eq!(&*spans, &[id]);

        // Test multiple spans.
        let (ids, spans) = record_root_spans(|| {
            let id1 = with_span(nop);
            let id2 = with_span(nop);
            vec![id1, id2]
        });
        assert_eq!(&*spans, &ids);

        // Test we don't record nested spans.
        let (id, spans) = record_root_spans(|| {
            with_span(|| {
                span(CommandStart::default(), || {
                    (with_span(nop), CommandEnd::default())
                });
            })
        });
        assert_eq!(&*spans, &[id]);

        // Test nested calls to record_root_spans work.
        let (id, spans) = record_root_spans(|| {
            with_span(|| {
                let (id, spans) = record_root_spans(|| with_span(nop));
                assert_eq!(&*spans, &[id]);
            })
        });
        assert_eq!(&*spans, &[id]);
    }

    #[tokio::test]
    async fn test_async_record_root_spans() {
        let nop = || ();

        // Easy case.
        let (id, spans) = async_record_root_spans(async { with_span(nop) }).await;
        assert_eq!(&*spans, &[id]);

        // Easy case (2).
        let (id, spans) = async_record_root_spans(span_async(CommandStart::default(), async {
            (current_span().unwrap(), CommandEnd::default())
        }))
        .await;
        assert_eq!(&*spans, &[id]);

        // With yields, check we don't re-enter.
        let (ids, spans) = async_record_root_spans(async {
            let id1 = span_async(CommandStart::default(), async {
                tokio::task::yield_now().await;
                (current_span().unwrap(), CommandEnd::default())
            })
            .await;

            tokio::task::yield_now().await;

            let id2 = span_async(CommandStart::default(), async {
                tokio::task::yield_now().await;
                (current_span().unwrap(), CommandEnd::default())
            })
            .await;

            vec![id1, id2]
        })
        .await;
        assert_eq!(&*spans, &ids);
    }
}


/*
 * 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.
 */

#![cfg_attr(feature = "gazebo_lint", feature(plugin))]
#![cfg_attr(feature = "gazebo_lint", allow(deprecated))] // :(
#![cfg_attr(feature = "gazebo_lint", plugin(gazebo_lint))]

//!
//! Events and event streams for Buck2.
//!
//! The `Event` enum is the set of events that Buck2 can produce. Events can be produced both during the course of
//! a command or through background operation (such as DICE invalidations of changed files).
//!
//! There are three critical nouns in this data model:
//!  * An **event**, which is a structure representing a point-in-time occurrence, with some additional data.
//!  * A **trace**, which is a collection of events that are to be interpreted as semantically linked. Traces are
//!    globally identified by a trace ID, which is a v4 UUID. Traces are often associated with individual commands,
//!    although they do not have to.
//!  * A **span**, which is a pair of two events that represent a start and stop pair. A span covers a range of time
//!    points. All events are parented to a span that was currently active at the location the event was emitted.

pub mod daemon_id;
pub mod dispatch;
pub mod metadata;
pub mod sink;
pub mod source;
pub mod span;

use std::num::NonZeroU64;
use std::str::FromStr;
use std::sync::Arc;
use std::time::SystemTime;

use anyhow::Context;
use async_trait::async_trait;
use buck2_cli_proto::CommandResult;
use buck2_cli_proto::PartialResult;
use buck2_wrapper_common::invocation_id::TraceId;
use derive_more::From;
use gazebo::variants::UnpackVariants;
use serde::Serialize;
use thiserror::Error;

use crate::sink::channel::ChannelEventSink;
use crate::source::ChannelEventSource;
use crate::span::SpanId;

/// An event that can be produced by Buck2. Events are points in time with additional metadata attached to them,
/// depending on the nature of the event.
///
/// Some events are special in that they represent points in time where an operation started or ended. These events
/// introduce new "spans". All events belong to a span except the first and last events of a trace. All spans except
/// the span created by the first and last events of the trace have a parent; as such, spans form a tree.
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct BuckEvent {
    /// Full event, the rest of the fields are caches.
    event: Box<buck2_data::BuckEvent>,

    /// A timestamp for when this event was emitted.
    timestamp: SystemTime,

    /// If this event starts a new span, the span ID assigned to this span, or None if this event is a leaf event
    /// that does not start a new span.
    span_id: Option<SpanId>,

    /// The ID of the span that contains this event. Will be non-None in all Events except the first and last events
    /// of a trace.
    parent_id: Option<SpanId>,
}

impl BuckEvent {
    pub fn new(
        timestamp: SystemTime,
        trace_id: TraceId,
        span_id: Option<SpanId>,
        parent_id: Option<SpanId>,
        data: buck2_data::buck_event::Data,
    ) -> BuckEvent {
        let event = buck2_data::BuckEvent {
            timestamp: Some(timestamp.into()),
            trace_id: trace_id.to_string(),
            span_id: span_id.map_or(0, |s| s.0.into()),
            parent_id: parent_id.map_or(0, |s| s.0.into()),
            data: Some(data),
        };
        BuckEvent {
            event: Box::new(event),
            timestamp,
            span_id,
            parent_id,
        }
    }

    pub fn timestamp(&self) -> SystemTime {
        self.timestamp
    }

    pub fn trace_id(&self) -> anyhow::Result<TraceId> {
        Ok(TraceId::from_str(&self.event.trace_id)?)
    }

    pub fn span_id(&self) -> Option<SpanId> {
        self.span_id
    }

    pub fn parent_id(&self) -> Option<SpanId> {
        self.parent_id
    }

    pub fn event(&self) -> &buck2_data::BuckEvent {
        &self.event
    }

    pub fn data(&self) -> &buck2_data::buck_event::Data {
        self.event
            .data
            .as_ref()
            .expect("data is set, it is validated")
    }

    pub fn data_mut(&mut self) -> &mut buck2_data::buck_event::Data {
        self.event
            .data
            .as_mut()
            .expect("data is set, it is validated")
    }

    pub fn span_start_event(&self) -> Option<&buck2_data::SpanStartEvent> {
        match self.data() {
            buck2_data::buck_event::Data::SpanStart(start) => Some(start),
            _ => None,
        }
    }

    pub fn span_end_event(&self) -> Option<&buck2_data::SpanEndEvent> {
        match self.data() {
            buck2_data::buck_event::Data::SpanEnd(end) => Some(end),
            _ => None,
        }
    }

    pub fn command_start(&self) -> anyhow::Result<Option<&buck2_data::CommandStart>> {
        match self.span_start_event() {
            None => Ok(None),
            Some(span_start_event) => {
                match span_start_event
                    .data
                    .as_ref()
                    .with_context(|| BuckEventError::MissingField(self.clone()))?
                {
                    buck2_data::span_start_event::Data::Command(command_start) => {
                        Ok(Some(command_start))
                    }
                    _ => Ok(None),
                }
            }
        }
    }
}

impl From<BuckEvent> for Box<buck2_data::BuckEvent> {
    fn from(e: BuckEvent) -> Self {
        e.event
    }
}

impl TryFrom<Box<buck2_data::BuckEvent>> for BuckEvent {
    type Error = anyhow::Error;

    fn try_from(event: Box<buck2_data::BuckEvent>) -> anyhow::Result<BuckEvent> {
        event.data.as_ref().ok_or(BuckEventError::MissingData)?;
        fn new_span_id(num: u64) -> Option<SpanId> {
            NonZeroU64::new(num).map(SpanId)
        }
        Ok(Self {
            timestamp: SystemTime::try_from(
                event
                    .timestamp
                    .clone()
                    .ok_or(BuckEventError::MissingTimestamp)?,
            )?,
            span_id: new_span_id(event.span_id),
            parent_id: new_span_id(event.parent_id),
            event,
        })
    }
}

/// An event that can be produced by Buck2 that is not intended to be presented to the user, but rather is used to
/// communicate with other parts of Buck2.
#[derive(Clone, From)]
pub enum ControlEvent {
    /// A command result, produced upon completion of a command.
    CommandResult(Box<CommandResult>),
    /// A progress event from this command. Different commands have different types.
    PartialResult(PartialResult),
}

/// The set of events that can flow out of an EventSource. Control events are not intended to be sent across the gRPC
/// boundary, while Buck events are.
#[derive(Clone, From, UnpackVariants)]
#[allow(clippy::large_enum_variant)]
pub enum Event {
    /// A control event: events that are not to be exposed across gRPC, but inform stream consumers of important control
    /// changes, such the CommandResult event, which implies that there will be no further events.
    Control(ControlEvent),
    /// A buck event: events that are to be exposed across gRPC verbatim.
    Buck(BuckEvent),
}

/// Statistics from this event sink on how messages were processed.
#[derive(Clone, Debug)]
pub struct EventSinkStats {
    /// Count of number of successful messages (e.g. those that have been processed by their downstream destination).
    pub successes: u64,
    /// Count of messages that failed to be submitted and will not be retried.
    pub failures: u64,
    /// How many messages are currently buffered by this sink.
    pub buffered: u64,
    /// How many messages were not even enqueued by this sink.
    pub dropped: u64,
}

impl EventSinkStats {
    /// Since there can be one or more sinks (e.g. with [`sink::TeeSink`](???)), we need to aggregate these into
    /// useful singular values.
    pub fn aggregate(&self, other: &Self) -> Self {
        Self {
            successes: self.successes + other.successes,
            failures: self.failures + other.failures,
            buffered: self.buffered + other.buffered,
            dropped: self.dropped + other.dropped,
        }
    }
}

/// A sink for events, easily plumbable to the guts of systems that intend to produce events consumeable by
/// higher-level clients. Sending an event is synchronous.
pub trait EventSink: Send + Sync {
    /// Sends an event into this sink, to be consumed elsewhere. Explicitly does not return a Result type; if sending
    /// an event does fail, implementations will handle the failure by panicking or performing some other graceful
    /// recovery; callers of EventSink are not expected to handle failures.
    fn send(&self, event: BuckEvent);

    /// Sends a control event into this sink, to be consumed elsewhere. Control events are not sent to gRPC clients.
    fn send_control(&self, control_event: ControlEvent);

    /// Collects stats on this sink (e.g. messages accepted, rejected).
    fn stats(&self) -> Option<EventSinkStats>;
}

impl EventSink for Arc<dyn EventSink> {
    fn send(&self, event: BuckEvent) {
        EventSink::send(self.as_ref(), event);
    }

    fn send_control(&self, control_event: ControlEvent) {
        EventSink::send_control(self.as_ref(), control_event);
    }

    fn stats(&self) -> Option<EventSinkStats> {
        EventSink::stats(self.as_ref())
    }
}

/// A source for events, suitable for reading a stream of events coming out of Buck.
#[async_trait]
pub trait EventSource: Send {
    fn receive(&mut self) -> Option<Event>;

    fn try_receive(&mut self) -> Option<Event>;
}

/// Creates a pair of an EventSource and an EventSink such that writes to the sink can be read by the event source.
pub fn create_source_sink_pair() -> (impl EventSource, impl EventSink) {
    let (send, recv) = crossbeam_channel::unbounded();
    let sink = ChannelEventSink::new(send);
    let source = ChannelEventSource::new(recv);
    (source, sink)
}

#[allow(clippy::large_enum_variant)]
#[derive(Error, Debug)]
enum BuckEventError {
    #[error("The `buck2_data::BuckEvent` provided has no `Timestamp`")]
    MissingTimestamp,
    #[error("The `buck2_data::BuckEvent` provided has no `Data`")]
    MissingData,
    #[error("Sent an event missing one or more fields: `{0:?}`")]
    MissingField(BuckEvent),
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;

    use buck2_data::CommandStart;
    use buck2_data::SpanStartEvent;

    use super::*;

    #[test]
    fn round_trip_success() {
        let test = BuckEvent::new(
            SystemTime::now(),
            TraceId::new(),
            Some(SpanId::new()),
            Some(SpanId::new()),
            SpanStartEvent {
                data: Some(
                    CommandStart {
                        data: None,
                        metadata: HashMap::new(),
                    }
                    .into(),
                ),
            }
            .into(),
        );
        assert_eq!(
            test,
            BuckEvent::try_from(Box::<buck2_data::BuckEvent>::from(test.clone())).unwrap()
        );
    }

    #[test]
    fn trace_id_hash_produces_a_reasonable_number() {
        let trace_id = TraceId::from_str("0436430c-2b02-624c-2032-570501212b57").unwrap();
        let hash = trace_id.hash();
        assert_eq!(5739261098605499414, hash);

        let other_trace_id = TraceId::from_str("586615bb-f57a-45a6-8804-3c6fcb0347de").unwrap();
        let hash = other_trace_id.hash();
        assert_eq!(8717222666446319742, hash);
    }
}


/*
 * 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.
 */

//! Metadata collection, for telemetry purposes.
use std::collections::HashMap;
use std::env;

use buck2_core::facebook_only;
use buck2_wrapper_common::BUCK2_WRAPPER_ENV_VAR;
use once_cell::sync::OnceCell;

use crate::daemon_id::DAEMON_UUID;

/// Collects metadata from the current binary and environment and writes it as map, suitable for telemetry purposes.
pub fn collect() -> HashMap<String, String> {
    facebook_only();
    fn add_env_var(map: &mut HashMap<String, String>, key: &'static str, var: &'static str) {
        if let Ok(data) = env::var(var) {
            map.insert(key.to_owned(), data);
        }
    }

    let mut map = HashMap::new();

    let info = system_info();
    if let Some(hostname) = info.hostname {
        map.insert("hostname".to_owned(), hostname);
    }
    if let Some(username) = info.username {
        map.insert("username".to_owned(), username);
    }

    if let Some(rev) = buck2_build_info::revision() {
        map.insert("buck2_revision".to_owned(), rev.to_owned());
    }

    if let Some(time) = buck2_build_info::time_iso8601() {
        map.insert("buck2_build_time".to_owned(), time.to_owned());
    }

    if let Some(ts) = buck2_build_info::release_timestamp() {
        map.insert("buck2_release_timestamp".to_owned(), ts.to_owned());
    }

    if is_proc_translated::is_proc_translated() {
        map.insert("rosetta".to_owned(), "1".to_owned());
    }

    // Global trace ID
    map.insert("daemon_uuid".to_owned(), DAEMON_UUID.to_string());

    map.insert("os".to_owned(), info.os);
    if let Some(version) = info.os_version {
        map.insert("os_version".to_owned(), version);
    }

    add_env_var(&mut map, "sandcastle_job_info", "SANDCASTLE_JOB_INFO");
    add_env_var(&mut map, "sandcastle_alias", "SANDCASTLE_ALIAS");
    add_env_var(&mut map, "launched_via_wrapper", BUCK2_WRAPPER_ENV_VAR);
    add_env_var(&mut map, "fbpackage_name", "FBPACKAGE_PACKAGE_NAME");
    add_env_var(&mut map, "fbpackage_version", "FBPACKAGE_PACKAGE_VERSION");
    add_env_var(&mut map, "fbpackage_release", "FBPACKAGE_PACKAGE_RELEASE");
    map
}

pub struct SystemInfo {
    pub username: Option<String>,
    pub hostname: Option<String>,
    pub os: String,
    pub os_version: Option<String>,
}

pub fn system_info() -> SystemInfo {
    let hostname = hostname();
    let username;
    #[cfg(any(fbcode_build, cargo_internal_build))]
    {
        username = user::current_username().ok();
    }
    #[cfg(not(any(fbcode_build, cargo_internal_build)))]
    {
        username = None;
    }

    SystemInfo {
        hostname,
        username,
        os: os_type(),
        os_version: sys_info::os_release().ok(),
    }
}

/// The operating system - "linux" "darwin" "windows" etc.
fn os_type() -> String {
    if cfg!(target_os = "linux") {
        "linux".to_owned()
    } else if cfg!(target_os = "macos") {
        "darwin".to_owned()
    } else if cfg!(target_os = "windows") {
        "windows".to_owned()
    } else {
        "unknown".to_owned()
    }
}

pub fn hostname() -> Option<String> {
    static CELL: OnceCell<Option<String>> = OnceCell::new();

    CELL.get_or_init(|| {
        hostname::get()
            .ok()
            .map(|res| res.to_string_lossy().into_owned())
    })
    .clone()
}


/*
 * 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.
 */

//! Implementations of `[crate::EventSink]` that are useful in different situations. Buck2 primarily uses the `channel`
//! sink during normal operation.
pub(crate) mod channel;
pub(crate) mod null;
pub mod scribe;
pub mod tee;


/*
 * 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 async_trait::async_trait;

use crate::Event;
use crate::EventSource;

/// A simple EventSource backed by an unbounded channel.
pub struct ChannelEventSource(crossbeam_channel::Receiver<Event>);

impl ChannelEventSource {
    pub fn new(recv: crossbeam_channel::Receiver<Event>) -> ChannelEventSource {
        ChannelEventSource(recv)
    }
}

#[async_trait]
impl EventSource for ChannelEventSource {
    fn receive(&mut self) -> Option<Event> {
        self.0.recv().ok()
    }

    fn try_receive(&mut self) -> Option<Event> {
        self.0.try_recv().ok()
    }
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::time::SystemTime;

    use buck2_data::buck_event::Data::SpanStart;
    use buck2_data::span_start_event::Data::Command;
    use buck2_data::CommandStart;
    use buck2_data::SpanStartEvent;

    use super::ChannelEventSource;
    use crate::sink::channel::ChannelEventSink;
    use crate::BuckEvent;
    use crate::EventSink;
    use crate::EventSource;
    use crate::TraceId;

    #[tokio::test]
    async fn receive_smoke() {
        let (send, recv) = crossbeam_channel::unbounded();
        let sink = ChannelEventSink::new(send);
        let mut source = ChannelEventSource::new(recv);
        sink.send(BuckEvent::new(
            SystemTime::now(),
            TraceId::new(),
            None,
            None,
            SpanStartEvent {
                data: Some(
                    CommandStart {
                        data: None,
                        metadata: HashMap::new(),
                    }
                    .into(),
                ),
            }
            .into(),
        ));
        let event = source.receive().unwrap().unpack_buck().unwrap().clone();
        assert!(matches!(
            event.data(),
            SpanStart(SpanStartEvent {
                data: Some(Command(CommandStart { .. }))
            })
        ));
    }
}


/*
 * 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::num::NonZeroU64;
use std::sync::atomic::AtomicU64;
use std::sync::atomic::Ordering;

use allocative::Allocative;
use dupe::Dupe;
use serde::Serialize;

/// A SpanId is a unique identifier for a span, which is a pair of events that represent a conceptual start and stop
/// of a particular operation.
#[derive(
    Debug,
    Copy,
    Clone,
    Dupe,
    Serialize,
    PartialEq,
    Eq,
    Hash,
    derive_more::Display,
    Allocative
)]
pub struct SpanId(pub(crate) NonZeroU64);

impl SpanId {
    /// Generates a new SpanId, suitable for identifying a particular span within the context of a trace. Span IDs are
    /// increasing nonzero 64-bit integers.
    pub fn new() -> SpanId {
        static NEXT_ID: AtomicU64 = AtomicU64::new(1);
        loop {
            let next_id = NEXT_ID.fetch_add(1, Ordering::AcqRel);
            match NonZeroU64::new(next_id) {
                Some(id) => return SpanId(id),
                // 64-bit wrap around; continue the loop to generate the next non-zero ID.
                None => continue,
            }
        }
    }
}

impl From<SpanId> for u64 {
    fn from(span_id: SpanId) -> Self {
        span_id.0.into()
    }
}


/*
 * 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 dupe::Dupe;

use crate::BuckEvent;
use crate::ControlEvent;
use crate::Event;
use crate::EventSink;
use crate::EventSinkStats;

/// An EventSink implementation that pushes events onto an unbounded channel, to be consumed by receivers on said
/// channel.
#[derive(Clone)]
pub struct ChannelEventSink(crossbeam_channel::Sender<Event>);

impl ChannelEventSink {
    pub fn new(send: crossbeam_channel::Sender<Event>) -> ChannelEventSink {
        ChannelEventSink(send)
    }
}

impl EventSink for ChannelEventSink {
    fn send(&self, event: BuckEvent) {
        let _ignore = self.0.send(event.into());
    }

    fn send_control(&self, control_event: ControlEvent) {
        if let Err(e) = self.0.send(control_event.into()) {
            panic!("failed to send control event to ChannelEventSink: {}", e);
        }
    }

    fn stats(&self) -> Option<EventSinkStats> {
        None
    }
}

// Justification: UnboundedSender is backed by a
// [`Tx`](https://github.com/tokio-rs/tokio/blob/tokio-0.2.25/tokio/src/sync/mpsc/chan.rs#L224-L232),
// whose Clone implementation bumps a few atomics and clones an Arc.
impl Dupe for ChannelEventSink {}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::time::SystemTime;

    use buck2_data::buck_event::Data::SpanStart;
    use buck2_data::span_start_event::Data::Command;
    use buck2_data::CommandStart;
    use buck2_data::SpanStartEvent;

    use super::ChannelEventSink;
    use crate::BuckEvent;
    use crate::EventSink;
    use crate::TraceId;

    #[tokio::test]
    async fn sending_event_smoke() {
        let (send, recv) = crossbeam_channel::unbounded();
        let sink = ChannelEventSink::new(send);
        sink.send(BuckEvent::new(
            SystemTime::now(),
            TraceId::new(),
            None,
            None,
            SpanStartEvent {
                data: Some(
                    CommandStart {
                        data: None,
                        metadata: HashMap::new(),
                    }
                    .into(),
                ),
            }
            .into(),
        ));
        let event = recv.recv().unwrap().unpack_buck().unwrap().clone();
        assert!(matches!(
            event.data(),
            SpanStart(SpanStartEvent {
                data: Some(Command(CommandStart { .. }))
            })
        ));
    }
}


/*
 * 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 crate::BuckEvent;
use crate::ControlEvent;
use crate::EventSink;
use crate::EventSinkStats;

/// A null EventSink that discards all messages sent to it.
pub struct NullEventSink;

impl NullEventSink {
    pub fn new() -> NullEventSink {
        NullEventSink
    }
}

impl Default for NullEventSink {
    fn default() -> Self {
        NullEventSink::new()
    }
}

impl EventSink for NullEventSink {
    fn send(&self, _: BuckEvent) {}

    fn send_control(&self, _: ControlEvent) {}

    fn stats(&self) -> Option<EventSinkStats> {
        None
    }
}