use super::catalog::WorkspaceChunkCatalog;
use super::eligibility::WorkspaceEligibilityPolicy;
use super::reconcile::{CatalogReconcileReport, WorkspaceCatalogReconciler};
use super::types::WorkspaceIndexError;
use crate::workspace::{LocalWorkspaceManifest, WorkspaceFileChange, WorkspaceFileSystem};
use std::sync::{Arc, Mutex};
use std::time::Duration;
use tokio::sync::{broadcast, mpsc};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
const SNAPSHOT_SETTLE_DELAY: Duration = Duration::from_millis(10);
const PERSISTENT_INDEX_SETTLE_DELAY: Duration = Duration::from_millis(50);
const PERSISTENT_INDEX_RETRY_DELAYS: &[Duration] = &[
Duration::from_millis(100),
Duration::from_millis(250),
Duration::from_millis(500),
];
struct PersistentIndexCoordinator {
index: Arc<super::persistent::WorkspacePersistentIndex>,
updates: mpsc::UnboundedSender<Arc<super::catalog::ChunkCatalogSnapshot>>,
task: Mutex<Option<JoinHandle<()>>>,
}
impl PersistentIndexCoordinator {
fn start(
persistent: Arc<super::persistent::WorkspacePersistentIndex>,
lifetime: CancellationToken,
) -> Arc<Self> {
let (updates, mut pending_updates) =
mpsc::unbounded_channel::<Arc<super::catalog::ChunkCatalogSnapshot>>();
let task_index = Arc::clone(&persistent);
let task = tokio::spawn(async move {
loop {
let mut pending = tokio::select! {
_ = lifetime.cancelled() => return,
next = pending_updates.recv() => match next {
Some(snapshot) => snapshot,
None => return,
},
};
let settle = tokio::time::sleep(PERSISTENT_INDEX_SETTLE_DELAY);
tokio::pin!(settle);
loop {
tokio::select! {
_ = lifetime.cancelled() => return,
_ = &mut settle => break,
next = pending_updates.recv() => match next {
Some(update) => {
if is_newer_snapshot(&update, &pending) {
pending = update;
}
}
None => break,
},
}
}
sync_snapshot_with_retry(
Arc::clone(&task_index),
pending,
&lifetime,
&mut pending_updates,
)
.await;
}
});
Arc::new(Self {
index: persistent,
updates,
task: Mutex::new(Some(task)),
})
}
fn submit(&self, snapshot: super::catalog::ChunkCatalogSnapshot) {
let _ = self.updates.send(Arc::new(snapshot));
}
async fn publish(&self, snapshot: super::catalog::ChunkCatalogSnapshot) {
if self.index.is_ready() {
self.submit(snapshot);
return;
}
let index = Arc::clone(&self.index);
let pending = Arc::new(snapshot.clone());
let (tx, rx) = tokio::sync::oneshot::channel();
if std::thread::Builder::new()
.name("a3s-persistent-publish".to_owned())
.spawn(move || {
let _ = tx.send(index.sync_snapshot(pending.as_ref()));
})
.is_err()
{
tracing::warn!("failed to spawn inline persistent publish worker");
self.submit(snapshot);
return;
}
match rx.await {
Ok(Ok(())) => {
self.submit(snapshot);
}
Ok(Err(error)) => {
tracing::warn!(%error, "inline persistent publish failed; queueing for retry");
self.submit(snapshot);
}
Err(_) => {
tracing::warn!("inline persistent publish worker dropped its result");
self.submit(snapshot);
}
}
}
fn shutdown(&self) {
if let Some(task) = self
.task
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take()
{
task.abort();
}
}
}
async fn sync_snapshot_with_retry(
persistent: Arc<super::persistent::WorkspacePersistentIndex>,
mut pending: Arc<super::catalog::ChunkCatalogSnapshot>,
lifetime: &CancellationToken,
pending_updates: &mut mpsc::UnboundedReceiver<Arc<super::catalog::ChunkCatalogSnapshot>>,
) {
let mut retries = 0usize;
loop {
let snapshot = Arc::clone(&pending);
let persistent = Arc::clone(&persistent);
let (tx, rx) = tokio::sync::oneshot::channel();
let spawn_result = std::thread::Builder::new()
.name("a3s-persistent-sync".to_owned())
.spawn(move || {
let _ = tx.send(persistent.sync_snapshot(snapshot.as_ref()));
});
let result = match spawn_result {
Ok(_join) => rx.await.map_err(|_| ()),
Err(_) => Err(()),
};
match result {
Ok(Ok(())) => return,
Ok(Err(error)) => {
let retryable = retryable_index_error(&error);
tracing::warn!(%error, retryable, retries, "workspace persistent index update failed");
if !retryable {
let Some(newer_snapshot) = wait_for_index_retry(
Duration::from_secs(1),
lifetime,
pending_updates,
&mut pending,
)
.await
else {
return;
};
if newer_snapshot {
retries = 0;
}
continue;
}
let delay = PERSISTENT_INDEX_RETRY_DELAYS
.get(retries.min(PERSISTENT_INDEX_RETRY_DELAYS.len().saturating_sub(1)))
.copied()
.unwrap_or(Duration::from_secs(1));
retries = retries.saturating_add(1);
let Some(newer_snapshot) =
wait_for_index_retry(delay, lifetime, pending_updates, &mut pending).await
else {
return;
};
if newer_snapshot {
retries = 0;
}
}
Err(()) => {
tracing::warn!(retries, "workspace persistent index sync worker failed");
retries = retries.saturating_add(1);
let Some(newer_snapshot) = wait_for_index_retry(
Duration::from_millis(250),
lifetime,
pending_updates,
&mut pending,
)
.await
else {
return;
};
if newer_snapshot {
retries = 0;
}
}
}
}
}
async fn wait_for_index_retry(
delay: Duration,
lifetime: &CancellationToken,
pending_updates: &mut mpsc::UnboundedReceiver<Arc<super::catalog::ChunkCatalogSnapshot>>,
pending: &mut Arc<super::catalog::ChunkCatalogSnapshot>,
) -> Option<bool> {
let retry = tokio::time::sleep(delay);
tokio::pin!(retry);
tokio::select! {
_ = lifetime.cancelled() => None,
_ = &mut retry => {
Some(drain_newer_snapshots(pending_updates, pending))
}
next = pending_updates.recv() => match next {
Some(update) => {
let mut newer_snapshot = false;
if is_newer_snapshot(&update, pending) {
*pending = update;
newer_snapshot = true;
}
if drain_newer_snapshots(pending_updates, pending) {
newer_snapshot = true;
}
Some(newer_snapshot)
}
None => None,
}
}
}
fn drain_newer_snapshots(
pending_updates: &mut mpsc::UnboundedReceiver<Arc<super::catalog::ChunkCatalogSnapshot>>,
pending: &mut Arc<super::catalog::ChunkCatalogSnapshot>,
) -> bool {
let mut newer_snapshot = false;
while let Ok(update) = pending_updates.try_recv() {
if is_newer_snapshot(&update, pending) {
*pending = update;
newer_snapshot = true;
}
}
newer_snapshot
}
fn retryable_index_error(error: &WorkspaceIndexError) -> bool {
match error {
WorkspaceIndexError::InvalidConfig(message) => {
message.starts_with("persistent zvec index failed:")
}
WorkspaceIndexError::InvalidQuery(_) | WorkspaceIndexError::StaleRevision { .. } => false,
_ => true,
}
}
impl Drop for PersistentIndexCoordinator {
fn drop(&mut self) {
self.shutdown();
}
}
fn is_newer_snapshot(
candidate: &super::catalog::ChunkCatalogSnapshot,
current: &super::catalog::ChunkCatalogSnapshot,
) -> bool {
(candidate.source_revision(), candidate.revision())
> (current.source_revision(), current.revision())
}
#[cfg(all(test, feature = "zvec-rust-fts"))]
mod tests {
use super::PersistentIndexCoordinator;
use crate::workspace::{
ChunkCatalogLimits, ChunkingConfig, WorkspaceChunkCatalog, WorkspaceIndexError,
WorkspaceLexicalEngine, WorkspacePath, WorkspacePersistentIndex,
};
use std::time::Duration;
use tokio_util::sync::CancellationToken;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn persistent_coordinator_survives_initial_none_before_first_submit() {
let temp = tempfile::tempdir().expect("temporary workspace");
let catalog = WorkspaceChunkCatalog::new_with_engine(
ChunkingConfig::default(),
ChunkCatalogLimits::default(),
WorkspaceLexicalEngine::ZvecRust,
)
.expect("catalog");
let path = WorkspacePath::from_normalized("src/late.rs");
catalog
.replace_file(&path, Some("rust"), 1, "pub fn late_marker() {}\n")
.expect("catalog replacement");
let index = WorkspacePersistentIndex::open(
temp.path().join(".a3s-code/index"),
WorkspaceLexicalEngine::ZvecRust,
)
.expect("persistent index");
let lifetime = CancellationToken::new();
let coordinator = PersistentIndexCoordinator::start(index.clone(), lifetime.clone());
tokio::time::sleep(Duration::from_millis(50)).await;
coordinator.submit(catalog.snapshot().expect("catalog snapshot"));
tokio::time::timeout(Duration::from_secs(15), async {
while !index.is_ready() {
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.expect("persistent coordinator exited after the initial None watch notification");
assert_eq!(index.status().source_revision, 1);
lifetime.cancel();
coordinator.shutdown();
}
#[tokio::test]
async fn persistent_coordinator_coalesces_a_save_burst_to_the_newest_snapshot() {
let temp = tempfile::tempdir().expect("temporary workspace");
let catalog = WorkspaceChunkCatalog::new_with_engine(
ChunkingConfig::default(),
ChunkCatalogLimits::default(),
WorkspaceLexicalEngine::ZvecRust,
)
.expect("catalog");
let path = WorkspacePath::from_normalized("src/burst.rs");
let index = WorkspacePersistentIndex::open(
temp.path().join(".a3s-code/index"),
WorkspaceLexicalEngine::ZvecRust,
)
.expect("persistent index");
let lifetime = CancellationToken::new();
let coordinator = PersistentIndexCoordinator::start(index.clone(), lifetime.clone());
for revision in 1..=4 {
catalog
.replace_file(
&path,
Some("rust"),
revision,
&format!("pub fn burst_marker_{revision}() {{}}\n"),
)
.expect("catalog replacement");
coordinator.submit(catalog.snapshot().expect("catalog snapshot"));
}
let latest_revision = catalog
.snapshot()
.expect("latest snapshot")
.source_revision();
tokio::time::timeout(Duration::from_secs(15), async {
loop {
if index.status().source_revision == latest_revision {
break;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.expect("coalesced persistent index did not catch up");
let generations = std::fs::read_dir(temp.path().join(".a3s-code/index"))
.expect("persistent index directory")
.filter_map(Result::ok)
.filter(|entry| {
entry
.file_name()
.to_str()
.is_some_and(|name| name.starts_with("generation-"))
})
.count();
assert_eq!(generations, 1, "save burst built intermediate generations");
lifetime.cancel();
coordinator.shutdown();
}
#[tokio::test]
async fn persistent_coordinator_retries_a_transient_publish_failure() {
let temp = tempfile::tempdir().expect("temporary workspace");
let catalog = WorkspaceChunkCatalog::new_with_engine(
ChunkingConfig::default(),
ChunkCatalogLimits::default(),
WorkspaceLexicalEngine::ZvecRust,
)
.expect("catalog");
let path = WorkspacePath::from_normalized("src/retry.rs");
catalog
.replace_file(&path, Some("rust"), 1, "pub fn retry_marker() {}\n")
.expect("catalog replacement");
let index = WorkspacePersistentIndex::open(
temp.path().join(".a3s-code/index"),
WorkspaceLexicalEngine::ZvecRust,
)
.expect("persistent index");
let destination = temp.path().join(".a3s-code/index/generation-1");
std::fs::create_dir_all(destination.parent().expect("index parent")).expect("index parent");
std::fs::write(&destination, "temporary publish blocker").expect("publish blocker");
let lifetime = CancellationToken::new();
let coordinator = PersistentIndexCoordinator::start(index.clone(), lifetime.clone());
let unblock = destination.clone();
tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(80)).await;
std::fs::remove_file(unblock).expect("remove publish blocker");
});
coordinator.submit(catalog.snapshot().expect("catalog snapshot"));
tokio::time::timeout(Duration::from_secs(15), async {
loop {
if index.status().source_revision == 1
&& index.status().phase
== crate::workspace::WorkspacePersistentIndexPhase::Ready
{
break;
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.expect("coordinator did not recover after transient publish failure");
lifetime.cancel();
coordinator.shutdown();
}
#[test]
fn retries_only_recoverable_persistent_index_errors() {
assert!(!super::retryable_index_error(
&WorkspaceIndexError::InvalidConfig("schema".to_owned())
));
assert!(super::retryable_index_error(
&WorkspaceIndexError::InvalidConfig(
"persistent zvec index failed: temporary native lock".to_owned()
)
));
assert!(!super::retryable_index_error(
&WorkspaceIndexError::InvalidQuery("query".to_owned())
));
assert!(!super::retryable_index_error(
&WorkspaceIndexError::StaleRevision {
requested: 1,
current: 2,
}
));
assert!(super::retryable_index_error(
&WorkspaceIndexError::ReadFailed {
path: "index".to_owned(),
message: "temporarily unavailable".to_owned(),
}
));
}
}
pub(crate) struct LocalWorkspaceCatalogRuntime {
catalog: Arc<WorkspaceChunkCatalog>,
lifetime: CancellationToken,
task: Mutex<Option<tokio::task::JoinHandle<()>>>,
persistent: Option<Arc<PersistentIndexCoordinator>>,
}
impl LocalWorkspaceCatalogRuntime {
pub(crate) fn start(
manifest: Arc<LocalWorkspaceManifest>,
file_system: Arc<dyn WorkspaceFileSystem>,
) -> Arc<Self> {
Self::start_with_catalog_and_persistent(
manifest,
file_system,
WorkspaceChunkCatalog::default_catalog(),
None,
)
}
pub(crate) fn start_with_catalog_and_persistent(
manifest: Arc<LocalWorkspaceManifest>,
file_system: Arc<dyn WorkspaceFileSystem>,
catalog: Arc<WorkspaceChunkCatalog>,
persistent: Option<Arc<super::persistent::WorkspacePersistentIndex>>,
) -> Arc<Self> {
let snapshots = manifest.subscribe();
let changes = manifest.subscribe_changes();
let lifetime = CancellationToken::new();
let persistent_coordinator = persistent
.map(|persistent| PersistentIndexCoordinator::start(persistent, lifetime.clone()));
let runtime = Arc::new(Self {
catalog: Arc::clone(&catalog),
lifetime: lifetime.clone(),
task: Mutex::new(None),
persistent: persistent_coordinator.clone(),
});
let task = tokio::spawn(run_catalog_updates(
manifest,
WorkspaceCatalogReconciler::new(
catalog,
WorkspaceEligibilityPolicy::default(),
file_system,
),
snapshots,
changes,
lifetime,
persistent_coordinator,
));
*runtime
.task
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner) = Some(task);
runtime
}
pub(crate) fn catalog(&self) -> Arc<WorkspaceChunkCatalog> {
Arc::clone(&self.catalog)
}
pub(crate) fn shutdown(&self) {
self.lifetime.cancel();
if let Some(task) = self
.task
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take()
{
task.abort();
}
if let Some(persistent) = &self.persistent {
persistent.shutdown();
}
}
}
impl Drop for LocalWorkspaceCatalogRuntime {
fn drop(&mut self) {
self.shutdown();
}
}
async fn run_catalog_updates(
manifest: Arc<LocalWorkspaceManifest>,
reconciler: WorkspaceCatalogReconciler,
mut snapshots: broadcast::Receiver<crate::workspace::LocalWorkspaceManifestSnapshot>,
mut changes: broadcast::Receiver<WorkspaceFileChange>,
lifetime: CancellationToken,
persistent: Option<Arc<PersistentIndexCoordinator>>,
) {
let mut seeded = false;
while !seeded {
let initial = manifest.snapshot();
if initial.version > 0 {
report_reconciliation(
reconciler.reconcile_snapshot(&initial).await,
&reconciler,
persistent.as_ref(),
)
.await;
seeded = true;
break;
}
tokio::select! {
_ = lifetime.cancelled() => return,
update = snapshots.recv() => match update {
Ok(snapshot) => {
if snapshot.version > 0 {
report_reconciliation(
reconciler.reconcile_snapshot(&snapshot).await,
&reconciler,
persistent.as_ref(),
)
.await;
seeded = true;
}
}
Err(broadcast::error::RecvError::Lagged(_)) => {
let snapshot = manifest.snapshot();
if snapshot.version > 0 {
report_reconciliation(
reconciler.reconcile_after_lag(&snapshot).await,
&reconciler,
persistent.as_ref(),
)
.await;
seeded = true;
}
}
Err(broadcast::error::RecvError::Closed) => return,
},
_ = tokio::time::sleep(Duration::from_millis(20)) => {}
}
}
let _ = seeded;
loop {
tokio::select! {
_ = lifetime.cancelled() => break,
update = snapshots.recv() => match update {
Ok(snapshot) => {
tokio::time::sleep(SNAPSHOT_SETTLE_DELAY).await;
let batch = drain_changes(&mut changes);
if batch.lagged {
report_reconciliation(reconciler.reconcile_after_lag(&manifest.snapshot()).await, &reconciler, persistent.as_ref()).await;
} else if batch.changes.is_empty() {
report_reconciliation(reconciler.reconcile_snapshot(&snapshot).await, &reconciler, persistent.as_ref()).await;
} else {
report_reconciliation(reconciler.reconcile_changes(&snapshot, &batch.changes).await, &reconciler, persistent.as_ref()).await;
}
}
Err(broadcast::error::RecvError::Lagged(skipped)) => {
tracing::warn!(skipped, "workspace retrieval snapshot stream lagged; rebuilding admitted files");
report_reconciliation(reconciler.reconcile_after_lag(&manifest.snapshot()).await, &reconciler, persistent.as_ref()).await;
}
Err(broadcast::error::RecvError::Closed) => break,
},
update = changes.recv() => match update {
Ok(change) => {
let mut batch = drain_changes(&mut changes);
batch.changes.insert(0, change);
if batch.lagged {
report_reconciliation(reconciler.reconcile_after_lag(&manifest.snapshot()).await, &reconciler, persistent.as_ref()).await;
} else {
report_reconciliation(
reconciler
.reconcile_changes(&manifest.snapshot(), &batch.changes)
.await,
&reconciler,
persistent.as_ref(),
).await;
}
}
Err(broadcast::error::RecvError::Lagged(skipped)) => {
tracing::warn!(skipped, "workspace retrieval change stream lagged; rebuilding admitted files");
report_reconciliation(reconciler.reconcile_after_lag(&manifest.snapshot()).await, &reconciler, persistent.as_ref()).await;
}
Err(broadcast::error::RecvError::Closed) => break,
},
}
}
}
fn drain_changes(changes: &mut broadcast::Receiver<WorkspaceFileChange>) -> DrainedChanges {
let mut batch = DrainedChanges::default();
loop {
match changes.try_recv() {
Ok(change) => batch.changes.push(change),
Err(broadcast::error::TryRecvError::Lagged(_)) => batch.lagged = true,
Err(broadcast::error::TryRecvError::Empty | broadcast::error::TryRecvError::Closed) => {
break;
}
}
}
batch
}
#[derive(Default)]
struct DrainedChanges {
changes: Vec<WorkspaceFileChange>,
lagged: bool,
}
async fn report_reconciliation(
result: Result<CatalogReconcileReport, super::types::WorkspaceIndexError>,
reconciler: &WorkspaceCatalogReconciler,
persistent: Option<&Arc<PersistentIndexCoordinator>>,
) {
match result {
Ok(report) => {
if let Some(persistent) = persistent {
match reconciler.catalog_snapshot() {
Ok(snapshot) => {
persistent.publish(snapshot).await;
}
Err(error) => {
tracing::warn!(%error, "workspace persistent index snapshot failed")
}
}
}
if !report.failures.is_empty() {
tracing::warn!(
source_revision = report.source_revision,
failed_files = report.failures.len(),
indexed_files = report.indexed_files,
"workspace retrieval catalog is partially indexed"
);
}
tracing::debug!(
source_revision = report.source_revision,
catalog_revision = report.catalog_revision,
indexed_files = report.indexed_files,
indexed_chunks = report.indexed_chunks,
eligible_files = report.eligible_files,
read_files = report.read_paths.len(),
removed_files = report.removed_paths.len(),
full_rebuild = report.full_rebuild,
"workspace retrieval catalog reconciled"
);
}
Err(error) => tracing::warn!(%error, "workspace retrieval catalog reconciliation failed"),
}
}