diff --git a/crates/tinymemory-tools/src/deletion/mod.rs b/crates/tinymemory-tools/src/deletion/mod.rs index 6130ee75..ac7f67be 100644 --- a/crates/tinymemory-tools/src/deletion/mod.rs +++ b/crates/tinymemory-tools/src/deletion/mod.rs @@ -88,7 +88,7 @@ impl Backoff { pub trait DeletionExecutor: Send + Sync { /// Whether the engine can be reached at all. When not, a drain stops /// before running anything. - fn is_reachable(&self) -> bool; + async fn is_reachable(&self) -> bool; /// Runs `deletion`, returning how many items went. /// @@ -252,7 +252,7 @@ impl> DeletionQueue { if queued.is_empty() { return 0; } - if !executor.is_reachable() { + if !executor.is_reachable().await { log::debug!( "[tinymemory:deletion] memory off; {} pending deletions wait", queued.len() diff --git a/crates/tinymemory-tools/src/deletion/mod_tests.rs b/crates/tinymemory-tools/src/deletion/mod_tests.rs index 946db09b..7f91caef 100644 --- a/crates/tinymemory-tools/src/deletion/mod_tests.rs +++ b/crates/tinymemory-tools/src/deletion/mod_tests.rs @@ -1,8 +1,12 @@ use super::*; +use std::future::Future; use std::sync::Mutex as StdMutex; +use std::task::Poll; use crate::list_store::MemoryListStore; +use futures::channel::oneshot; +use futures::future::poll_fn; fn thread(id: &str) -> PendingDeletion { PendingDeletion::Thread { @@ -22,7 +26,7 @@ struct Engine { #[async_trait] impl DeletionExecutor for Engine { - fn is_reachable(&self) -> bool { + async fn is_reachable(&self) -> bool { self.reachable } async fn run(&self, deletion: &PendingDeletion) -> Result { @@ -35,6 +39,27 @@ impl DeletionExecutor for Engine { } } +struct DeferredReachability { + release: StdMutex>>, + ran: StdMutex>, +} + +#[async_trait] +impl DeletionExecutor for DeferredReachability { + async fn is_reachable(&self) -> bool { + let release = self.release.lock().unwrap().take(); + if let Some(release) = release { + let _ = release.await; + } + true + } + + async fn run(&self, deletion: &PendingDeletion) -> Result { + self.ran.lock().unwrap().push(deletion.clone()); + Ok(1) + } +} + fn engine(reachable: bool, failing: bool) -> Engine { Engine { reachable, @@ -64,6 +89,29 @@ async fn nothing_runs_while_the_engine_is_unreachable() { assert_eq!(queue.pending(), vec![thread("t")]); } +#[tokio::test] +async fn drain_awaits_async_reachability_before_running_deletions() { + let queue = queue(); + queue.enqueue(thread("t")); + let (release, wait) = oneshot::channel(); + let engine = DeferredReachability { + release: StdMutex::new(Some(wait)), + ran: StdMutex::new(Vec::new()), + }; + let mut drain = Box::pin(queue.drain(&engine, Drain::Force, Utc::now())); + + poll_fn(|cx| match drain.as_mut().poll(cx) { + Poll::Pending => Poll::Ready(()), + Poll::Ready(_) => panic!("drain completed before async reachability resolved"), + }) + .await; + assert!(engine.ran.lock().unwrap().is_empty()); + + release.send(()).unwrap(); + assert_eq!(drain.await, 1); + assert_eq!(engine.ran.lock().unwrap().as_slice(), &[thread("t")]); +} + #[tokio::test] async fn a_reachable_engine_settles_the_queue() { let queue = queue(); diff --git a/crates/tinymemory-tools/src/layout_migration/job.rs b/crates/tinymemory-tools/src/layout_migration/job.rs index 97fc8d95..3ac1b5fb 100644 --- a/crates/tinymemory-tools/src/layout_migration/job.rs +++ b/crates/tinymemory-tools/src/layout_migration/job.rs @@ -94,7 +94,7 @@ pub trait LayoutHost: Send + Sync { /// # Errors /// /// When memory is off or the account has no per-user root. - fn engines(&self) -> Result; + async fn engines(&self) -> Result; /// Where items go in the per-user tree. /// @@ -128,7 +128,7 @@ pub trait LayoutHost: Send + Sync { /// # Errors /// /// When memory is off or the account has no per-user root. - fn legacy_claim(&self) -> Result>>; + async fn legacy_claim(&self) -> Result>>; /// Whether the last import of old local memory ended with its last batch /// not confirmed listed. The migration then forgets by id, never erasing @@ -167,7 +167,7 @@ where if state.phase == Phase::Cleaned && (state.rechecked || !host.listed_unconfirmed()) { return Ok(Outcome::Done); } - let engines = host.engines()?; + let engines = host.engines().await?; if state.phase == Phase::Idle && !state.switched && !legacy_present(&*engines.legacy).await? { return nothing_to_move(host, &mut state).await; } @@ -175,7 +175,7 @@ where state.takeover = true; store.save(&state)?; } - let shared = host.legacy_claim()?; + let shared = host.legacy_claim().await?; if let Some(claim) = &shared { if claim.held_by_other()? { return claimed_elsewhere(host, &mut state).await; diff --git a/crates/tinymemory-tools/src/layout_migration/service.rs b/crates/tinymemory-tools/src/layout_migration/service.rs index 85fb124d..0aabedd9 100644 --- a/crates/tinymemory-tools/src/layout_migration/service.rs +++ b/crates/tinymemory-tools/src/layout_migration/service.rs @@ -197,7 +197,7 @@ pub struct MigrationStatus { /// tree cannot be read. pub async fn scan(host: &dyn LayoutHost) -> Result { let state = host.store().load()?; - let claim = match host.legacy_claim() { + let claim = match host.legacy_claim().await { // Memory off or signed out: nothing to show. Err(Error::Off(_)) => None, claim => claim?, @@ -215,7 +215,7 @@ pub async fn scan(host: &dyn LayoutHost) -> Result { { false } - Phase::Idle if !state.switched => match host.engines() { + Phase::Idle if !state.switched => match host.engines().await { // Memory off or signed out: nothing to show. Err(Error::Off(_)) => false, Err(error) => return Err(error), diff --git a/crates/tinymemory-tools/src/layout_migration/service_tests.rs b/crates/tinymemory-tools/src/layout_migration/service_tests.rs index 8f76232f..8dc74c2c 100644 --- a/crates/tinymemory-tools/src/layout_migration/service_tests.rs +++ b/crates/tinymemory-tools/src/layout_migration/service_tests.rs @@ -208,7 +208,7 @@ impl LayoutHost for CountingHost { &self.store } - fn engines(&self) -> crate::layout_migration::Result { + async fn engines(&self) -> crate::layout_migration::Result { if self.fail_engines { Err(crate::layout_migration::Error::Off( "engine unavailable".to_string(), @@ -219,7 +219,7 @@ impl LayoutHost for CountingHost { tree: self.inner.tree.clone(), }) } else { - self.inner.engines() + self.inner.engines().await } } @@ -239,8 +239,8 @@ impl LayoutHost for CountingHost { self.inner.free_now().await } - fn legacy_claim(&self) -> crate::layout_migration::Result>> { - self.inner.legacy_claim() + async fn legacy_claim(&self) -> crate::layout_migration::Result>> { + self.inner.legacy_claim().await } fn listed_unconfirmed(&self) -> bool { diff --git a/crates/tinymemory-tools/src/layout_migration/test_host.rs b/crates/tinymemory-tools/src/layout_migration/test_host.rs index d5083fbd..8211948a 100644 --- a/crates/tinymemory-tools/src/layout_migration/test_host.rs +++ b/crates/tinymemory-tools/src/layout_migration/test_host.rs @@ -100,7 +100,7 @@ impl LayoutHost for FakeHost { &self.store } - fn engines(&self) -> Result { + async fn engines(&self) -> Result { Ok(Engines { legacy: self.legacy.clone(), tree: self.tree.clone(), @@ -134,7 +134,7 @@ impl LayoutHost for FakeHost { self.free.load(Ordering::SeqCst) && left > 0 } - fn legacy_claim(&self) -> Result>> { + async fn legacy_claim(&self) -> Result>> { Ok(self .claim .clone() diff --git a/crates/tinymemory-tools/src/lifecycle/turn.rs b/crates/tinymemory-tools/src/lifecycle/turn.rs index a8f98cd0..be0cc1eb 100644 --- a/crates/tinymemory-tools/src/lifecycle/turn.rs +++ b/crates/tinymemory-tools/src/lifecycle/turn.rs @@ -56,13 +56,14 @@ pub trait Clock: Send + Sync { } /// Resolves the configured engine and identity to their lifecycle instance. +#[async_trait] pub trait EngineResolver: Send + Sync { /// Returns the configured memory instance for this host hook. /// /// # Errors /// /// Returns the host's engine or identity resolution failure. - fn resolve(&self) -> tinymemory_api::Result; + async fn resolve(&self) -> tinymemory_api::Result; } /// Receives background jobs returned by a successful post-turn log. @@ -310,7 +311,7 @@ impl TurnHooks { let work_clock = clock.clone(); let model = self.date_hint_model.clone(); let operation = async move { - let mut memory = resolver.resolve().ok()?; + let mut memory = resolver.resolve().await.ok()?; if deadline.is_some_and(|deadline| work_clock.now() >= deadline) { return None; } @@ -430,7 +431,7 @@ impl TurnHooks { if !logging { return; } - let Ok(memory) = self.resolver.resolve() else { + let Ok(memory) = self.resolver.resolve().await else { return; }; let Ok(report) = memory.post_turn(turn).await else { @@ -486,7 +487,7 @@ impl TurnHooks { let resolver = self.resolver.clone(); let clock = self.clock.clone(); let operation = async move { - let memory = resolver.resolve().ok()?; + let memory = resolver.resolve().await.ok()?; if deadline.is_some_and(|deadline| clock.now() >= deadline) { return None; } diff --git a/crates/tinymemory-tools/src/lifecycle/turn_runtime_tests.rs b/crates/tinymemory-tools/src/lifecycle/turn_runtime_tests.rs index dee04197..fb3ce185 100644 --- a/crates/tinymemory-tools/src/lifecycle/turn_runtime_tests.rs +++ b/crates/tinymemory-tools/src/lifecycle/turn_runtime_tests.rs @@ -138,8 +138,27 @@ struct StaticResolver { advance: Duration, } +struct DeferredResolver { + memory: AgentMemory, + release: Mutex>>, + calls: AtomicUsize, +} + +#[async_trait] +impl EngineResolver for DeferredResolver { + async fn resolve(&self) -> Result { + self.calls.fetch_add(1, Ordering::SeqCst); + let release = self.release.lock().unwrap().take(); + if let Some(release) = release { + let _ = release.await; + } + Ok(self.memory.clone()) + } +} + +#[async_trait] impl EngineResolver for StaticResolver { - fn resolve(&self) -> Result { + async fn resolve(&self) -> Result { self.calls.fetch_add(1, Ordering::SeqCst); self.clock.advance_by(self.advance); if self.fail { @@ -296,6 +315,41 @@ fn fixture_with_memory_and_advance( } } +#[tokio::test] +async fn pre_turn_awaits_async_engine_resolution() { + let engine = Arc::new(ReferenceEngine::new()); + let (release, wait) = oneshot::channel(); + let resolver = Arc::new(DeferredResolver { + memory: memory(engine), + release: Mutex::new(Some(wait)), + calls: AtomicUsize::new(0), + }); + let clock = FakeClock::at(Duration::ZERO); + let citations = Arc::new(RecordingCitations::default()); + let hooks = TurnHooks::new( + resolver.clone(), + Arc::new(clock.clone()), + None, + Arc::new(RecordingJobs::default()), + citations.clone(), + ); + let mut work = Box::pin(hooks.pre_turn( + PreTurn::new("thread", 0, "remember this"), + options(true, true), + )); + + poll_fn(|cx| match work.as_mut().poll(cx) { + Poll::Pending => Poll::Ready(()), + Poll::Ready(_) => panic!("pre-turn completed before async engine resolution"), + }) + .await; + assert_eq!(resolver.calls.load(Ordering::SeqCst), 1); + release.send(()).unwrap(); + + let _ = work.await; + assert_eq!(citations.calls.load(Ordering::SeqCst), 0); +} + fn options(logging: bool, recall: bool) -> PreTurnOptions { PreTurnOptions { logging,