laburnum 1.17.1

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

use {
  crate::{
    Partitions,
    database::chunk::RecordWriter,
    protocol::lsp::LanguageServer,
    scheduler::task::TaskContext,
  },
  std::{
    future::Future,
    pin::Pin,
    sync::{
      Arc,
      atomic::{
        AtomicBool,
        Ordering,
      },
    },
    time::Duration,
  },
};

type IdleMonitorTaskFn<P, T> = Box<
  dyn FnOnce(
      TaskContext<P, T>,
    ) -> Pin<Box<dyn Future<Output = Option<RecordWriter<P>>> + Send>>
    + Send
    + 'static,
>;

/// Polling interval for the idle monitor.
///
/// Chosen to be small enough that idle shutdown feels responsive but large
/// enough to keep the worker free for real work.
const POLL_INTERVAL: Duration = Duration::from_secs(1);

/// Creates a task that watches for daemon idleness and signals shutdown.
///
/// The monitor wakes once per `POLL_INTERVAL` and asks the registry whether
/// any clients are connected. If the registry has been client-free for at
/// least `idle_timeout`, it sets `idle_triggered` and calls
/// `Scheduler::request_shutdown` so the daemon's accept loop observes the
/// shutdown on its next pass and reports `ShutdownReason::IdleTimeout`.
pub(crate) fn idle_monitor_task<P: Partitions, T: LanguageServer<P>>(
  shutdown_flag: Arc<AtomicBool>,
  idle_triggered: Arc<AtomicBool>,
  idle_timeout: Duration,
) -> IdleMonitorTaskFn<P, T> {
  Box::new(move |ctx| {
    Box::pin({
      let scheduler = ctx.scheduler();

      async move {
        loop {
          smol::Timer::after(POLL_INTERVAL).await;

          if shutdown_flag.load(Ordering::Acquire) {
            return None;
          }

          let registry = scheduler.registry();
          if registry.client_count() != 0 {
            continue;
          }

          let Some(idle) = registry.idle_duration() else {
            continue;
          };

          if idle < idle_timeout {
            continue;
          }

          otel::event!(
            "daemon_idle_monitor_trigger",
            "idle_seconds" = idle.as_secs() as i64,
            "timeout_seconds" = idle_timeout.as_secs() as i64
          );
          idle_triggered.store(true, Ordering::Release);
          scheduler.request_shutdown();
          return None;
        }
      }
    })
  })
}