laburnum 1.17.2

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use {
  crate::{
    Ident,
    Partitions,
    TRACER,
    connect::{
      ipc::Connection,
      lsp::{
        ClientId,
        notification::Notification,
        request::Request,
      },
    },
    database::{
      Database,
      GenerationEpoch,
      RecordKey,
      chunk::RecordWriter,
      query::{QueryClient, SortKeyCondition},
    },
    progress::ProgressTracker,
    protocol::{
      jsonrpc,
      lsp::LanguageServer,
      otel::TraceContext,
    },
    scheduler::{
      Scheduler,
      lanes::Lane,
    },
  },
  parking_lot::Mutex,
  std::{
    future::Future,
    marker::PhantomData,
    pin::Pin,
    sync::{
      Arc,
      atomic::{
        AtomicU64,
        Ordering,
      },
    },
    task::{
      Context,
      Poll,
    },
  },
};

static TASK_COUNTER: AtomicU64 = AtomicU64::new(0);

fn generate_execution_id() -> Ident {
  let id = TASK_COUNTER.fetch_add(1, Ordering::Relaxed);
  Ident::from_hash(id)
}

/// Whether a task counts toward scheduler quiescence.
///
/// `User` tasks are compilation work — parse, watchers, follow-ups. They
/// increment the scheduler's `work_in_flight` counter at construction and
/// decrement it on completion; when the count hits zero the pipeline is
/// idle and progress can end.
///
/// `System` tasks are laburnum's own infrastructure — the RPC message
/// pump, the daemon accept loop, periodic GC, the idle debounce. These
/// run for the lifetime of the server (or never complete), so they are
/// excluded from the counter: counting them would pin `work_in_flight`
/// above zero forever and quiescence would never be observed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum TaskClass {
  System,
  User,
}

/// A Laburnum task executes an async future that processes source files and
/// writes records to the database.
///
/// Tasks are the fundamental unit of work in the Laburnum scheduler. They:
/// - Execute async futures that process language server operations
/// - Write results as chunks to the database
/// - Support automatic cancellation when source versions become unavailable
/// - Are scheduled on priority lanes for execution order control
///
/// ## Task Lifecycle
///
/// 1. **Creation**: Tasks are spawned via `Scheduler::queue()` or
///    `TaskContext::spawn_task()`
/// 2. **Execution**: The scheduler polls tasks via `poll_once()` on worker
///    threads
/// 3. **Completion**: When the future completes with `Some(RecordWriter)`, a
///    chunk is written to the database
/// 4. **Cancellation**: Tasks return `None` if cancelled or if source data is
///    no longer available
///
/// ## Task Cancellation
///
/// Tasks automatically cancel when they attempt to access source data that has
/// been evicted from the cache. This prevents wasted computation on stale data
/// when newer versions of files are available.
///
/// ### Cancellation Example
///
/// ```rust
/// use laburnum::{
///   LaburnumError,
///   database::chunk::RecordWriter,
/// };
///
/// async fn on_file_version(
///   uri: laburnum::Uri,
///   source: std::sync::Arc<laburnum::source::Source>,
///   source_key: laburnum::SourceKey,
///   ctx: laburnum::scheduler::task::TaskContext<P, T>,
///   writer: RecordWriter<P>,
/// ) -> Result<RecordWriter<P>, LaburnumError> {
///   let source_cache = ctx.source_cache_reader();
///
///   // Option 1: Explicit availability check
///   ctx.check_source_available(source_key)?;
///
///   // Option 2: Get source (returns Err(TaskCancelled) if unavailable)
///   let source = source_cache.get_source_or_cancel(source_key)?;
///
///   // Continue processing...
///   Ok(writer)
/// }
/// ```
///
/// When a task returns `Err(LaburnumError::TaskCancelled)`, the error is caught
/// in the task wrapper and logged at DEBUG level. No chunk is written,
/// preventing stale data from entering the database.
#[allow(clippy::type_complexity)]
pub(crate) struct LaburnumTask<P: Partitions, T: LanguageServer<P>> {
  future: Mutex<
    Option<Pin<Box<dyn Future<Output = Option<RecordWriter<P>>> + Send>>>,
  >,
  pub(crate) lane:      Lane,
  pub(crate) scheduler: Arc<Scheduler<P, T>>,
  trace_context:        TraceContext,
  /// Whether this task counts toward scheduler quiescence (see
  /// [`TaskClass`]). `User` tasks bump `work_in_flight`; `System` tasks
  /// do not.
  task_class:           TaskClass,
  /// The epoch at which this task was first polled. Used for snapshot
  /// isolation — the reaper won't reclaim records from this epoch until
  /// this task completes.
  birth_epoch:          Mutex<Option<GenerationEpoch>>,
  _phantom:             PhantomData<T>,
}

impl<P: Partitions, T: LanguageServer<P>> LaburnumTask<P, T> {
  /// Polls the task future once and handles the result.
  ///
  /// # Returns
  ///
  /// - `Poll::Ready(())`: Task completed (chunk written or cancelled)
  /// - `Poll::Pending`: Task yielded, needs more polling
  ///
  /// # Behavior
  ///
  /// If the future completes:
  /// - `Some(RecordWriter)`: Builds chunk, writes to database, triggers
  ///   watchers
  /// - `None`: Task cancelled, no chunk written
  ///
  /// If the future is pending, it remains in the scheduler for future polling.
  pub(crate) fn poll_once(self: &Arc<Self>) -> Poll<()> {
    let _guard = self.trace_context.attach();

    // Register birth epoch on first poll for snapshot isolation
    {
      let mut epoch_slot = self.birth_epoch.lock();
      if epoch_slot.is_none() {
        let epoch = self.scheduler.db.get_current_epoch();
        *epoch_slot = Some(epoch);
        self.scheduler.register_active_epoch(epoch);
      }
    }

    let mut future_slot = self.future.lock();

    if let Some(mut future) = future_slot.take() {
      let waker = futures::task::waker_ref(self);

      let mut cx = Context::from_waker(&waker);

      match future.as_mut().poll(&mut cx) {
        | Poll::Ready(Some(chunk_builder)) => {
          let mut chunk = chunk_builder.build();

          if !chunk.is_empty() || !chunk.clear_prefixes().is_empty() || chunk.has_pending_sources() {
            otel::span!("laburnum.task.complete", in |cx|{
              let task_id = chunk.task_id();
              let db = self.scheduler.db.clone();

              // Extract data from chunk before commit takes ownership
              let error_source_keys = chunk.error_diagnostic_source_keys();
              let pending_sources = chunk.take_sources();

              // Single call — all index logic is inside commit_chunk now
              let source_cache_reader = self.scheduler.source_cache.read().reader();
              let mut result = db.commit_chunk(chunk, &source_cache_reader);

              // Feed deferred decrements to the reaper before on_new_chunk
              // consumes the result for watcher dispatch
              let deferred_decrements = std::mem::take(&mut result.deferred_decrements);
              if !deferred_decrements.is_empty() {
                self.scheduler.reaper.queue_decrements(deferred_decrements);
              }

              // Write barrier: if GC is marking, add new hashes to gray set
              // to prevent the collector from missing newly referenced records
              if self.scheduler.gc.is_marking() {
                self.scheduler.gc.add_to_gray(result.new_hashes.iter().copied());
              }

              // Complete any pending sources carried through the chunk
              if !pending_sources.is_empty() {
                let mut source_cache = self.scheduler.source_cache.write();
                for (source_key, source) in pending_sources {
                  source_cache.complete_version(source_key, source);
                }
              }

              // Mark source keys with error diagnostics
              if !error_source_keys.is_empty() {
                let mut source_cache = self.scheduler.source_cache.write();
                for source_key in &error_source_keys {
                  source_cache.set_has_errors(*source_key, true);
                }
              }

              // Watcher notifications
              if result.has_changes() {
                self.scheduler.on_new_chunk(task_id, result);
              }
            });
          }

          self.deregister_epoch();
          self.on_complete();
          Poll::Ready(())
        },
        | Poll::Ready(None) => {
          self.deregister_epoch();
          self.on_complete();
          Poll::Ready(())
        },
        | Poll::Pending => {
          *future_slot = Some(future);
          Poll::Pending
        },
      }
    } else {
      Poll::Ready(())
    }
  }

  /// Called once when the task reaches `Poll::Ready`. For `User` tasks this
  /// decrements `work_in_flight`, arming the idle debounce if the count
  /// reaches zero (the scheduler is now quiescent). `System` tasks do not
  /// participate in the counter.
  fn on_complete(&self) {
    if self.task_class == TaskClass::User {
      self.scheduler.user_task_finished();
    }
  }

  /// Deregister this task's birth epoch from the scheduler's active epoch registry.
  fn deregister_epoch(&self) {
    if let Some(epoch) = self.birth_epoch.lock().take() {
      self.scheduler.deregister_active_epoch(epoch);
    }
  }

  pub(crate) fn new<F, Fut>(
    scheduler: Arc<Scheduler<P, T>>,
    task_fn: F,
    lane: Lane,
    client_id: ClientId,
  ) -> Arc<Self>
  where
    F: FnOnce(TaskContext<P, T>) -> Fut + Send + 'static,
    Fut: Future<Output = Option<RecordWriter<P>>> + Send + 'static,
  {
    Self::build(scheduler, task_fn, lane, None, client_id, TaskClass::User)
  }

  pub(crate) fn new_with_parent<F, Fut>(
    scheduler: Arc<Scheduler<P, T>>,
    task_fn: F,
    lane: Lane,
    parent_task_id: Option<crate::Ident>,
    client_id: ClientId,
  ) -> Arc<Self>
  where
    F: FnOnce(TaskContext<P, T>) -> Fut + Send + 'static,
    Fut: Future<Output = Option<RecordWriter<P>>> + Send + 'static,
  {
    Self::build(
      scheduler,
      task_fn,
      lane,
      parent_task_id,
      client_id,
      TaskClass::User,
    )
  }

  /// Construct a `System` task — laburnum infrastructure that runs for the
  /// server's lifetime and must NOT count toward quiescence (see
  /// [`TaskClass`]).
  pub(crate) fn new_system_with_parent<F, Fut>(
    scheduler: Arc<Scheduler<P, T>>,
    task_fn: F,
    lane: Lane,
    parent_task_id: Option<crate::Ident>,
    client_id: ClientId,
  ) -> Arc<Self>
  where
    F: FnOnce(TaskContext<P, T>) -> Fut + Send + 'static,
    Fut: Future<Output = Option<RecordWriter<P>>> + Send + 'static,
  {
    Self::build(
      scheduler,
      task_fn,
      lane,
      parent_task_id,
      client_id,
      TaskClass::System,
    )
  }

  fn build<F, Fut>(
    scheduler: Arc<Scheduler<P, T>>,
    task_fn: F,
    lane: Lane,
    parent_task_id: Option<crate::Ident>,
    client_id: ClientId,
    task_class: TaskClass,
  ) -> Arc<Self>
  where
    F: FnOnce(TaskContext<P, T>) -> Fut + Send + 'static,
    Fut: Future<Output = Option<RecordWriter<P>>> + Send + 'static,
  {
    let trace_context = TraceContext::from_current_span();
    let execution_id = generate_execution_id();

    // Count user work the moment the task exists, not at queue time: a
    // `Pending` task is re-queued on wake through `queue_task`, so counting
    // there would double-count. Construction happens exactly once per task.
    if task_class == TaskClass::User {
      scheduler.user_task_started();
    }

    let ctx = TaskContext::new(
      scheduler.db.clone(),
      scheduler.clone(),
      client_id,
    )
    .with_execution_id(execution_id);

    let ctx = if let Some(parent_id) = parent_task_id {
      ctx.with_task_id(parent_id)
    } else {
      ctx
    };

    Arc::new(Self {
      future: Mutex::new(Some(Box::pin(task_fn(ctx)))),
      lane,
      scheduler,
      trace_context,
      task_class,
      birth_epoch: Mutex::new(None),
      _phantom: PhantomData,
    })
  }
}

impl<P: Partitions, T: LanguageServer<P>> futures::task::ArcWake
  for LaburnumTask<P, T>
{
  fn wake_by_ref(arc_self: &Arc<Self>) {
    use crate::scheduler::lanes::is_rpc_lane;
    if is_rpc_lane(arc_self.lane) {
      arc_self.scheduler.queue_rpc_task(arc_self.clone());
    } else {
      arc_self.scheduler.queue_task(arc_self.clone());
    }
  }
}

/// Execution context provided to tasks.
///
/// The task context provides access to:
/// - **Database queries** with snapshot isolation
/// - **Source cache** for reading file content
/// - **Child task spawning** for parallel work
/// - **Matched keys** for watchers (which keys changed)
/// - **Notifications** for sending LSP messages
///
/// # Snapshot Isolation
///
/// Each context captures a database snapshot timestamp. All queries through
/// this context see a consistent view of the database, even if new chunks are
/// added concurrently.
///
/// # Example
///
/// ```ignore
/// scheduler.queue(move |mut ctx| async move {
///   // Query database
///   let records = ctx.query_client()
///     .query(PARTITION_KEY)
///     .execute()
///     .await;
///
///   // Read source file
///   let source = ctx.source_cache_reader()
///     .get_source(source_key)?;
///
///   // Write results
///   let mut writer = ctx.new_record_writer(task_id);
///   writer.insert(partition_key, sort_key, result);
///
///   Ok(Some(writer))
/// }, DEFAULT_LANE);
/// ```
pub struct TaskContext<P: Partitions, T: LanguageServer<P>> {
  db:                   Database<P>,
  scheduler:            Arc<Scheduler<P, T>>,
  connection:           Connection,
  query_clients:        Vec<QueryClient<P>>,
  matched_keys_updated: Vec<RecordKey<P>>,
  matched_keys_deleted: Vec<RecordKey<P>>,
  /// Dependencies this run read while absent. A watcher that needs a record
  /// that is not present yet records `(partition, condition)` here via
  /// [`defer_until`](Self::defer_until); after the handler returns, the
  /// scheduler registers a re-dispatch so the watcher re-runs when a matching
  /// key lands. See GLD-0035.
  pending_deps:         Vec<(Ident, SortKeyCondition<P::SortKey>)>,
  task_id:              Option<Ident>,
  execution_id:         Ident,
  client_id:            ClientId,
  _phantom:             PhantomData<T>,
}

impl<P: Partitions, T: LanguageServer<P>> TaskContext<P, T> {
  pub(crate) fn new(
    db: Database<P>,
    scheduler: Arc<Scheduler<P, T>>,
    client_id: ClientId,
  ) -> Self {
    let connection = scheduler.connection.clone();
    Self {
      db,
      scheduler,
      connection,
      query_clients: Vec::new(),
      matched_keys_updated: Vec::new(),
      matched_keys_deleted: Vec::new(),
      pending_deps: Vec::new(),
      task_id: None,
      execution_id: Ident::new("unset"),
      client_id,
      _phantom: PhantomData,
    }
  }

  pub(crate) fn with_execution_id(mut self, execution_id: Ident) -> Self {
    self.execution_id = execution_id;
    self
  }

  pub(crate) fn with_task_id(mut self, task_id: Ident) -> Self {
    self.task_id = Some(task_id);
    self
  }

  pub fn client_id(&self) -> ClientId {
    self.client_id
  }

  pub fn position_encoding(&self) -> crate::protocol::lsp::PositionEncodingKind {
    self
      .scheduler
      .registry()
      .get(self.client_id)
      .map(|c| c.position_encoding().clone())
      .unwrap_or(crate::protocol::lsp::PositionEncodingKind::DEFAULT)
  }

  pub(crate) fn new_record_writer(&self, task_id: Ident) -> RecordWriter<P> {
    RecordWriter::new(task_id).with_parent_task_id(self.task_id)
  }

  /// Creates a new query client for reading the database.
  ///
  /// Each call creates a separate client that tracks its own dependencies and
  /// captures a database snapshot. All queries through this client see a
  /// consistent view of the database.
  ///
  /// # Snapshot Isolation
  ///
  /// The client captures the current database timestamp. It only sees chunks
  /// with `commit_time <=` snapshot time, ensuring consistent reads even as
  /// new chunks are added.
  ///
  /// # Example
  ///
  /// ```ignore
  /// let all_symbols = ctx.query_client()
  ///   .query(SYMBOL_TABLE_PK)
  ///   .execute()
  ///   .await;
  ///
  /// for (metadata, record_ref) in all_symbols.iter() {
  ///   // Process each record
  /// }
  /// ```
  pub fn query_client(&mut self) -> &mut QueryClient<P> {
    let client = QueryClient::new(self.db.clone());
    self.query_clients.push(client);
    let index = self.query_clients.len() - 1;
    &mut self.query_clients[index]
  }

  /// Spawns a child task on the specified priority lane.
  ///
  /// The spawned task will be executed by the worker pool. Child tasks
  /// automatically have their parent task ID set for dependency tracking.
  ///
  /// # Example
  ///
  /// ```ignore
  /// // Spawn parallel typecheck tasks for multiple modules
  /// for module in modules {
  ///   ctx.spawn_task(move |mut ctx| async move {
  ///     let types = typecheck_module(module, &ctx).await?;
  ///     let mut writer = ctx.new_record_writer(task_id);
  ///     writer.insert(TYPES_PK, sort_key, types);
  ///     Ok(Some(writer))
  ///   }, ASYNC_LANE1);
  /// }
  /// ```
  pub(crate) fn spawn_task<F, Fut>(
    &self,
    task_fn: F,
    lane: Lane,
  ) -> Arc<LaburnumTask<P, T>>
  where
    F: FnOnce(TaskContext<P, T>) -> Fut + Send + 'static,
    Fut: std::future::Future<
        Output = Option<crate::database::chunk::RecordWriter<P>>,
      > + Send
      + 'static,
  {
    self.spawn_task_for_client(task_fn, lane, self.client_id)
  }

  /// Spawns a child task with a specific client_id.
  ///
  /// This is used by RpcTask to spawn tasks with the client_id extracted from
  /// the incoming message, rather than using the RpcTask's own client_id.
  pub(crate) fn spawn_task_for_client<F, Fut>(
    &self,
    task_fn: F,
    lane: Lane,
    client_id: ClientId,
  ) -> Arc<LaburnumTask<P, T>>
  where
    F: FnOnce(TaskContext<P, T>) -> Fut + Send + 'static,
    Fut: std::future::Future<
        Output = Option<crate::database::chunk::RecordWriter<P>>,
      > + Send
      + 'static,
  {
    let parent_task_id = self.task_id;

    let task = LaburnumTask::new_with_parent(
      self.scheduler.clone(),
      task_fn,
      lane,
      parent_task_id,
      client_id,
    );
    self.scheduler.queue_task(task.clone());
    task
  }

  pub fn filesystems(&self) -> Arc<parking_lot::RwLock<Vec<crate::fs::FS>>> {
    self.scheduler.filesystems.clone()
  }

  pub fn source_cache(
    &self,
  ) -> Arc<parking_lot::RwLock<crate::SourceCache<P, T>>> {
    self.scheduler.source_cache.clone()
  }

  pub fn source_cache_reader(
    &self,
  ) -> crate::source::cache::reporter::SourceCacheReader {
    let guard = self.scheduler.source_cache.read();
    guard.reader()
  }

  pub fn scheduler(&self) -> Arc<Scheduler<P, T>> {
    self.scheduler.clone()
  }

  pub fn server(&self) -> Arc<T> {
    self.scheduler.server.clone()
  }

  /// Send a notification to a specific client or broadcast to all clients.
  ///
  /// When `client_id` is `Some`, the notification is sent only to that client.
  /// When `client_id` is `None`, the notification is broadcast to all clients.
  ///
  /// For `ClientId::INTERNAL` (used in stdio/test mode), the notification is sent
  /// via the direct connection since INTERNAL clients are not registered in the
  /// client registry.
  pub async fn send_notification<N: Notification>(
    &self,
    params: N::Params,
    client_id: Option<ClientId>,
  ) {
    let params_value = serde_json::to_value(&params).unwrap_or_default();
    let notification = jsonrpc::Notification::build(N::METHOD)
      .params(params_value)
      .finish();
    let message = jsonrpc::Message::Notification(notification);

    if let Some(id) = client_id {
      // For INTERNAL client_id (stdio/test mode), use the direct connection
      // since INTERNAL is not registered in the client registry.
      if id == ClientId::INTERNAL {
        if let Err(e) = self.connection.sender.send(message).await {
          otel::error!(
            "notification_send_failed",
            format!("Failed to send notification {} to INTERNAL client: {:?}", N::METHOD, e)
          );
        }
      } else {
        // Send to specific registered client via registry
        if let Err(e) = self.scheduler.registry().send_to(id, message).await {
          otel::error!(
            "notification_send_failed",
            format!("Failed to send notification {} to client {:?}: {:?}", N::METHOD, id, e)
          );
        }
      }
    } else {
      // Broadcast to all clients
      if let Err(e) = self.connection.sender.send(message).await {
        otel::error!(
          "notification_send_failed",
          format!("Failed to send notification {}: {:?}", N::METHOD, e)
        );
      }
    }
  }

  /// Broadcast a notification to all clients subscribed to the given topic.
  ///
  /// Unlike [`send_notification`](Self::send_notification), which targets a
  /// specific client or all clients unconditionally, this method uses the
  /// topic-based subscription system: only clients that have called
  /// [`ClientRegistry::subscribe`] for the given topic will receive the
  /// notification.
  ///
  /// Use this for non-LSP-spec notifications that clients opt into, such as
  /// build progress, test results, or server status changes.
  pub async fn broadcast_notification<N: Notification>(
    &self,
    topic: impl crate::connect::lsp::Topic,
    params: N::Params,
  ) {
    let params_value = serde_json::to_value(&params).unwrap_or_default();
    let notification = jsonrpc::Notification::build(N::METHOD)
      .params(params_value)
      .finish();

    self
      .scheduler
      .registry()
      .broadcast(topic, notification)
      .await;
  }

  pub(crate) async fn send_request_fire_and_forget<R: Request>(
    &self,
    params: R::Params,
  ) {
    use std::sync::atomic::{
      AtomicI64,
      Ordering,
    };
    static REQUEST_ID_COUNTER: AtomicI64 = AtomicI64::new(0);

    let params_value = serde_json::to_value(&params).unwrap_or_default();
    let id =
      jsonrpc::Id::Number(REQUEST_ID_COUNTER.fetch_add(1, Ordering::Relaxed));
    let request = jsonrpc::Request::build(R::METHOD, id)
      .params(params_value)
      .finish();
    let message = jsonrpc::Message::Request(request);
    let _ = self.connection.sender.send(message).await;
  }

  pub fn progress_tracker(&self) -> Arc<ProgressTracker> {
    self.scheduler.progress_tracker.clone()
  }

  /// Begin a progress operation (idempotent — a second begin while one is
  /// open coalesces). The matching `end` is sent automatically when the
  /// scheduler goes idle (all user work drained); callers emit cosmetic
  /// updates with [`progress_report`](Self::progress_report).
  pub fn progress(&self, progress_id: Ident, title: &str) {
    self.scheduler.progress_tracker.begin(progress_id, title);
  }

  /// Send a cosmetic `report` update (a text label) for an open progress.
  /// No-op if the progress isn't open.
  pub fn progress_report(&self, progress_id: Ident, message: &str) {
    self
      .scheduler
      .progress_tracker
      .report(progress_id, message, None);
  }

  /// Returns the keys that were updated (created or modified) in the triggering
  /// chunk.
  ///
  /// This is populated by watcher tasks spawned via pattern matching.
  /// For regular tasks, this will be empty.
  pub fn matched_keys_updated(&mut self) -> Vec<RecordKey<P>> {
    std::mem::take(&mut self.matched_keys_updated)
  }

  /// Returns the keys that were deleted in the triggering chunk.
  ///
  /// This is populated by watcher tasks spawned via pattern matching.
  /// For regular tasks, this will be empty.
  pub fn matched_keys_deleted(&mut self) -> Vec<RecordKey<P>> {
    std::mem::take(&mut self.matched_keys_deleted)
  }

  /// Internal method to set the matched keys for watcher tasks.
  ///
  /// This should only be called by the scheduler when spawning watcher tasks.
  pub(crate) fn set_matched_keys(
    &mut self,
    updated: Vec<RecordKey<P>>,
    deleted: Vec<RecordKey<P>>,
  ) {
    self.matched_keys_updated = updated;
    self.matched_keys_deleted = deleted;
  }

  /// Record a dependency this run needed but did not find. After the watcher
  /// handler returns, the scheduler registers a re-dispatch keyed by
  /// `(partition_key, condition)`; when a later commit writes a key matching
  /// `condition` in `partition_key`, this watcher is re-run with the same
  /// matched keys.
  ///
  /// This is the non-blocking alternative to `wait_until`: the task completes
  /// now (so it does not pin scheduler quiescence) and re-fires only when the
  /// missing dependency actually lands. See GLD-0035.
  pub fn defer_until(
    &mut self,
    partition_key: Ident,
    condition: SortKeyCondition<P::SortKey>,
  ) {
    self.pending_deps.push((partition_key, condition));
  }

  /// Drain the dependencies recorded via [`defer_until`](Self::defer_until).
  /// Called by the scheduler after the watcher handler returns.
  pub(crate) fn take_pending_deps(
    &mut self,
  ) -> Vec<(Ident, SortKeyCondition<P::SortKey>)> {
    std::mem::take(&mut self.pending_deps)
  }

  pub fn check_source_available(
    &self,
    source_key: crate::SourceKey,
  ) -> Result<(), crate::LaburnumError> {
    let source_cache_reader = self.source_cache_reader();
    if source_cache_reader.get_source(source_key).is_none() {
      Err(crate::LaburnumError::TaskCancelled)
    } else {
      Ok(())
    }
  }

  /// Returns the number of connected clients.
  pub fn client_count(&self) -> usize {
    self.scheduler.registry().client_count()
  }

  /// Returns the IDs of all connected clients.
  pub fn client_ids(&self) -> Vec<crate::connect::lsp::ClientId> {
    self.scheduler.registry().client_ids()
  }

  /// Returns the duration since the last client disconnected, if the server is
  /// idle.
  pub fn idle_duration(&self) -> Option<std::time::Duration> {
    self.scheduler.registry().idle_duration()
  }

  /// Returns info about all connected clients.
  pub fn connected_clients(&self) -> Vec<crate::connect::lsp::ConnectedClientInfo> {
    self.scheduler.registry().connected_clients()
  }

  /// Request graceful daemon shutdown.
  ///
  /// The shutdown will happen after the current request completes.
  /// This is used by the `lsp/shutdown` command.
  pub fn request_shutdown(&self) {
    self.scheduler.request_shutdown();
  }
}

#[cfg(test)]
mod tests {
  use {
    super::*,
    crate::{
      connect::lsp::ClientId,
      database::tests::storage::{TestPartitions, TestSortKey},
      protocol::lsp::LspVersion,
      server::LaburnumLanguageServer,
    },
  };

  fn test_client_id() -> ClientId {
    ClientId::from_raw(1)
  }

  fn create_test_context() -> TaskContext<TestPartitions, LaburnumLanguageServer> {
    use crate::{
      connect::lsp::ClientId,
      database::Database,
      scheduler::Scheduler,
    };

    let (server_conn, _client_conn) = crate::connect::ipc::Connection::memory();
    let filesystems = std::sync::Arc::new(parking_lot::RwLock::new(Vec::new()));
    let source_cache =
      std::sync::Arc::new(parking_lot::RwLock::new(crate::SourceCache::new()));
    let scheduler = Scheduler::new_with_worker_count(
      server_conn,
      Arc::new(LaburnumLanguageServer),
      filesystems,
      source_cache,
      1,
    );

    TaskContext::new(Database::new(), scheduler, ClientId::INTERNAL)
  }

  #[test]
  fn test_task_context_default_empty_matched_keys() {
    let mut ctx = create_test_context();

    assert_eq!(ctx.matched_keys_updated().len(), 0);
    assert_eq!(ctx.matched_keys_deleted().len(), 0);
  }

  #[test]
  fn test_set_matched_keys() {
    use crate::database::RecordKey;

    let mut ctx = create_test_context();

    let updated = vec![RecordKey::new(Ident::new("pk"), TestSortKey::Test1("sk1".to_string()))];
    let deleted = vec![RecordKey::new(Ident::new("pk"), TestSortKey::Test1("sk2".to_string()))];

    ctx.set_matched_keys(updated.clone(), deleted.clone());

    assert_eq!(ctx.matched_keys_updated(), updated);
    assert_eq!(ctx.matched_keys_deleted(), deleted);
  }

  #[test]
  fn test_check_source_available_returns_error_when_not_found() {
    let ctx = create_test_context();
    let key = crate::SourceKey::new(1, 0);
    let result = ctx.check_source_available(key);
    assert!(matches!(result, Err(crate::LaburnumError::TaskCancelled)));
  }

  #[test]
  fn test_check_source_available_returns_ok_when_found() {
    let ctx = create_test_context();

    let uri = crate::Uri::parse("file:///test.txt").unwrap();
    let source_cache_arc = ctx.source_cache();
    let key = {
      let mut source_cache = source_cache_arc.write();
      let (key, _) = source_cache
        .upsert_with_version(uri, "content".to_string(), LspVersion::new(0), test_client_id())
        .unwrap();
      source_cache.complete_pending_for_test(key);
      key
    };

    let result = ctx.check_source_available(key);
    assert!(result.is_ok());
  }
}