use std::io;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::{Arc, Mutex, mpsc};
use std::thread::JoinHandle;
use std::time::{Duration, Instant};
use tau_cli_term::{ArgCompleter, CompletionItem, TermHandle};
use tau_harness::runtime_dir::RunningSessionSnapshot;
use tau_proto::SessionId;
const FRESH_TTL: Duration = Duration::from_secs(5);
const FAILURE_BACKOFF: Duration = Duration::from_secs(2);
const STALE_TTL: Duration = Duration::from_secs(30);
type Clock = Arc<dyn Fn() -> Instant + Send + Sync>;
type Discovery = Arc<dyn Fn() -> io::Result<RunningSessionSnapshot> + Send + Sync>;
#[cfg(test)]
type WorkerTestHook = Arc<dyn Fn() + Send + Sync>;
#[cfg(test)]
#[derive(Clone, Default)]
struct WorkerTestHooks {
after_shutdown_check: Option<WorkerTestHook>,
after_shutdown_action: Option<WorkerTestHook>,
}
pub(super) struct SessionCompletion {
control: Arc<WorkerControl>,
}
impl SessionCompletion {
pub(super) fn new(current_session: SessionId, term: TermHandle) -> Self {
Self::new_with(
current_session,
term,
Arc::new(Instant::now),
Arc::new(tau_harness::runtime_dir::list_running_sessions_tolerant),
)
}
pub(super) fn completer(&self) -> ArgCompleter {
let control = Arc::downgrade(&self.control);
Arc::new(move |args: &[&str]| {
let [partial] = args else {
return Vec::new();
};
let Some(control) = control.upgrade() else {
return Vec::new();
};
control.complete(partial)
})
}
fn new_with(
current_session: SessionId,
term: TermHandle,
clock: Clock,
discovery: Discovery,
) -> Self {
#[cfg(test)]
{
Self::new_with_test_hooks(
current_session,
term,
clock,
discovery,
WorkerTestHooks::default(),
)
}
#[cfg(not(test))]
{
Self {
control: Arc::new(WorkerControl {
current_session,
term,
clock,
discovery,
cache: Arc::new(Mutex::new(Cache::default())),
worker: Mutex::new(None),
latest_generation: Arc::new(AtomicU64::new(0)),
lifecycle: Arc::new(WorkerLifecycle {
shutdown: AtomicBool::new(false),
}),
}),
}
}
}
#[cfg(test)]
fn new_with_test_hooks(
current_session: SessionId,
term: TermHandle,
clock: Clock,
discovery: Discovery,
test_hooks: WorkerTestHooks,
) -> Self {
Self {
control: Arc::new(WorkerControl {
current_session,
term,
clock,
discovery,
cache: Arc::new(Mutex::new(Cache::default())),
worker: Mutex::new(None),
latest_generation: Arc::new(AtomicU64::new(0)),
lifecycle: Arc::new(WorkerLifecycle {
shutdown: AtomicBool::new(false),
test_hooks,
}),
}),
}
}
}
impl Drop for SessionCompletion {
fn drop(&mut self) {
self.control.shutdown();
}
}
struct WorkerControl {
current_session: SessionId,
term: TermHandle,
clock: Clock,
discovery: Discovery,
cache: Arc<Mutex<Cache>>,
worker: Mutex<Option<Worker>>,
latest_generation: Arc<AtomicU64>,
lifecycle: Arc<WorkerLifecycle>,
}
struct WorkerLifecycle {
shutdown: AtomicBool,
#[cfg(test)]
test_hooks: WorkerTestHooks,
}
impl WorkerControl {
fn complete(&self, partial: &str) -> Vec<CompletionItem> {
let generation = self.term.completion_refresh_generation();
self.latest_generation.store(generation, Ordering::Release);
let now = (self.clock)();
let (items, refresh) = self
.cache
.lock()
.expect("session completion cache")
.read(now);
if refresh {
self.request_discovery();
}
let needle = partial.to_lowercase();
let mut prefix_matches = Vec::new();
let mut substring_matches = Vec::new();
for indexed in items.iter() {
if needle.is_empty() || indexed.lower_value.starts_with(&needle) {
prefix_matches.push(indexed.item.clone());
} else if indexed.lower_value.contains(&needle) {
substring_matches.push(indexed.item.clone());
}
}
prefix_matches.extend(substring_matches);
prefix_matches
}
fn request_discovery(&self) {
let mut worker = self.worker.lock().expect("session completion worker");
if worker.is_none() {
*worker = Some(Worker::spawn(
self.current_session.clone(),
self.term.clone(),
Arc::clone(&self.clock),
Arc::clone(&self.discovery),
Arc::clone(&self.cache),
Arc::clone(&self.latest_generation),
Arc::clone(&self.lifecycle),
));
}
if let Some(worker) = worker.as_ref() {
let _ = worker.demand.try_send(());
}
}
fn shutdown(&self) {
self.lifecycle.shutdown.store(true, Ordering::Release);
let worker = self
.worker
.lock()
.expect("session completion worker")
.take();
if let Some(worker) = worker {
let Worker { demand, join } = worker;
let _ = demand.try_send(());
#[cfg(test)]
if let Some(hook) = &self.lifecycle.test_hooks.after_shutdown_action {
hook();
}
drop(demand);
let _ = join.join();
}
}
}
struct Worker {
demand: mpsc::SyncSender<()>,
join: JoinHandle<()>,
}
impl Worker {
fn spawn(
current_session: SessionId,
term: TermHandle,
clock: Clock,
discovery: Discovery,
cache: Arc<Mutex<Cache>>,
latest_generation: Arc<AtomicU64>,
lifecycle: Arc<WorkerLifecycle>,
) -> Self {
let (demand, demand_rx) = mpsc::sync_channel(1);
let join = std::thread::spawn(move || {
while demand_rx.recv().is_ok() {
if lifecycle.shutdown.load(Ordering::Acquire) {
break;
}
let result = discovery();
let observed_at = clock();
{
let mut cache = cache.lock().expect("session completion cache");
match result {
Ok(snapshot) => {
if 0 < snapshot.incomplete_claims {
tracing::debug!(
incomplete_claims = snapshot.incomplete_claims,
"session completion discovery omitted incomplete claims"
);
}
cache.record_success(
observed_at,
completion_items(snapshot, ¤t_session),
);
}
Err(error) => {
cache.record_failure(observed_at);
tracing::debug!(
%error,
"session completion discovery failed"
);
}
}
}
if lifecycle.shutdown.load(Ordering::Acquire) {
break;
}
#[cfg(test)]
if let Some(hook) = &lifecycle.test_hooks.after_shutdown_check {
hook();
}
while demand_rx.try_recv().is_ok() {}
let generation = latest_generation.load(Ordering::Acquire);
term.request_completion_refresh_if_generation(generation);
}
});
Self { demand, join }
}
}
#[derive(Default)]
struct Cache {
items: Arc<Vec<IndexedCompletion>>,
last_success: Option<Instant>,
last_failure: Option<Instant>,
}
impl Cache {
fn read(&self, now: Instant) -> (Arc<Vec<IndexedCompletion>>, bool) {
let fresh = self
.last_success
.is_some_and(|at| now.saturating_duration_since(at) <= FRESH_TTL);
let backed_off = self
.last_failure
.is_some_and(|at| now.saturating_duration_since(at) < FAILURE_BACKOFF);
let usable = self
.last_success
.is_some_and(|at| now.saturating_duration_since(at) <= STALE_TTL);
(
if usable {
Arc::clone(&self.items)
} else {
Arc::new(Vec::new())
},
!fresh && !backed_off,
)
}
fn record_success(&mut self, observed_at: Instant, items: Vec<IndexedCompletion>) {
self.items = Arc::new(items);
self.last_success = Some(observed_at);
self.last_failure = None;
}
fn record_failure(&mut self, observed_at: Instant) {
self.last_failure = Some(observed_at);
}
}
struct IndexedCompletion {
item: CompletionItem,
lower_value: String,
}
fn completion_items(
snapshot: RunningSessionSnapshot,
current_session: &SessionId,
) -> Vec<IndexedCompletion> {
snapshot
.sessions
.into_iter()
.map(|session| {
let mut description = session.project_root.display().to_string();
if &session.session_id == current_session {
description.push_str(" (current)");
}
let value = session.session_id.to_string();
IndexedCompletion {
lower_value: value.to_lowercase(),
item: CompletionItem::new(value, description),
}
})
.collect()
}
#[cfg(test)]
mod tests;