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

//! Handles command concurrency.
//!
//! `buck2` supports limited concurrency for commands.
//! If there are no buckconfig changes, nor file changes, then commands can be allowed to execute
//! concurrently. Otherwise, `buck2` will block waiting for other commands to finish.

use std::collections::VecDeque;
use std::fmt::Debug;
use std::ops::Deref;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use async_condvar_fair::Condvar;
use async_trait::async_trait;
use buck2_core::soft_error;
use buck2_data::DiceBlockConcurrentCommandEnd;
use buck2_data::DiceBlockConcurrentCommandStart;
use buck2_data::DiceConcurrentCommands;
use buck2_data::DiceEqualityCheck;
use buck2_data::DiceSynchronizeSectionEnd;
use buck2_data::DiceSynchronizeSectionStart;
use buck2_data::ExclusiveCommandWaitEnd;
use buck2_data::ExclusiveCommandWaitStart;
use buck2_data::ExitWhenDifferentState;
use buck2_data::NoActiveDiceState;
use buck2_events::dispatch::EventDispatcher;
use buck2_util::truncate::truncate;
use buck2_wrapper_common::invocation_id::TraceId;
use derive_more::Display;
use dice::Dice;
use dice::DiceComputations;
use dice::DiceEquality;
use dice::DiceTransaction;
use dice::DiceTransactionUpdater;
use dice::UserComputationData;
use dupe::Dupe;
use futures::future::BoxFuture;
use futures::future::Future;
use futures::future::FutureExt;
use futures::future::Shared;
use itertools::Itertools;
use more_futures::cancellation::ExplicitCancellationContext;
use starlark_map::small_map::SmallMap;
use starlark_map::small_set::SmallSet;
use thiserror::Error;
use tokio::sync::Mutex;
use tokio::sync::MutexGuard;

#[derive(Error, Debug)]
enum ConcurrencyHandlerError {
    #[error(
        "Recursive invocation of Buck, which is discouraged, but will probably work (using the same state). Trace Ids: {0}. Recursive invocation command: `{1}`"
    )]
    NestedInvocationWithSameStates(String, String),
    #[error(
        "Recursive invocation of Buck, with a different state. Use `--isolation-dir` on the inner invocation to fix this. Trace Ids: {0}. Recursive invocation command: `{1}`"
    )]
    NestedInvocationWithDifferentStates(String, String),
    #[error("`--exit-when-different-state` was set")]
    ExitWhenDifferentState,
}

#[derive(Clone, Dupe, Copy, Debug)]
pub enum RunState {
    NestedSameState,
    ParallelSameState,
}

#[derive(Clone, Dupe, Copy, Debug)]
pub enum BypassSemaphore {
    Run(RunState),
    Block,
    Error,
}

/// Manages concurrent commands, blocking when appropriate.
///
/// Currently, we allow concurrency if two `DiceTransactions` are deemed equivalent, such that
/// any computation result that occurs in one is directly reusable by another.
#[derive(Clone, Dupe, Allocative)]
pub struct ConcurrencyHandler {
    data: Arc<Mutex<ConcurrencyHandlerData>>,
    // use an async condvar because the `wait` to `notify` spans across an async function (namely
    // the entire command execution).
    #[allocative(skip)]
    cond: Arc<Condvar>,
    dice: Arc<Dice>,
    /// Used to prevent commands (clean --stale) from running in parallel with dice commands
    exclusive_command_lock: Arc<ExclusiveCommandLock>,
}

#[derive(Allocative)]
struct ConcurrencyHandlerData {
    /// the currently active `Dice` being used. Commands can only run concurrently if these are
    /// "equivalent".
    dice_status: DiceStatus,
    /// A list of the currently running commands.
    active_commands: SmallMap<CommandId, CommandData>,
    /// When a command enters
    next_command_id: CommandId,
    /// The epoch of the last ActiveDice we assigned.
    cleanup_epoch: usize,
    /// Whether this has been tainted previously.
    previously_tainted: bool,
}

#[derive(Allocative, Display, Copy, Clone, Dupe, PartialEq, Eq, Hash)]
struct CommandId(usize);

impl CommandId {
    /// Increment this counter and return the next command.
    fn increment(&mut self) -> CommandId {
        let res = CommandId(self.0);
        self.0 += 1;
        res
    }
}

#[derive(Allocative)]
struct CommandData {
    trace_id: TraceId,
    argv: Vec<String>,
    dispatcher: EventDispatcher,
}

impl CommandData {
    fn format_argv(&self) -> String {
        let mut iter = self.argv.iter();
        // Skip the "/path/to/buck2" part so we can just emit "buck2" for the start of the cmd
        iter.next();

        let cmd = format!("buck2 {}", iter.join(" "));
        truncate(&cmd, 500)
    }

    fn notify_tainted(&self) {
        self.dispatcher.instant_event(buck2_data::TagEvent {
            tags: vec!["concurrency-tainted".to_owned()],
        });
    }

    fn notify_previously_tainted(&self) {
        self.dispatcher.instant_event(buck2_data::TagEvent {
            tags: vec!["concurrency-previously-tainted".to_owned()],
        });
    }
}

#[derive(Allocative)]
enum DiceStatus {
    Available {
        active: Option<ActiveDice>,
    },
    Cleanup {
        future: Shared<BoxFuture<'static, ()>>,
        epoch: usize,
    },
}

#[derive(Allocative)]
struct ActiveDice {
    version: DiceEquality,

    /// Whether this DICE version had concurrent commands that executed on it.
    tainted: bool,
}

impl DiceStatus {
    fn idle() -> Self {
        Self::Available { active: None }
    }

    fn active(version: DiceEquality) -> Self {
        Self::Available {
            active: Some(ActiveDice {
                version,
                tainted: false,
            }),
        }
    }
}

impl ConcurrencyHandlerData {
    fn has_no_active_commands(&self) -> bool {
        self.active_commands.is_empty()
    }

    /// Attempt a transition to cleanup, or straight to idle if cleanup can be skipped. Returns
    /// whether the transition was done.
    fn transition_to_cleanup(&mut self, dice: &Dice) -> bool {
        if !self.has_no_active_commands() {
            return false;
        }

        tracing::info!("Transitioning ActiveDice to cleanup");

        // When releasing the active DICE, if any work is ongoing, place it in a clean up
        // state. Callers will wait until it goes idle.
        self.cleanup_epoch += 1;
        self.dice_status = DiceStatus::Cleanup {
            future: dice.wait_for_idle().boxed().shared(),
            epoch: self.cleanup_epoch,
        };

        true
    }

    /// Attempt a transition to available assuming the cleanup future at `cleanup_epoch` has been
    /// awaited already.
    fn transition_to_idle(&mut self, cleanup_epoch: usize) {
        if !matches!(self.dice_status, DiceStatus::Cleanup { .. }) {
            // Noop: we already transitioned to available.
            return;
        }

        if self.cleanup_epoch != cleanup_epoch {
            // Noop: we already transitioned to available then back to cleanup.
            return;
        }

        self.dice_status = DiceStatus::idle();
    }

    fn notify_tainted(&self) {
        for command in self.active_commands.values() {
            command.notify_tainted()
        }
    }
}

#[async_trait]
pub trait DiceUpdater: Send + Sync {
    async fn update(
        &self,
        mut ctx: DiceTransactionUpdater,
        _user_data: &mut UserComputationData,
    ) -> anyhow::Result<DiceTransactionUpdater>;
}

#[async_trait]
pub trait DiceDataProvider: Send + Sync + 'static {
    async fn provide(&self, ctx: &DiceComputations) -> anyhow::Result<UserComputationData>;
}

#[derive(Allocative)]
struct ExclusiveCommandLock {
    lock: tokio::sync::RwLock<()>,
    owning_command: Arc<parking_lot::Mutex<VecDeque<String>>>,
}

enum ExclusiveCommandLockGuard<'a> {
    Shared(tokio::sync::RwLockReadGuard<'a, ()>),
    Exclusive(
        tokio::sync::RwLockWriteGuard<'a, ()>,
        Arc<parking_lot::Mutex<VecDeque<String>>>,
    ),
}

impl<'a> Drop for ExclusiveCommandLockGuard<'a> {
    fn drop(&mut self) {
        if let ExclusiveCommandLockGuard::Exclusive(_, owner) = self {
            let mut own = owner.lock();
            own.pop_front();
        }
    }
}

impl ExclusiveCommandLock {
    pub fn new() -> Self {
        ExclusiveCommandLock {
            lock: tokio::sync::RwLock::new(()),
            owning_command: Arc::new(parking_lot::Mutex::new(VecDeque::new())),
        }
    }

    pub async fn exclusive_lock<'a>(&'a self, cmd_name: String) -> ExclusiveCommandLockGuard<'a> {
        {
            let mut owning_command = self.owning_command.lock();
            owning_command.push_back(cmd_name);
            drop(owning_command);
        }
        ExclusiveCommandLockGuard::Exclusive(self.lock.write().await, self.owning_command.dupe())
    }

    pub async fn shared_lock<'a>(&'a self) -> ExclusiveCommandLockGuard<'a> {
        ExclusiveCommandLockGuard::Shared(self.lock.read().await)
    }

    pub fn owning_command(&self) -> Option<String> {
        // owning command is not unset when exclusive lock is dropped, just ignored
        if self.lock.try_read().is_ok() {
            None
        } else {
            self.owning_command.lock().front().cloned()
        }
    }
}

impl ConcurrencyHandler {
    pub fn new(dice: Arc<Dice>) -> Self {
        ConcurrencyHandler {
            data: Arc::new(Mutex::new(ConcurrencyHandlerData {
                dice_status: DiceStatus::idle(),
                active_commands: SmallMap::new(),
                next_command_id: CommandId(0),
                cleanup_epoch: 0,
                previously_tainted: false,
            })),
            cond: Default::default(),
            dice,
            exclusive_command_lock: Arc::new(ExclusiveCommandLock::new()),
        }
    }

    /// Enters a critical section that requires concurrent command synchronization,
    /// and runs the given `exec` function in the critical section.
    pub async fn enter<F, Fut, R>(
        &self,
        event_dispatcher: EventDispatcher,
        data: &dyn DiceDataProvider,
        updates: &dyn DiceUpdater,
        exec: F,
        is_nested_invocation: bool,
        sanitized_argv: Vec<String>,
        exclusive_cmd: Option<String>,
        exit_when_different_state: bool,
        cancellations: &ExplicitCancellationContext,
    ) -> anyhow::Result<R>
    where
        F: FnOnce(DiceTransaction) -> Fut,
        Fut: Future<Output = R> + Send,
    {
        let _exclusive_command_guard = event_dispatcher
            .span_async(
                ExclusiveCommandWaitStart {
                    command_name: self.exclusive_command_lock.owning_command(),
                },
                async move {
                    let guard = if let Some(cmd_name) = exclusive_cmd {
                        let guard = self.exclusive_command_lock.exclusive_lock(cmd_name).await;
                        self.dice.wait_for_idle().await;
                        guard
                    } else {
                        self.exclusive_command_lock.shared_lock().await
                    };
                    (guard, ExclusiveCommandWaitEnd {})
                },
            )
            .await;

        let events = event_dispatcher.dupe();
        let (_guard, transaction) = event_dispatcher
            .span_async(DiceSynchronizeSectionStart {}, async move {
                (
                    cancellations
                        .critical_section(|| {
                            self.wait_for_others(
                                data,
                                updates,
                                events,
                                is_nested_invocation,
                                sanitized_argv,
                                exit_when_different_state,
                            )
                        })
                        .await,
                    DiceSynchronizeSectionEnd {},
                )
            })
            .await?;

        Ok(exec(transaction).await)
    }

    // this is normally super unsafe, but because we are using an async condvar that takes care
    // of unlocking this mutex, this mutex is actually essentially never held across awaits.
    // The async condvar will handle properly allowing under threads to proceed, avoiding
    // starvation.
    async fn wait_for_others(
        &self,
        user_data: &dyn DiceDataProvider,
        updates: &dyn DiceUpdater,
        event_dispatcher: EventDispatcher,
        is_nested_invocation: bool,
        sanitized_argv: Vec<String>,
        exit_when_different_state: bool,
    ) -> anyhow::Result<(OnExecExit, DiceTransaction)> {
        let trace = event_dispatcher.trace_id().dupe();

        let span = tracing::span!(tracing::Level::DEBUG, "wait_for_others", trace = %trace);
        let _enter = span.enter();

        let mut data = self.data.lock().await;

        let command_id = data.next_command_id.increment();

        let command_data = CommandData {
            trace_id: trace.dupe(),
            argv: sanitized_argv,
            dispatcher: event_dispatcher.dupe(),
        };

        let (transaction, tainted) = loop {
            match &data.dice_status {
                DiceStatus::Cleanup { future, epoch } => {
                    tracing::debug!("ActiveDice is in cleanup");
                    let future = future.clone();
                    let epoch = *epoch;

                    // block while dice cleans up
                    drop(data);
                    event_dispatcher
                        .span_async(
                            buck2_data::DiceCleanupStart { epoch: epoch as _ },
                            async move { (future.await, buck2_data::DiceCleanupEnd {}) },
                        )
                        .await;
                    data = self.data.lock().await;

                    data.transition_to_idle(epoch);
                }
                DiceStatus::Available { active } => {
                    tracing::debug!("ActiveDice is available");

                    let dice_was_idle = self.dice.is_idle().await;

                    // we rerun the updates in case that files on disk have changed between commands.
                    // this might cause some churn, but concurrent commands don't happen much and
                    // isn't a big perf bottleneck. Dice should be able to resurrect nodes properly.

                    let transaction = async {
                        let updater = self.dice.updater();
                        let mut user_data = user_data
                            .provide(updater.existing_state().await.deref())
                            .await?;

                        let transaction = updates.update(updater, &mut user_data).await?;

                        event_dispatcher
                            .span_async(buck2_data::DiceStateUpdateStart {}, async {
                                (
                                    async {
                                        let transaction =
                                            transaction.commit_with_data(user_data).await;
                                        anyhow::Ok(transaction)
                                    }
                                    .await,
                                    buck2_data::DiceStateUpdateEnd {},
                                )
                            })
                            .await
                    }
                    .await?;

                    if let Some(active) = active {
                        let is_same_state = transaction.equivalent(&active.version);

                        event_dispatcher.instant_event(DiceConcurrentCommands {
                            total_concurrent_commands: data.active_commands.len() as u32,
                        });

                        // If we have a different state, attempt to transition to cleanup. This will
                        // succeed only if the current state is not in use.
                        if !is_same_state {
                            if data.transition_to_cleanup(&self.dice) {
                                continue;
                            }
                        }

                        tracing::debug!("ActiveDice has an active_transaction");

                        event_dispatcher.instant_event(DiceEqualityCheck {
                            is_equal: is_same_state,
                        });

                        let bypass_semaphore =
                            self.determine_bypass_semaphore(is_same_state, is_nested_invocation);

                        match bypass_semaphore {
                            BypassSemaphore::Error => {
                                return Err(anyhow::Error::new(
                                    ConcurrencyHandlerError::NestedInvocationWithDifferentStates(
                                        format_traces(&data.active_commands, &command_data),
                                        command_data.format_argv(),
                                    ),
                                ));
                            }
                            BypassSemaphore::Run(state) => {
                                self.emit_logs(state, &data.active_commands, &command_data)?;
                                break (transaction, false);
                            }
                            BypassSemaphore::Block => {
                                if exit_when_different_state {
                                    event_dispatcher.instant_event(ExitWhenDifferentState {});

                                    return Err(anyhow::Error::new(
                                        ConcurrencyHandlerError::ExitWhenDifferentState,
                                    ))
                                    .context(buck2_data::ErrorCause::DaemonIsBusy);
                                }
                                // We should probably show more than the first here, but for now
                                // this is what we have.
                                //
                                // Note: unwrap here relies on the fact that transition_to_cleanup
                                // would have transitioned if we had no active commands.

                                let active_command = data.active_commands.first().unwrap().1;
                                let trace_id = active_command.trace_id.dupe();
                                let argv = active_command.format_argv();

                                data = event_dispatcher
                                    .span_async(
                                        DiceBlockConcurrentCommandStart {
                                            current_active_trace_id: trace_id.to_string(),
                                            cmd_args: argv,
                                        },
                                        async {
                                            (
                                                self.cond.wait((data, &*self.data)).await,
                                                DiceBlockConcurrentCommandEnd {
                                                    ending_active_trace_id: trace_id.to_string(),
                                                },
                                            )
                                        },
                                    )
                                    .await;
                            }
                        }
                    } else {
                        tracing::debug!("ActiveDice has no active_transaction");
                        event_dispatcher.instant_event(NoActiveDiceState {});
                        data.dice_status = DiceStatus::active(transaction.equality_token());
                        break (transaction, !dice_was_idle);
                    }
                }
            }
        };

        tracing::info!("Acquired access to DICE");

        if data.previously_tainted {
            command_data.notify_previously_tainted();
        }

        if tainted {
            command_data.notify_tainted();
            data.notify_tainted();
            data.previously_tainted = true;
        }

        // create the on exit drop handler, which will take care of notifying tasks.
        let drop_guard = OnExecExit::new(self.dupe(), command_id, command_data, data);

        Ok((drop_guard, transaction))
    }

    /// Access dice without locking for dumps.
    pub fn unsafe_dice(&self) -> &Arc<Dice> {
        &self.dice
    }

    fn determine_bypass_semaphore(
        &self,
        is_same_state: bool,
        is_nested_invocation: bool,
    ) -> BypassSemaphore {
        if is_same_state {
            if is_nested_invocation {
                BypassSemaphore::Run(RunState::NestedSameState)
            } else {
                BypassSemaphore::Run(RunState::ParallelSameState)
            }
        } else if is_nested_invocation {
            BypassSemaphore::Error
        } else {
            BypassSemaphore::Block
        }
    }

    fn emit_logs(
        &self,
        state: RunState,
        active_commands: &SmallMap<CommandId, CommandData>,
        current_command: &CommandData,
    ) -> anyhow::Result<()> {
        let active_commands = format_traces(active_commands, current_command);

        match state {
            RunState::NestedSameState => {
                soft_error!(
                    "nested_invocation_same_dice_state",
                    anyhow::anyhow!(ConcurrencyHandlerError::NestedInvocationWithSameStates(
                        active_commands,
                        current_command.format_argv(),
                    ))
                )?;
            }
            _ => {}
        }

        Ok(())
    }
}

fn format_traces(
    active_commands: &SmallMap<CommandId, CommandData>,
    current: &CommandData,
) -> String {
    let trace_ids = active_commands
        .values()
        .chain(std::iter::once(current))
        .map(|cmd| &cmd.trace_id)
        .collect::<SmallSet<_>>();

    trace_ids.iter().join(", ")
}

/// Held to execute a command so that when the command is canceled, we properly remove its state
/// from the handler so that it's no longer registered as a ongoing command.
struct OnExecExit(Option<(ConcurrencyHandler, CommandId)>);

impl OnExecExit {
    pub fn new(
        handler: ConcurrencyHandler,
        command: CommandId,
        data: CommandData,
        mut guard: MutexGuard<'_, ConcurrencyHandlerData>,
    ) -> Self {
        guard.active_commands.insert(command, data);
        Self(Some((handler, command)))
    }
}

impl Drop for OnExecExit {
    fn drop(&mut self) {
        let this = self.0.take().expect("dropped twice");
        tracing::info!("Command has exited: {}", this.1);

        tokio::task::spawn(async move {
            let mut data = this.0.data.lock().await;
            data.active_commands
                .remove(&this.1)
                .expect("command was active but not in active_commands");
            tracing::info!("Active command was removed: {}", this.1);

            if data.has_no_active_commands() {
                // we notify all commands since we don't know how many can actually wake up and run
                // concurrently as several of the currently waiting commands could be "equivalent".
                // This could cause commands to wake up out of order and race, such that the longest
                // waiting command might not still be forced to wait. In reality, it is probably not
                // a terrible issue, as we are unlikely to have many concurrent commands, and people
                // are unlikely to usually care about the precise order they get to run.
                this.0.cond.notify_all()
            }
        });
    }
}

#[cfg(test)]
mod tests {
    use std::sync::atomic::AtomicBool;
    use std::sync::atomic::Ordering;
    use std::sync::Arc;
    use std::task::Poll;
    use std::time::Duration;

    use allocative::Allocative;
    use anyhow::Context;
    use assert_matches::assert_matches;
    use async_trait::async_trait;
    use buck2_core::is_open_source;
    use buck2_events::create_source_sink_pair;
    use buck2_events::dispatch::EventDispatcher;
    use buck2_events::span::SpanId;
    use buck2_events::BuckEvent;
    use buck2_events::EventSource;
    use buck2_wrapper_common::invocation_id::TraceId;
    use derivative::Derivative;
    use dice::DetectCycles;
    use dice::Dice;
    use dice::DiceComputations;
    use dice::DiceTransactionUpdater;
    use dice::InjectedKey;
    use dice::Key;
    use dice::UserComputationData;
    use dupe::Dupe;
    use futures::pin_mut;
    use futures::poll;
    use more_futures::cancellation::CancellationContext;
    use parking_lot::Mutex;
    use tokio::sync::Barrier;
    use tokio::sync::RwLock;
    use tokio::sync::Semaphore;

    use super::*;

    struct NoChanges;

    #[async_trait]
    impl DiceUpdater for NoChanges {
        async fn update(
            &self,
            ctx: DiceTransactionUpdater,
            _user_data: &mut UserComputationData,
        ) -> anyhow::Result<DiceTransactionUpdater> {
            Ok(ctx)
        }
    }

    struct CtxDifferent;

    #[async_trait]
    impl DiceUpdater for CtxDifferent {
        async fn update(
            &self,
            mut ctx: DiceTransactionUpdater,
            _user_data: &mut UserComputationData,
        ) -> anyhow::Result<DiceTransactionUpdater> {
            ctx.changed_to(vec![(K, ())])?;
            Ok(ctx)
        }
    }

    #[derive(Clone, Dupe, Display, Debug, Hash, Eq, PartialEq, Allocative)]
    struct K;

    #[async_trait]
    impl InjectedKey for K {
        type Value = ();

        fn equality(_x: &Self::Value, _y: &Self::Value) -> bool {
            false
        }
    }

    struct TestDiceDataProvider;

    #[async_trait]
    impl DiceDataProvider for TestDiceDataProvider {
        async fn provide(&self, _ctx: &DiceComputations) -> anyhow::Result<UserComputationData> {
            Ok(Default::default())
        }
    }

    #[tokio::test]
    async fn nested_invocation_same_transaction() {
        // FIXME: This times out on open source, and we don't know why
        if is_open_source() {
            return;
        }

        let dice = Dice::builder().build(DetectCycles::Enabled);

        let concurrency = ConcurrencyHandler::new(dice);

        let traces1 = TraceId::new();
        let traces2 = TraceId::new();
        let traces3 = TraceId::new();

        let barrier = Arc::new(Barrier::new(3));

        let fut1 = concurrency.enter(
            EventDispatcher::null_sink_with_trace(traces1),
            &TestDiceDataProvider,
            &NoChanges,
            |_| {
                let b = barrier.dupe();
                async move {
                    b.wait().await;
                }
            },
            true,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );
        let fut2 = concurrency.enter(
            EventDispatcher::null_sink_with_trace(traces2),
            &TestDiceDataProvider,
            &NoChanges,
            |_| {
                let b = barrier.dupe();
                async move {
                    b.wait().await;
                }
            },
            true,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );
        let fut3 = concurrency.enter(
            EventDispatcher::null_sink_with_trace(traces3),
            &TestDiceDataProvider,
            &NoChanges,
            |_| {
                let b = barrier.dupe();
                async move {
                    b.wait().await;
                }
            },
            true,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );

        let (r1, r2, r3) = futures::future::join3(fut1, fut2, fut3).await;
        r1.unwrap();
        r2.unwrap();
        r3.unwrap();
    }

    #[tokio::test]
    async fn nested_invocation_should_error() {
        let dice = Dice::builder().build(DetectCycles::Enabled);

        let concurrency = ConcurrencyHandler::new(dice);

        let traces1 = TraceId::new();
        let traces2 = TraceId::new();

        let barrier = Arc::new(Barrier::new(2));

        let fut1 = concurrency.enter(
            EventDispatcher::null_sink_with_trace(traces1),
            &TestDiceDataProvider,
            &NoChanges,
            |_| {
                let b = barrier.dupe();
                async move {
                    b.wait().await;
                }
            },
            true,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );

        let fut2 = concurrency.enter(
            EventDispatcher::null_sink_with_trace(traces2),
            &TestDiceDataProvider,
            &CtxDifferent,
            |_| {
                let b = barrier.dupe();
                async move {
                    b.wait().await;
                }
            },
            true,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );

        match futures::future::try_join(fut1, fut2).await {
            Err(e) => assert!(e.to_string().contains("Recursive invocation")),
            Ok(_) => {
                panic!("Futures should not have completed successfully")
            }
        }
    }

    #[tokio::test]
    async fn parallel_invocation_same_transaction() {
        let dice = Dice::builder().build(DetectCycles::Enabled);

        let concurrency = ConcurrencyHandler::new(dice);

        let traces1 = TraceId::new();
        let traces2 = TraceId::new();
        let traces3 = TraceId::new();

        let barrier = Arc::new(Barrier::new(3));

        let fut1 = concurrency.enter(
            EventDispatcher::null_sink_with_trace(traces1),
            &TestDiceDataProvider,
            &NoChanges,
            |_| {
                let b = barrier.dupe();
                async move {
                    b.wait().await;
                }
            },
            false,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );
        let fut2 = concurrency.enter(
            EventDispatcher::null_sink_with_trace(traces2),
            &TestDiceDataProvider,
            &NoChanges,
            |_| {
                let b = barrier.dupe();
                async move {
                    b.wait().await;
                }
            },
            false,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );
        let fut3 = concurrency.enter(
            EventDispatcher::null_sink_with_trace(traces3),
            &TestDiceDataProvider,
            &NoChanges,
            |_| {
                let b = barrier.dupe();
                async move {
                    b.wait().await;
                }
            },
            false,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );

        let (r1, r2, r3) = futures::future::join3(fut1, fut2, fut3).await;
        r1.unwrap();
        r2.unwrap();
        r3.unwrap();
    }

    #[tokio::test]
    async fn parallel_invocation_different_traceid_blocks() -> anyhow::Result<()> {
        let dice = Dice::builder().build(DetectCycles::Enabled);

        let concurrency = ConcurrencyHandler::new(dice.dupe());

        let traces1 = TraceId::new();
        let traces2 = traces1.dupe();
        let traces_different = TraceId::new();

        let block1 = Arc::new(RwLock::new(()));
        let blocked1 = block1.write().await;

        let block2 = Arc::new(RwLock::new(()));
        let blocked2 = block2.write().await;

        let barrier1 = Arc::new(Barrier::new(3));
        let barrier2 = Arc::new(Barrier::new(2));

        let arrived = Arc::new(AtomicBool::new(false));

        let fut1 = tokio::spawn({
            let concurrency = concurrency.dupe();
            let barrier = barrier1.dupe();
            let b = block1.dupe();

            async move {
                concurrency
                    .enter(
                        EventDispatcher::null_sink_with_trace(traces1),
                        &TestDiceDataProvider,
                        &NoChanges,
                        |_| async move {
                            barrier.wait().await;
                            let _g = b.read().await;
                        },
                        false,
                        Vec::new(),
                        None,
                        false,
                        ExplicitCancellationContext::testing(),
                    )
                    .await
            }
        });

        let fut2 = tokio::spawn({
            let concurrency = concurrency.dupe();
            let barrier = barrier1.dupe();
            let b = block2.dupe();

            async move {
                concurrency
                    .enter(
                        EventDispatcher::null_sink_with_trace(traces2),
                        &TestDiceDataProvider,
                        &NoChanges,
                        |_| async move {
                            barrier.wait().await;
                            let _g = b.read().await;
                        },
                        false,
                        Vec::new(),
                        None,
                        false,
                        ExplicitCancellationContext::testing(),
                    )
                    .await
            }
        });

        barrier1.wait().await;

        let fut3 = tokio::spawn({
            let concurrency = concurrency.dupe();
            let barrier = barrier2.dupe();
            let arrived = arrived.dupe();

            async move {
                barrier.wait().await;
                concurrency
                    .enter(
                        EventDispatcher::null_sink_with_trace(traces_different),
                        &TestDiceDataProvider,
                        &CtxDifferent,
                        |_| async move {
                            arrived.store(true, Ordering::Relaxed);
                        },
                        false,
                        Vec::new(),
                        None,
                        false,
                        ExplicitCancellationContext::testing(),
                    )
                    .await
            }
        });

        barrier2.wait().await;

        assert!(!arrived.load(Ordering::Relaxed));

        drop(blocked1);
        fut1.await??;

        assert!(!arrived.load(Ordering::Relaxed));

        drop(blocked2);
        fut2.await??;

        fut3.await??;

        assert!(arrived.load(Ordering::Relaxed));

        Ok(())
    }

    #[tokio::test]
    async fn parallel_invocation_exit_when_different_state() -> anyhow::Result<()> {
        let dice = Dice::builder().build(DetectCycles::Enabled);

        let concurrency = ConcurrencyHandler::new(dice.dupe());

        let traces1 = TraceId::new();
        let traces2 = traces1.dupe();
        let traces_different = TraceId::new();

        let block1 = Arc::new(RwLock::new(()));
        let blocked1 = block1.write().await;

        let block2 = Arc::new(RwLock::new(()));
        let blocked2 = block2.write().await;

        let barrier1 = Arc::new(Barrier::new(3));
        let barrier2 = Arc::new(Barrier::new(2));

        let arrived = Arc::new(AtomicBool::new(false));

        let fut1 = tokio::spawn({
            let concurrency = concurrency.dupe();
            let barrier = barrier1.dupe();
            let b = block1.dupe();

            async move {
                concurrency
                    .enter(
                        EventDispatcher::null_sink_with_trace(traces1),
                        &TestDiceDataProvider,
                        &NoChanges,
                        |_| async move {
                            barrier.wait().await;
                            let _g = b.read().await;
                        },
                        false,
                        Vec::new(),
                        None,
                        true,
                        ExplicitCancellationContext::testing(),
                    )
                    .await
            }
        });

        let fut2 = tokio::spawn({
            let concurrency = concurrency.dupe();
            let barrier = barrier1.dupe();
            let b = block2.dupe();

            async move {
                concurrency
                    .enter(
                        EventDispatcher::null_sink_with_trace(traces2),
                        &TestDiceDataProvider,
                        &NoChanges,
                        |_| async move {
                            barrier.wait().await;
                            let _g = b.read().await;
                        },
                        false,
                        Vec::new(),
                        None,
                        true,
                        ExplicitCancellationContext::testing(),
                    )
                    .await
            }
        });

        barrier1.wait().await;

        let fut3 = tokio::spawn({
            let concurrency = concurrency.dupe();
            let barrier = barrier2.dupe();
            let arrived = arrived.dupe();

            async move {
                barrier.wait().await;
                concurrency
                    .enter(
                        EventDispatcher::null_sink_with_trace(traces_different),
                        &TestDiceDataProvider,
                        &CtxDifferent,
                        |_| async move {
                            arrived.store(true, Ordering::Relaxed);
                        },
                        false,
                        Vec::new(),
                        None,
                        true,
                        ExplicitCancellationContext::testing(),
                    )
                    .await
            }
        });

        barrier2.wait().await;

        assert!(!arrived.load(Ordering::Relaxed));

        drop(blocked1);
        fut1.await??;

        assert!(!arrived.load(Ordering::Relaxed));

        drop(blocked2);
        fut2.await??;

        let fut3_result = fut3.await?;

        assert!(fut3_result.is_err());
        assert!(
            fut3_result
                .err()
                .unwrap()
                .to_string()
                .contains("daemon is busy")
        );

        Ok(())
    }

    #[derive(Clone, Dupe, Derivative, Allocative, Display)]
    #[derivative(Hash, Eq, PartialEq, Debug)]
    #[display(fmt = "CleanupTestKey")]
    struct CleanupTestKey {
        #[derivative(Debug = "ignore", Hash = "ignore", PartialEq = "ignore")]
        is_executing: Arc<Mutex<()>>,
    }

    #[async_trait::async_trait]
    impl Key for CleanupTestKey {
        type Value = ();

        #[allow(clippy::await_holding_lock)]
        async fn compute(
            &self,
            _ctx: &DiceComputations,
            cancellation: &CancellationContext,
        ) -> Self::Value {
            let _guard = self.is_executing.lock();

            // TODO: use critical_section as it's simpler, but this stack doesn't have it and
            // this works equally well here :)
            cancellation
                .with_structured_cancellation(|_obs| tokio::time::sleep(Duration::from_secs(1)))
                .await;
        }

        fn equality(_me: &Self::Value, _other: &Self::Value) -> bool {
            true
        }
    }

    #[tokio::test]
    async fn test_cleanup_stage() -> anyhow::Result<()> {
        let key = CleanupTestKey {
            is_executing: Arc::new(Mutex::new(())),
        };

        let key = &key;

        let dice = Dice::builder().build(DetectCycles::Enabled);

        let concurrency = ConcurrencyHandler::new(dice.dupe());

        // Kick off our computation and wait until it's running.

        concurrency
            .enter(
                EventDispatcher::null(),
                &TestDiceDataProvider,
                &NoChanges,
                |dice| async move {
                    let compute = dice.compute(key).fuse();

                    let started = async {
                        while !key.is_executing.is_locked() {
                            tokio::task::yield_now().await;
                        }
                    }
                    .fuse();

                    // NOTE: We still need to poll `compute` for it to actually spawn, hence the
                    // select below.

                    futures::pin_mut!(compute);
                    futures::pin_mut!(started);

                    futures::select! {
                        _ = compute => panic!("compute finished before started?"),
                        _ = started => {}
                    }
                },
                false,
                Vec::new(),
                None,
                false,
                ExplicitCancellationContext::testing(),
            )
            .await?;

        // Now, re-enter. We expect to reuse and therefore to not wait.

        concurrency
            .enter(
                EventDispatcher::null(),
                &TestDiceDataProvider,
                &NoChanges,
                |_dice| async move {
                    // The key should still be evaluating by now.
                    assert!(key.is_executing.is_locked());
                },
                false,
                Vec::new(),
                None,
                false,
                ExplicitCancellationContext::testing(),
            )
            .await?;

        // Now, enter with a different context. This time, we expect to not reuse.

        concurrency
            .enter(
                EventDispatcher::null(),
                &TestDiceDataProvider,
                &CtxDifferent,
                |_dice| async move {
                    assert!(!key.is_executing.is_locked());
                },
                false,
                Vec::new(),
                None,
                false,
                ExplicitCancellationContext::testing(),
            )
            .await?;

        Ok(())
    }

    async fn wait_for_event<T, F>(source: &mut T, matcher: Box<F>) -> anyhow::Result<BuckEvent>
    where
        T: EventSource,
        F: Fn(&BuckEvent) -> bool + Send,
    {
        tokio::time::timeout(Duration::from_millis(2), async {
            loop {
                if let Some(event) = source.try_receive() {
                    if let Some(event) = event.unpack_buck() {
                        if matcher(event) {
                            break event.clone();
                        }
                    }
                }
                tokio::task::yield_now().await;
            }
        })
        .await
        .context("Time out waiting for matching buck event")
    }

    async fn wait_for_exclusive_span_start<T>(
        source: &mut T,
        cmd: Option<&str>,
    ) -> anyhow::Result<Option<SpanId>>
    where
        T: EventSource,
    {
        let cmd = cmd.map(|c| c.to_owned());
        Ok(wait_for_event(
            source,
            Box::new(|e: &BuckEvent| {
                if let Some(span_start) = &e.span_start_event() {
                    if let Some(buck2_data::span_start_event::Data::ExclusiveCommandWait(data)) =
                        &span_start.data
                    {
                        let ExclusiveCommandWaitStart {
                            command_name: event_cmd,
                        } = data;
                        return event_cmd == &cmd;
                    }
                }
                false
            }),
        )
        .await?
        .span_id())
    }

    async fn wait_for_exclusive_span_end<T>(
        source: &mut T,
        span_id: Option<SpanId>,
    ) -> anyhow::Result<BuckEvent>
    where
        T: EventSource,
    {
        wait_for_event(
            source,
            Box::new(|e: &BuckEvent| {
                if let Some(span_end) = &e.span_end_event() {
                    if let Some(buck2_data::span_end_event::Data::ExclusiveCommandWait(_)) =
                        &span_end.data
                    {
                        return e.span_id() == span_id || span_id.is_none();
                    }
                }
                false
            }),
        )
        .await
    }

    #[tokio::test]
    async fn exclusive_command_lock() -> anyhow::Result<()> {
        let dice = Dice::builder().build(DetectCycles::Enabled);
        let concurrency = ConcurrencyHandler::new(dice.dupe());
        let (mut source, sink) = create_source_sink_pair();
        let dispatcher = EventDispatcher::new(TraceId::new(), sink);

        let mutex = Arc::new(Mutex::new(()));
        let command = |exclusive_cmd: Option<&str>, barriers: Option<&Arc<(Barrier, Barrier)>>| {
            tokio::spawn({
                let concurrency = concurrency.dupe();
                let dispatcher = dispatcher.dupe();
                let barriers = barriers.map(|b| b.dupe());
                let exclusive_cmd = exclusive_cmd.map(|b| b.to_owned());
                let mutex = mutex.dupe();
                async move {
                    concurrency
                        .enter(
                            dispatcher,
                            &TestDiceDataProvider,
                            &NoChanges,
                            |_| async move {
                                let _guard = mutex.try_lock().expect("Not exclusive!");
                                if let Some(barriers) = barriers {
                                    barriers.0.wait().await;
                                    barriers.1.wait().await;
                                }
                                tokio::task::yield_now().await;
                            },
                            false,
                            Vec::new(),
                            exclusive_cmd,
                            false,
                            ExplicitCancellationContext::testing(),
                        )
                        .await
                }
            })
        };

        let non_exclusive_barriers = Arc::new((Barrier::new(2), Barrier::new(2)));
        // Start non_exclusive command and enter critical section
        let non_exclusive_fut = command(None, Some(&non_exclusive_barriers.dupe()));
        non_exclusive_barriers.0.wait().await;

        let span_id_non_exclusive = wait_for_exclusive_span_start(&mut source, None).await?;
        wait_for_exclusive_span_end(&mut source, span_id_non_exclusive).await?;

        let command_barriers = Arc::new((Barrier::new(2), Barrier::new(2)));
        // Start exclusive command, blocked by non_exclusive
        let exclusive_fut_1 = command(Some("exclusive_1"), Some(&command_barriers.dupe()));

        let span_id_exclusive_1 = wait_for_exclusive_span_start(&mut source, None).await?;

        // Finish non_exclusive, enter exclusive_1 critical section
        non_exclusive_barriers.1.wait().await;
        non_exclusive_fut.await??;
        command_barriers.0.wait().await;

        wait_for_exclusive_span_end(&mut source, span_id_exclusive_1).await?;

        // Start series of exclusive commands and another second non_exclusive
        let exclusive_fut_2 = command(Some("exclusive_2"), None);
        let span_id_exclusive_2 =
            wait_for_exclusive_span_start(&mut source, Some("exclusive_1")).await?;
        let exclusive_fut_3 = command(Some("exclusive_3"), None);
        let span_id_exclusive_3 =
            wait_for_exclusive_span_start(&mut source, Some("exclusive_1")).await?;
        let non_exclusive_fut = command(None, None);
        let span_id_non_exclusive =
            wait_for_exclusive_span_start(&mut source, Some("exclusive_1")).await?;

        // Unblock first exclusive command, remaining commands are unblocked
        command_barriers.1.wait().await;
        exclusive_fut_1.await??;
        exclusive_fut_2.await??;
        exclusive_fut_3.await??;
        non_exclusive_fut.await??;

        wait_for_exclusive_span_end(&mut source, span_id_exclusive_2).await?;
        wait_for_exclusive_span_end(&mut source, span_id_exclusive_3).await?;
        wait_for_exclusive_span_end(&mut source, span_id_non_exclusive).await?;
        Ok(())
    }

    #[tokio::test]
    async fn test_thundering_herd() -> anyhow::Result<()> {
        let dice = Dice::builder().build(DetectCycles::Enabled);

        let concurrency = ConcurrencyHandler::new(dice.dupe());

        let concurrency = &concurrency;

        let tasks = (0..3).map(|_i| async {
            concurrency
                .enter(
                    EventDispatcher::null(),
                    &TestDiceDataProvider,
                    &CtxDifferent,
                    |dice| async move {
                        // NOTE: We need to actually compute something for DICE to be not-idle.
                        dice.compute(&K).await.unwrap();
                        tokio::task::yield_now().await;
                    },
                    false,
                    Vec::new(),
                    None,
                    false,
                    ExplicitCancellationContext::testing(),
                )
                .await
        });

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

        assert!(!concurrency.data.lock().await.previously_tainted);

        Ok(())
    }

    #[tokio::test]
    async fn test_updates_are_synchronized() -> anyhow::Result<()> {
        let dice = Dice::builder().build(DetectCycles::Enabled);

        let concurrency = ConcurrencyHandler::new(dice.dupe());

        struct Updater {
            should_be_able_to_run: AtomicBool,
            arrived_update: Semaphore,
        }

        #[async_trait]
        impl DiceUpdater for Updater {
            async fn update(
                &self,
                ctx: DiceTransactionUpdater,
                _user_data: &mut UserComputationData,
            ) -> anyhow::Result<DiceTransactionUpdater> {
                self.arrived_update.add_permits(1);
                tokio::task::yield_now().await;

                if self.should_be_able_to_run.load(Ordering::SeqCst) {
                    Ok(ctx)
                } else {
                    panic!("shouldn't be running")
                }
            }
        }

        let updater1 = Updater {
            should_be_able_to_run: AtomicBool::new(false),
            arrived_update: Semaphore::new(0),
        };
        let fut1 = concurrency.enter(
            EventDispatcher::null(),
            &TestDiceDataProvider,
            &updater1,
            |_dice| async move {
                tokio::task::yield_now().await;
            },
            false,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );
        pin_mut!(fut1);

        let updater2 = Updater {
            should_be_able_to_run: AtomicBool::new(false),
            arrived_update: Semaphore::new(0),
        };
        let fut2 = concurrency.enter(
            EventDispatcher::null(),
            &TestDiceDataProvider,
            &updater2,
            |_dice| async move {
                tokio::task::yield_now().await;
            },
            false,
            Vec::new(),
            None,
            false,
            ExplicitCancellationContext::testing(),
        );
        pin_mut!(fut2);

        // poll once will arrive at the `yield` in updater
        assert_matches!(poll!(&mut fut1), Poll::Pending);
        let _g = updater1.arrived_update.acquire().await?;

        // polling multiple times on the second command will all be pending and not complete
        assert_matches!(poll!(&mut fut2), Poll::Pending);
        assert_matches!(poll!(&mut fut2), Poll::Pending);
        assert_matches!(poll!(&mut fut2), Poll::Pending);
        assert_matches!(poll!(&mut fut2), Poll::Pending);

        // now make the first command runnable
        updater1.should_be_able_to_run.store(true, Ordering::SeqCst);
        fut1.await?;

        // 1 is done and dropped, so `2` can now finish
        updater2.should_be_able_to_run.store(true, Ordering::SeqCst);
        fut2.await?;

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

use std::collections::HashMap;
use std::future::Future;
use std::sync::Arc;

use async_trait::async_trait;
use buck2_build_signals::DeferredBuildSignals;
use buck2_build_signals::HasCriticalPathBackend;
use buck2_common::result::SharedResult;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::working_dir::WorkingDir;
use buck2_core::pattern::pattern_type::ConfiguredProvidersPatternExtra;
use buck2_core::pattern::ParsedPattern;
use buck2_data::CommandCriticalEnd;
use buck2_data::CommandCriticalStart;
use buck2_data::DiceCriticalSectionEnd;
use buck2_data::DiceCriticalSectionStart;
use buck2_events::dispatch::EventDispatcher;
use buck2_execute::materialize::materializer::Materializer;
use dice::DiceComputations;
use dice::DiceTransaction;
use dupe::Dupe;
use more_futures::cancellation::ExplicitCancellationContext;

use crate::concurrency::ConcurrencyHandler;
use crate::concurrency::DiceDataProvider;
use crate::concurrency::DiceUpdater;
use crate::stderr_output_guard::StderrOutputGuard;

#[async_trait]
pub trait ServerCommandContextTrait: Send + Sync {
    fn working_dir(&self) -> &ProjectRelativePath;

    fn working_dir_abs(&self) -> &WorkingDir;

    fn project_root(&self) -> &ProjectRoot;

    fn materializer(&self) -> Arc<dyn Materializer>;

    /// exposes the dice for scoped access, but isn't intended to be callable by anyone
    async fn dice_accessor(&self, private: PrivateStruct) -> SharedResult<DiceAccessor>;

    fn events(&self) -> &EventDispatcher;

    fn stderr(&self) -> anyhow::Result<StderrOutputGuard<'_>>;

    async fn request_metadata(&self) -> anyhow::Result<HashMap<String, String>>;

    async fn config_metadata(
        &self,
        ctx: &DiceComputations,
    ) -> anyhow::Result<HashMap<String, String>>;

    fn log_target_pattern(
        &self,
        providers_patterns: &[ParsedPattern<ConfiguredProvidersPatternExtra>],
    );

    fn cancellation_context(&self) -> &ExplicitCancellationContext;
}

pub struct PrivateStruct(());

pub struct DiceAccessor {
    pub dice_handler: ConcurrencyHandler,
    pub data: Box<dyn DiceDataProvider>,
    pub setup: Box<dyn DiceUpdater>,
    pub is_nested_invocation: bool,
    pub sanitized_argv: Vec<String>,
    pub exit_when_different_state: bool,
    pub build_signals: Box<dyn DeferredBuildSignals>,
}

#[async_trait]
pub trait ServerCommandDiceContext {
    async fn with_dice_ctx<'v, F, Fut, R>(&'v self, exec: F) -> anyhow::Result<R>
    where
        F: FnOnce(&'v dyn ServerCommandContextTrait, DiceTransaction) -> Fut + Send,
        Fut: Future<Output = anyhow::Result<R>> + Send,
        R: Send;

    async fn with_dice_ctx_maybe_exclusive<'v, F, Fut, R>(
        &'v self,
        exec: F,
        exclusive_cmd: Option<String>,
    ) -> anyhow::Result<R>
    where
        F: FnOnce(&'v dyn ServerCommandContextTrait, DiceTransaction) -> Fut + Send,
        Fut: Future<Output = anyhow::Result<R>> + Send,
        R: Send;
}

#[async_trait]
impl ServerCommandDiceContext for dyn ServerCommandContextTrait + '_ {
    /// Allows running a section of code that uses the shared DiceTransaction
    async fn with_dice_ctx<'v, F, Fut, R>(&'v self, exec: F) -> anyhow::Result<R>
    where
        F: FnOnce(&'v dyn ServerCommandContextTrait, DiceTransaction) -> Fut + Send,
        Fut: Future<Output = anyhow::Result<R>> + Send,
        R: Send,
    {
        self.with_dice_ctx_maybe_exclusive(exec, None).await
    }

    async fn with_dice_ctx_maybe_exclusive<'v, F, Fut, R>(
        &'v self,
        exec: F,
        exclusive_cmd: Option<String>,
    ) -> anyhow::Result<R>
    where
        F: FnOnce(&'v dyn ServerCommandContextTrait, DiceTransaction) -> Fut + Send,
        Fut: Future<Output = anyhow::Result<R>> + Send,
        R: Send,
    {
        let DiceAccessor {
            dice_handler,
            data,
            setup,
            is_nested_invocation,
            sanitized_argv,
            exit_when_different_state,
            build_signals,
        } = self.dice_accessor(PrivateStruct(())).await?;

        let events = self.events().dupe();
        events
            .span_async(DiceCriticalSectionStart {}, async move {
                (
                    dice_handler
                        .enter(
                            self.events().dupe(),
                            &*data,
                            &*setup,
                            |dice| async move {
                                let events = self.events().dupe();

                                let metadata = self.config_metadata(&dice).await?;

                                events
                                    .span_async(
                                        CommandCriticalStart {
                                            metadata: metadata.clone(),
                                            dice_version: dice.equality_token().to_string(),
                                        },
                                        async move {
                                            let res = buck2_build_signals::scope(
                                                build_signals,
                                                self.events().dupe(),
                                                dice.per_transaction_data()
                                                    .get_critical_path_backend(),
                                                || exec(self, dice),
                                            )
                                            .await;

                                            (res, CommandCriticalEnd { metadata })
                                        },
                                    )
                                    .await
                            },
                            is_nested_invocation,
                            sanitized_argv,
                            exclusive_cmd,
                            exit_when_different_state,
                            self.cancellation_context(),
                        )
                        .await,
                    DiceCriticalSectionEnd {},
                )
            })
            .await?
    }
}


/*
 * 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 command_end;
pub mod concurrency;
pub mod ctx;
pub mod logging;
pub mod partial_result_dispatcher;
pub mod pattern;
pub mod stderr_output_guard;
pub mod stdout_partial_output;
pub mod streaming_request_handler;
pub mod template;


/*
 * 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::fs::File;
use std::io;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::Mutex;

use dupe::Dupe;
use tracing_subscriber::fmt::MakeWriter;

static TRACING_LOG: Mutex<Option<TracingLogFile>> = Mutex::new(None);

#[derive(Clone, Dupe)]
pub struct TracingLogFile {
    file: Arc<PathBuf>,
    writer: Arc<Mutex<File>>,
}

impl TracingLogFile {
    pub fn new(file: PathBuf) -> anyhow::Result<Self> {
        let mut global = TRACING_LOG.lock().unwrap();
        Ok(if let Some(this) = &*global {
            this.dupe()
        } else {
            let writer = File::create(&file)?;

            let this = Self {
                file: Arc::new(file),
                writer: Arc::new(Mutex::new(writer)),
            };

            *global = Some(this.dupe());
            this
        })
    }

    pub fn refresh() -> anyhow::Result<()> {
        let this = TRACING_LOG.lock().unwrap();

        if let Some(this) = this.as_ref() {
            *this.writer.lock().unwrap() = File::create(&*this.file)?;
        }

        Ok(())
    }
}

impl<'a> io::Write for &'a TracingLogFile {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.writer.lock().unwrap().write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.writer.lock().unwrap().flush()
    }
}

impl<'a> MakeWriter<'a> for TracingLogFile {
    type Writer = &'a Self;

    fn make_writer(&'a self) -> Self::Writer {
        self
    }
}


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

use buck2_cli_proto::partial_result;
use buck2_cli_proto::PartialResult;
use buck2_events::dispatch::EventDispatcher;

use crate::stdout_partial_output::StdoutPartialOutput;

/// A typed partial result dispatcher. Each command can only send one kind of partial result, hence
/// the typing.
pub struct PartialResultDispatcher<T> {
    dispatcher: EventDispatcher,
    result_type: PhantomData<T>,
}

impl<T> PartialResultDispatcher<T>
where
    T: Into<partial_result::PartialResult>,
{
    pub fn new(dispatcher: EventDispatcher) -> Self {
        Self {
            dispatcher,
            result_type: PhantomData,
        }
    }

    /// NOTE: This doesn't actually require &mut self but that's been reasonable to have for the
    /// predecessor to this (stdout) so keeping it this way.
    pub fn emit(&mut self, res: T) {
        self.dispatcher.control_event(PartialResult {
            partial_result: Some(res.into()),
        });
    }
}

impl PartialResultDispatcher<buck2_cli_proto::StdoutBytes> {
    pub fn as_writer(&mut self) -> StdoutPartialOutput<'_> {
        StdoutPartialOutput::new(self)
    }
}

/// An uninhabited type for methods that do not produce partial results.
pub enum NoPartialResult {}

impl From<NoPartialResult> for partial_result::PartialResult {
    fn from(v: NoPartialResult) -> Self {
        match v {}
    }
}


/*
 * 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_cli_proto::ClientContext;
use buck2_common::dice::cells::HasCellResolver;
use buck2_common::target_aliases::BuckConfigTargetAliasResolver;
use buck2_common::target_aliases::HasTargetAliasResolver;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::CellResolver;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::pattern::pattern_type::PatternType;
use buck2_core::pattern::ParsedPattern;
use buck2_core::target::label::TargetLabel;
use dice::DiceComputations;
use dice::DiceTransaction;
use gazebo::prelude::*;

use crate::ctx::ServerCommandContextTrait;

pub struct PatternParser {
    cell_resolver: CellResolver,
    cwd: CellPath,
    target_alias_resolver: BuckConfigTargetAliasResolver,
}

impl PatternParser {
    pub async fn new(
        ctx: &mut DiceComputations,
        cwd: &ProjectRelativePath,
    ) -> anyhow::Result<Self> {
        let cell_resolver = ctx.get_cell_resolver().await?;

        let cwd = cell_resolver.get_cell_path(&cwd)?;
        let cell_name = cwd.cell();

        let target_alias_resolver = ctx.target_alias_resolver_for_cell(cell_name).await?;

        Ok(Self {
            cell_resolver,
            cwd,
            target_alias_resolver,
        })
    }

    pub fn parse_pattern<T: PatternType>(&self, pattern: &str) -> anyhow::Result<ParsedPattern<T>> {
        ParsedPattern::parse_relaxed(
            &self.target_alias_resolver,
            self.cwd.as_ref(),
            pattern,
            &self.cell_resolver,
        )
    }
}

/// Parse target patterns out of command line arguments.
///
/// The format allowed here is more relaxed than in build files and elsewhere, so only use this
/// with strings passed by the user on the CLI.
/// See `ParsedPattern::parse_relaxed` for details.
pub async fn parse_patterns_from_cli_args<T: PatternType>(
    ctx: &mut DiceComputations,
    target_patterns: &[buck2_data::TargetPattern],
    cwd: &ProjectRelativePath,
) -> anyhow::Result<Vec<ParsedPattern<T>>> {
    let parser = PatternParser::new(ctx, cwd).await?;

    target_patterns.try_map(|value| parser.parse_pattern(&value.value))
}

/// Extract target configuration (platform) label from [`ClientContext`].
pub async fn target_platform_from_client_context(
    client_ctx: &ClientContext,
    server_ctx: &dyn ServerCommandContextTrait,
    dice_ctx: &DiceTransaction,
) -> anyhow::Result<Option<TargetLabel>> {
    target_platform_from_client_context_impl(
        client_ctx,
        &dice_ctx.get_cell_resolver().await?,
        server_ctx.working_dir(),
    )
    .await
}

async fn target_platform_from_client_context_impl(
    client_context: &ClientContext,
    cell_resolver: &CellResolver,
    working_dir: &ProjectRelativePath,
) -> anyhow::Result<Option<TargetLabel>> {
    let cwd = cell_resolver.get_cell_path(working_dir)?;

    let target_platform = &client_context.target_platform;
    if !target_platform.is_empty() {
        Ok(Some(
            ParsedPattern::parse_precise(target_platform, cwd.cell(), cell_resolver)?
                .as_target_label(target_platform)?,
        ))
    } else {
        Ok(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.
 */

use std::io;
use std::io::BufWriter;
use std::io::Write;

use crate::partial_result_dispatcher::PartialResultDispatcher;

/// A wrapper that implements Write for a PartialResultDispatcher that emits StdoutBytes.
pub struct StdoutPartialOutput<'a> {
    inner: BufWriter<WriterWrapper<'a>>,
}

impl<'a> StdoutPartialOutput<'a> {
    pub fn new(dispatcher: &'a mut PartialResultDispatcher<buck2_cli_proto::StdoutBytes>) -> Self {
        Self {
            inner: BufWriter::new(WriterWrapper { inner: dispatcher }),
        }
    }
}

impl<'a> Write for StdoutPartialOutput<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.write(buf)
    }

    fn flush(&mut self) -> io::Result<()> {
        self.inner.flush()
    }
}

struct WriterWrapper<'a> {
    inner: &'a mut PartialResultDispatcher<buck2_cli_proto::StdoutBytes>,
}

impl<'a> Write for WriterWrapper<'a> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
        self.inner.emit(buck2_cli_proto::StdoutBytes {
            data: buf.to_owned(),
        });

        Ok(buf.len())
    }

    fn flush(&mut self) -> io::Result<()> {
        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.
 */

use std::marker::PhantomData;
use std::pin::Pin;
use std::task::Context;
use std::task::Poll;

use buck2_cli_proto::StreamingRequest;
use futures::Stream;
use pin_project::pin_project;
use tonic::Status;

#[derive(Debug, thiserror::Error)]
enum StreamingRequestError {
    #[error("Request returned error status: {0}")]
    GrpcStatus(Status),
    #[error("Unexpected EOF reading gRPC request")]
    UnexpectedEof,
}

/// Simple container that holds onto a stream of incoming client requests.
///
/// The primary use for this is pulling messages of a specific type from
/// the client via [`StreamingRequestHandler::message`]
#[pin_project]
pub struct StreamingRequestHandler<T: TryFrom<StreamingRequest, Error = anyhow::Error>> {
    #[pin]
    client_stream: tonic::Streaming<StreamingRequest>,
    _phantom: PhantomData<T>,
}

impl<T: TryFrom<StreamingRequest, Error = anyhow::Error>> StreamingRequestHandler<T> {
    pub fn new(client_stream: tonic::Streaming<StreamingRequest>) -> Self {
        Self {
            client_stream,
            _phantom: PhantomData,
        }
    }

    /// Get a message of type `T` from inside of a [`StreamingRequest`] envelope.
    ///
    /// Returns an error if the message is of the wrong type.
    pub async fn message(&mut self) -> anyhow::Result<T> {
        let request = match self.client_stream.message().await {
            Err(e) => return Err(StreamingRequestError::GrpcStatus(e).into()),
            Ok(Some(m)) => m,
            Ok(None) => return Err(StreamingRequestError::UnexpectedEof.into()),
        };
        request.try_into()
    }

    fn map_poll(v: Option<Result<StreamingRequest, Status>>) -> Option<anyhow::Result<T>> {
        match v {
            None => None,
            Some(Err(e)) => Some(Err(StreamingRequestError::GrpcStatus(e).into())),
            Some(Ok(v)) => match v.try_into() {
                Ok(v) => Some(Ok(v)),
                Err(e) => Some(Err(e)),
            },
        }
    }
}

impl<T: TryFrom<StreamingRequest, Error = anyhow::Error>> Stream for StreamingRequestHandler<T> {
    type Item = anyhow::Result<T>;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let this = self.project();
        let inner: Pin<&mut _> = this.client_stream;

        match inner.poll_next(cx) {
            Poll::Ready(v) => Poll::Ready(Self::map_poll(v)),
            Poll::Pending => Poll::Pending,
        }
    }
}


/*
 * 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 buck2_events::dispatch::span_async;
use dice::DiceTransaction;

use crate::command_end::command_end_ext;
use crate::ctx::ServerCommandContextTrait;
use crate::ctx::ServerCommandDiceContext;
use crate::logging::TracingLogFile;
use crate::partial_result_dispatcher::PartialResultDispatcher;

/// Typical server command with DICE and span.
#[async_trait]
pub trait ServerCommandTemplate: Send + Sync {
    /// Event to send in the beginning of command.
    type StartEvent: Into<buck2_data::command_start::Data> + Default;
    /// Event to send in the end of command.
    type EndEvent: Into<buck2_data::command_end::Data> + Default;
    /// Command return type.
    /// TODO: This is called `Result` everywhere, we should probably be consistent.
    type Response: Send;
    /// Command partial response.
    type PartialResult: Send + Sync;

    /// Create start event. Called before command is invoked.
    fn start_event(&self) -> Self::StartEvent {
        Self::StartEvent::default()
    }

    /// Create end event. Called after command is invoked.
    fn end_event(&self, _response: &anyhow::Result<Self::Response>) -> Self::EndEvent {
        Self::EndEvent::default()
    }

    /// Set `buck2_data::CommandEnd::is_success` to
    /// * `command` returns `Ok`
    /// * and this function returns `true`
    fn is_success(&self, response: &Self::Response) -> bool;

    /// If not `None`, command will block and be blocked by any concurrent commands.
    fn exclusive_command_name(&self) -> Option<String> {
        None
    }

    /// Command implementation.
    async fn command(
        &self,
        server_ctx: &dyn ServerCommandContextTrait,
        partial_result_dispatcher: PartialResultDispatcher<Self::PartialResult>,
        ctx: DiceTransaction,
    ) -> anyhow::Result<Self::Response>;
}

/// Call this function to run the command template implementation.
pub async fn run_server_command<T: ServerCommandTemplate>(
    command: T,
    server_ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<<T as ServerCommandTemplate>::PartialResult>,
) -> anyhow::Result<T::Response> {
    let metadata = server_ctx.request_metadata().await?;
    let start_event = buck2_data::CommandStart {
        metadata: metadata.clone(),
        data: Some(command.start_event().into()),
    };

    // refresh our tracing log per command
    TracingLogFile::refresh()?;

    span_async(start_event, async {
        let result = server_ctx
            .with_dice_ctx_maybe_exclusive(
                |server_ctx, dice| command.command(server_ctx, partial_result_dispatcher, dice),
                command.exclusive_command_name(),
            )
            .await;
        let end_event = command_end_ext(metadata, &result, command.end_event(&result), |result| {
            command.is_success(result)
        });
        (result, end_event)
    })
    .await
}


/*
 * 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 buck2_interpreter::starlark_debug::StarlarkDebugController;
use dupe::Dupe;
use starlark::debug::DapAdapterEvalHook;
use tokio::sync::OwnedSemaphorePermit;
use tracing::debug;

use crate::error::StarlarkDebuggerInternalError;
use crate::server::BuckStarlarkDebuggerServer;
use crate::HookId;

/// The debug controller is created once for each starlark evaluation. It adds the hooks to the
/// starlark Evaluator needed for debugging. Internally, the server will send and receive messages
/// with the DapAdapterEvalHooks that this injects to the Evaluator.
///
/// When evaluation finishes and the debug  controller is dropped, it will notify the server that it
/// has been dropped.
#[derive(Debug)]
pub struct BuckStarlarkDebugController {
    eval_wrapper: EvalWrapperHolder,
    server: Arc<BuckStarlarkDebuggerServer>,
    hook_id: HookId,
    description: String,
    _permit: OwnedSemaphorePermit,
    /// We don't want this object to be held across await points for a couple reasons: (1) it holds
    /// the semaphore permit and we want to limit the scope of where that needs to be held and (2) it's
    /// intended to wrap the starlark Evaluator which itself should not be sent across threads as it may
    /// hold a large starlark heap.
    _unsend: std::marker::PhantomData<*mut ()>,
}

impl StarlarkDebugController for BuckStarlarkDebugController {
    /// Initializes the Evaluator. This can only be used once for a particular controller.
    fn initialize(&mut self, eval: &mut starlark::eval::Evaluator) -> anyhow::Result<()> {
        match self.eval_wrapper.take()? {
            Some(v) => {
                debug!("adding dap hooks for {}", &self.description);
                v.add_dap_hooks(eval);
            }
            None => {
                // server is shutting down
            }
        };
        Ok(())
    }
}

impl BuckStarlarkDebugController {
    pub(crate) fn new(
        eval_wrapper: Option<Box<dyn DapAdapterEvalHook>>,
        hook_id: HookId,
        description: &str,
        server: &Arc<BuckStarlarkDebuggerServer>,
        permit: tokio::sync::OwnedSemaphorePermit,
    ) -> Self {
        Self {
            eval_wrapper: match eval_wrapper {
                Some(v) => EvalWrapperHolder::Ready(v),
                None => EvalWrapperHolder::ShuttingDown,
            },
            hook_id,
            description: description.to_owned(),
            server: server.dupe(),
            _permit: permit,
            _unsend: std::marker::PhantomData,
        }
    }
}

impl Drop for BuckStarlarkDebugController {
    fn drop(&mut self) {
        self.server.drop_hook(self.hook_id)
    }
}

#[derive(Debug)]
enum EvalWrapperHolder {
    /// The normal case, we have the eval hook ready to initialize an Evaluator.
    Ready(Box<dyn DapAdapterEvalHook>),
    /// Once initialize() has been called we'll change to the Used state (in which it'll be an
    /// error to initialize again).
    Used,
    /// Indicates that this was created while the debug server was already shutting down. That's
    /// an expected case and we just ignore it and don't add hooks to the Evaluator in that case.
    ShuttingDown,
}

impl EvalWrapperHolder {
    fn take(&mut self) -> anyhow::Result<Option<Box<dyn DapAdapterEvalHook>>> {
        match self {
            EvalWrapperHolder::ShuttingDown => Ok(None),
            EvalWrapperHolder::Used => {
                Err(StarlarkDebuggerInternalError::EvalWrapperAlreadyUsed.into())
            }
            EvalWrapperHolder::Ready(_) => {
                let owned = std::mem::replace(self, EvalWrapperHolder::Used);
                match owned {
                    EvalWrapperHolder::Ready(v) => Ok(Some(v)),
                    _ => unreachable!(),
                }
            }
        }
    }
}


/*
 * 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 contains the [DebugServer] trait which represents the part of the debug adapter
//! protocol implemented by the buck debug server and some utilities for working with
//! the debug adapter response/request/etc types.

use debugserver_types as dap;
use serde::Deserialize;
use serde::Serialize;

// TODO(cjhopman): debugserver_types is pretty out of date, we should switch to another
// crate (debug_types looked more up to date). Also it'd probably be nice to use a crate
// that has the dap types pre-generated (i.e. as part of the source) rather than
// one that generates them at build time so that the IDE experience is better.

/// This trait contains all the debug adapter requests that our debugserver supports.
///
/// See <https://microsoft.github.io/debug-adapter-protocol/specification> for the specification
/// of the debug adapter protocol.
///
/// This trait only has the one implementation, but it's split out here to make it clearer
/// what part of the protocol is implemented and for helpers like `dispatch` not needing to
/// handle the concrete type.
pub(crate) trait DebugServer {
    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Initialize>
    fn initialize(
        &mut self,
        x: dap::InitializeRequestArguments,
    ) -> anyhow::Result<Option<serde_json::Value>>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_SetBreakpoints>
    fn set_breakpoints(
        &mut self,
        x: dap::SetBreakpointsArguments,
    ) -> anyhow::Result<dap::SetBreakpointsResponseBody>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_SetExceptionBreakpoints>
    fn set_exception_breakpoints(
        &mut self,
        x: dap::SetExceptionBreakpointsArguments,
    ) -> anyhow::Result<()>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Attach>
    fn attach(&mut self, x: dap::AttachRequestArguments) -> anyhow::Result<()>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Threads>
    fn threads(&mut self) -> anyhow::Result<dap::ThreadsResponseBody>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_ConfigurationDone>
    fn configuration_done(&mut self) -> anyhow::Result<()>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_StackTrace>
    fn stack_trace(
        &mut self,
        x: dap::StackTraceArguments,
    ) -> anyhow::Result<dap::StackTraceResponseBody>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Scopes>
    fn scopes(&mut self, x: dap::ScopesArguments) -> anyhow::Result<dap::ScopesResponseBody>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Variables>
    fn variables(
        &mut self,
        x: dap::VariablesArguments,
    ) -> anyhow::Result<dap::VariablesResponseBody>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Continue>
    fn continue_(&mut self, x: ContinueArguments) -> anyhow::Result<dap::ContinueResponseBody>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Next>
    fn next(&mut self, x: dap::NextArguments) -> anyhow::Result<()>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_StepIn>
    fn step_in(&mut self, x: dap::StepInArguments) -> anyhow::Result<()>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_StepOut>
    fn step_out(&mut self, x: dap::StepOutArguments) -> anyhow::Result<()>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Evaluate>
    fn evaluate(&mut self, x: dap::EvaluateArguments) -> anyhow::Result<dap::EvaluateResponseBody>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Disconnect>
    fn disconnect(&mut self, x: dap::DisconnectArguments) -> anyhow::Result<()>;

    /// See <https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Source>
    fn source(&mut self, x: dap::SourceArguments) -> anyhow::Result<dap::SourceResponseBody>;
}

/// DAP ContinueArguments (debugserver_types is missing single_thread field)
#[derive(Debug, Eq, PartialEq, Clone, Deserialize, Serialize)]
#[serde(rename_all = "camelCase")]
pub(crate) struct ContinueArguments {
    pub(crate) thread_id: i64,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub(crate) single_thread: Option<bool>,
}

/// Create a dap Event with the given body. The user is responsible for updating the `seq` field.
pub(crate) fn dap_event<T: Serialize>(event: &str, body: Option<&T>) -> dap::Event {
    dap::Event {
        type_: "event".to_owned(),
        seq: 0,
        event: event.to_owned(),
        body: body.map(|v| serde_json::json!(v)),
    }
}

/// Create a Response for the Request indicating an error. The user is responsible for updating the `seq` field.
pub(crate) fn err_response(req: &dap::Request, err: &anyhow::Error) -> dap::Response {
    dap::Response {
        type_: "response".to_owned(),
        command: req.command.clone(),
        request_seq: req.seq,
        seq: 0,
        success: false,
        message: Some(format!("{:#}", err)),
        body: None,
    }
}

pub(crate) fn dispatch(
    server: &mut impl DebugServer,
    r: &dap::Request,
) -> anyhow::Result<dap::Response> {
    fn arg<T: for<'a> Deserialize<'a>>(r: &dap::Request) -> anyhow::Result<T> {
        Ok(serde_json::from_value(r.arguments.clone().ok_or_else(
            || anyhow::anyhow!("missing expected argument in DAP request"),
        )?)?)
    }

    fn ret<T: Serialize>(
        r: &dap::Request,
        v: anyhow::Result<Option<T>>,
    ) -> anyhow::Result<dap::Response> {
        let v = v?;
        Ok(dap::Response {
            type_: "response".to_owned(),
            command: r.command.clone(),
            request_seq: r.seq,
            seq: 0,
            success: true,
            message: None,
            body: v.map(|v| serde_json::to_value(v).unwrap()),
        })
    }

    fn ret_some<T: Serialize>(
        r: &dap::Request,
        v: anyhow::Result<T>,
    ) -> anyhow::Result<dap::Response> {
        ret(r, v.map(|v| Some(v)))
    }

    fn ret_none(r: &dap::Request, v: anyhow::Result<()>) -> anyhow::Result<dap::Response> {
        ret::<()>(r, v.map(|_| None))
    }

    match r.command.as_str() {
        "initialize" => ret(r, server.initialize(arg(r)?)),
        "setBreakpoints" => ret_some(r, server.set_breakpoints(arg(r)?)),
        "setExceptionBreakpoints" => ret_none(r, server.set_exception_breakpoints(arg(r)?)),
        "attach" => ret_none(r, server.attach(arg(r)?)),
        "threads" => ret_some(r, server.threads()),
        "configurationDone" => ret_none(r, server.configuration_done()),
        "stackTrace" => ret_some(r, server.stack_trace(arg(r)?)),
        "scopes" => ret_some(r, server.scopes(arg(r)?)),
        "variables" => ret_some(r, server.variables(arg(r)?)),
        "continue" => ret_some(r, server.continue_(arg(r)?)),
        "evaluate" => ret_some(r, server.evaluate(arg(r)?)),
        "disconnect" => ret_none(r, server.disconnect(arg(r)?)),
        "source" => ret_some(r, server.source(arg(r)?)),
        "next" => ret_none(r, server.next(arg(r)?)),
        "stepIn" => ret_none(r, server.step_in(arg(r)?)),
        "stepOut" => ret_none(r, server.step_out(arg(r)?)),
        _ => Err(anyhow::anyhow!(
            "Buck2 debugserver didn't recognize command: {}",
            r.command
        )),
    }
}


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

/// Errors from buck's starlark debugger
#[derive(Debug, Error)]
pub(crate) enum StarlarkDebuggerError {
    #[error("starlark debugger has not yet implemented this functionality")]
    Unimplemented,
    #[error("another debug-attach process is already attached")]
    DebuggerAlreadyAttached,
}

/// Internal errors from buck's starlark debugger
#[derive(Debug, Error)]
pub(crate) enum StarlarkDebuggerInternalError {
    #[error("Internal error: debbugger server shutdown unexpectedly")]
    UnexpectedDebuggerShutdown,
    #[error("Internal error: debug controller already initialized with Evaluator")]
    EvalWrapperAlreadyUsed,
}


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

#![allow(rustdoc::private_intra_doc_links)]

//! Provides the (daemon-side) support for buck2's starlark debugging.
//!
//! Components of the debugger:
//!
//! [BuckStarlarkDebuggerServer] is the main way that the core of buck2 integrates
//! the starlark debugger. This provides the hooks to wrap a Starlark evaluation and
//! enable the debugger (and handle communication between that starlark evaluation
//! and the debugger server/state)
//!
//! [BuckStarlarkDebuggerHandle] is a "handle" the to the debugger server. One of these
//! will be created for each buck2 command and put in the dice per-transaction data. Code
//! that needs to do starlark evaluation can then use this to setup their Evaluator
//! appropriately (though this is really just an implementation detail hidden in the
//! helper [buck2_interpreter::dice::starlark_provider::with_starlark_eval_provider]).
//!
//! [CURRENT_DEBUGGER] holds a global reference for the currently attached debugger and
//! vends out handles (through [create_debugger_handle]).
//!
//! ServerState the internal implementation of the starlark debugger server sends all
//! requests and starlark evaluation events to the single-threaded ServerState.
//!
//! Information about the debugger behavior:
//!
//! The debugger server exposes all current starlark evaluations as separate threads to
//! the DAP client. Then, one of it's main responsibility is as a mux/demux for requests/events/etc
//! flowing between the DAP client and the starlark evaluation.
//!
//! The server also tracks the state of all evaluations so that it can send regularly
//! debugger snapshots to all current buck commands. Those commands in turn use those snapshots
//! to render a debugger superconsole component so that users aren't left confused about
//! why their commands are seemingly hanging.
//!
//! The server handles a variety of other small adapter functionality. For example,
//! the DAP client may send source file references as absolute paths, but our starlark
//! evaluators will only understand them as the project-relative paths that buck
//! uses and so the server must convert them from/to the client/starlark.

use std::sync::Arc;
use std::sync::Mutex;

use async_trait::async_trait;
use buck2_core::fs::project::ProjectRoot;
use buck2_events::dispatch::EventDispatcher;
use buck2_interpreter::starlark_debug::StarlarkDebugController;
use buck2_interpreter::starlark_debug::StarlarkDebuggerHandle;
use derive_more::Display;
use dupe::Dupe;
use tokio::sync::mpsc;

use crate::error::StarlarkDebuggerError;
use crate::run::ToClientMessage;
use crate::server::BuckStarlarkDebuggerServer;

mod controller;
mod dap_api;
mod error;
pub mod run;
mod server;

/// A handle to the debugger server.
#[derive(Debug, Clone, Dupe)]
pub struct BuckStarlarkDebuggerHandle(Arc<HandleData>);

#[derive(Debug)]
pub struct HandleData {
    id: HandleId,
    server: Arc<BuckStarlarkDebuggerServer>,
}

impl Drop for HandleData {
    fn drop(&mut self) {
        self.server.drop_handle(self.id)
    }
}

#[async_trait]
impl StarlarkDebuggerHandle for BuckStarlarkDebuggerHandle {
    async fn start_eval(
        &self,
        description: &str,
    ) -> anyhow::Result<Box<dyn StarlarkDebugController>> {
        self.0.server.start_eval(self, description).await
    }
}

#[derive(Debug, Clone, Dupe, Copy, Eq, PartialEq, Hash)]
pub(crate) struct HandleId(u32);

#[derive(Display, Debug, Clone, Dupe, Copy, Eq, PartialEq, Hash)]
pub(crate) struct HookId(u32);

/// We allow only a single debugger to be attached at a time. While it's attached, this will hold the server instance.
static CURRENT_DEBUGGER: Mutex<Option<Arc<BuckStarlarkDebuggerServer>>> = Mutex::new(None);

/// Used by each command to get a handle to the current debugger. The debugger server will capture the
/// event dispatcher to send back debugger state snapshots (which indicate that the debugger is attached
/// and which, if any, threads are paused) while the command is running.
pub fn create_debugger_handle(events: EventDispatcher) -> Option<BuckStarlarkDebuggerHandle> {
    CURRENT_DEBUGGER
        .lock()
        .unwrap()
        .as_ref()
        .and_then(|v| v.new_handle(events))
}

/// Manages setting/unsetting the CURRENT_DEBUGGER while the starlark debug-attach command is running.
struct ServerConnection(Arc<BuckStarlarkDebuggerServer>);

impl ServerConnection {
    fn new(
        to_client_send: mpsc::UnboundedSender<ToClientMessage>,
        project_root: ProjectRoot,
    ) -> anyhow::Result<Self> {
        let mut locked = CURRENT_DEBUGGER.lock().unwrap();
        if locked.is_some() {
            return Err(StarlarkDebuggerError::DebuggerAlreadyAttached.into());
        }

        let server = Arc::new(BuckStarlarkDebuggerServer::new(
            to_client_send,
            project_root,
        ));
        *locked = Some(server.dupe());
        Ok(Self(server))
    }
}

impl Drop for ServerConnection {
    fn drop(&mut self) {
        *CURRENT_DEBUGGER.lock().unwrap() = None;
        let _ignored = self.0.detach();
    }
}


/*
 * 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::span_async;
use buck2_server_ctx::command_end::command_end;
use buck2_server_ctx::ctx::ServerCommandContextTrait;
use buck2_server_ctx::partial_result_dispatcher::PartialResultDispatcher;
use buck2_server_ctx::streaming_request_handler::StreamingRequestHandler;
use debugserver_types as dap;
use tokio::select;
use tokio::sync::mpsc;
use tokio_stream::StreamExt;
use tracing::debug;

use crate::error::StarlarkDebuggerInternalError;
use crate::ServerConnection;

/// Messages from the debugger server to its client (the cli `buck2 starlark debug-attach` which
/// then forwards them along through its stdout).
#[derive(Debug)]
pub(crate) enum ToClientMessage {
    Event(dap::Event),
    Response(dap::Response),
    Shutdown(anyhow::Result<()>),
}

impl ToClientMessage {
    fn pretty_string(&self) -> anyhow::Result<String> {
        match self {
            ToClientMessage::Event(ev) => Ok(serde_json::to_string_pretty(&ev)?),
            ToClientMessage::Response(resp) => Ok(serde_json::to_string_pretty(&resp)?),
            ToClientMessage::Shutdown(res) => Ok(format!("{:?}", res)),
        }
    }
}

/// Wraps `run_dap_server` with a command start/end span.
pub async fn run_dap_server_command(
    ctx: &dyn ServerCommandContextTrait,
    partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::DapMessage>,
    req: StreamingRequestHandler<buck2_cli_proto::DapRequest>,
) -> anyhow::Result<buck2_cli_proto::DapResponse> {
    let metadata = ctx.request_metadata().await?;
    let start_event = buck2_data::CommandStart {
        metadata: metadata.clone(),
        data: Some(buck2_data::StarlarkDebugAttachCommandStart {}.into()),
    };
    span_async(start_event, async move {
        let result = run_dap_server(ctx, partial_result_dispatcher, req).await;
        let end_event = command_end(
            metadata,
            &result,
            buck2_data::StarlarkDebugAttachCommandEnd {},
        );
        (result, end_event)
    })
    .await
}

async fn run_dap_server(
    ctx: &dyn ServerCommandContextTrait,
    mut partial_result_dispatcher: PartialResultDispatcher<buck2_cli_proto::DapMessage>,
    mut req: StreamingRequestHandler<buck2_cli_proto::DapRequest>,
) -> anyhow::Result<buck2_cli_proto::DapResponse> {
    let (to_client_send, mut to_client_recv) = mpsc::unbounded_channel();
    let server_connection = ServerConnection::new(to_client_send, ctx.project_root().clone())?;

    let mut seq = 0;

    // We can get requests from the client (e.g. vscode) and we can get responses/events from the server
    // and mostly we just forward them along to the other side.
    // If an error is encountered or the server sends the Shutdown message, we'll break out, disconnect
    // from the server and return the final response.
    let response = loop {
        select! {
            request = req.next() => {
                let request = match request {
                    Some(Err(e)) => return Err(e),
                    Some(Ok(v)) => v,
                    None => {
                        // client disconnected.
                        break buck2_cli_proto::DapResponse {};
                    }
                };

                let debugserver_req: dap::Request = serde_json::from_str(&request.dap_json)?;
                debug!("received request {}", &serde_json::to_string_pretty(&debugserver_req)?);
                server_connection.0.send_request(debugserver_req)?;
            }
            message = to_client_recv.recv() => {
                match message {
                    Some(v) => {
                        if let Some(v) = handle_outgoing_message(&mut seq, &mut partial_result_dispatcher, v)? {
                            break v;
                        }
                    }
                    None => return Err(StarlarkDebuggerInternalError::UnexpectedDebuggerShutdown.into())
                }
            }
        }
    };

    debug!("returning debugserver response");
    Ok(response)
}

/// Converts the ToClientMessage to either a partial result (and dispatches it) and/or the final
/// response (for a Shutdown message).
fn handle_outgoing_message(
    seq: &mut u32,
    partial_result_dispatcher: &mut PartialResultDispatcher<buck2_cli_proto::DapMessage>,
    message: ToClientMessage,
) -> anyhow::Result<Option<buck2_cli_proto::DapResponse>> {
    debug!("sending message {}", &message.pretty_string()?);

    let this_seq = *seq as i64;
    *seq += 1;

    let (dap_json, response) = match message {
        ToClientMessage::Event(mut ev) => {
            ev.seq = this_seq;
            (serde_json::to_vec(&ev)?, None)
        }
        ToClientMessage::Response(mut resp) => {
            resp.seq = this_seq;
            (serde_json::to_vec(&resp)?, None)
        }
        ToClientMessage::Shutdown(res) => {
            let exit_code = match res {
                Ok(..) => 0,
                Err(e) => {
                    debug!("server shutdown with error {:?}", &e);
                    1
                }
            };

            (
                serde_json::to_vec(&dap::ExitedEvent {
                    type_: "event".to_owned(),
                    seq: this_seq,
                    event: "exited".to_owned(),
                    body: dap::ExitedEventBody { exit_code },
                })?,
                Some(buck2_cli_proto::DapResponse {}),
            )
        }
    };

    partial_result_dispatcher.emit(buck2_cli_proto::DapMessage { dap_json });
    Ok(response)
}