diff --git a/.env.example b/.env.example index 481a6aab..1d9e9c3c 100644 --- a/.env.example +++ b/.env.example @@ -296,6 +296,8 @@ # MEMORY_MAX_MESSAGES_PER_RUN=200 # MEMORY_MIN_MESSAGES_PER_RUN=20 # MEMORY_MAX_QUEUED_RUNS=5000 +# Visible-card lists the memory pipeline keeps in RAM (LRU); 0 disables it. +# MEMORY_CARD_CACHE_ENTRIES=4096 # MEMORY_MAX_INPUT_TOKENS=10000 # MEMORY_TOKENIZER_MODEL=google/gemma-4-26B-A4B-it # MEMORY_TOKENIZER_FILE= diff --git a/crates/openplotva-app/src/dialog_jobs/input.rs b/crates/openplotva-app/src/dialog_jobs/input.rs index 92548adc..ff61ffdb 100644 --- a/crates/openplotva-app/src/dialog_jobs/input.rs +++ b/crates/openplotva-app/src/dialog_jobs/input.rs @@ -129,6 +129,7 @@ pub struct PostgresDialogInputMaterializer { history: PostgresHistoryStore, members: PostgresChatMemberStore, memory: Option, + memory_usage: Option, memory_embedder: Option>, memory_embedding_dim: i32, shield: Option, @@ -156,6 +157,7 @@ impl PostgresDialogInputMaterializer { history, members, memory: None, + memory_usage: None, memory_embedder: None, memory_embedding_dim: 0, shield: None, @@ -168,9 +170,16 @@ impl PostgresDialogInputMaterializer { } } + /// Recall memory from `memory` and record recalled cards in `usage`, the + /// write-behind journal that applies them to Postgres. #[must_use] - pub fn with_memory_store(mut self, memory: PostgresMemoryStore) -> Self { + pub fn with_memory_store( + mut self, + memory: PostgresMemoryStore, + usage: crate::memory_card_usage::MemoryCardUsageJournal, + ) -> Self { self.memory = Some(memory); + self.memory_usage = Some(usage); self } @@ -757,19 +766,9 @@ impl PostgresDialogInputMaterializer { .await { Ok(memory) => { - let card_ids = memory.cards.iter().map(|card| card.id).collect::>(); - if let Some(store) = self.memory.as_ref() - && let Err(error) = store - .record_card_retrieval(&card_ids, OffsetDateTime::now_utc()) - .await - { - tracing::warn!( - %error, - chat_id = params.chat_id, - user_id = params.user_id, - card_count = card_ids.len(), - "failed to record recalled memory cards" - ); + if let Some(journal) = self.memory_usage.as_ref() { + let card_ids = memory.cards.iter().map(|card| card.id).collect::>(); + journal.record(&card_ids, OffsetDateTime::now_utc()); } let captured = captured_memories_from_retrieval(&memory); ( diff --git a/crates/openplotva-app/src/lib.rs b/crates/openplotva-app/src/lib.rs index 181befde..d7066d1a 100644 --- a/crates/openplotva-app/src/lib.rs +++ b/crates/openplotva-app/src/lib.rs @@ -33,6 +33,7 @@ mod lane_checks; mod maintenance_api; pub mod media; pub mod members; +pub mod memory_card_usage; pub mod memory_runtime; pub mod message_gate; pub mod model_routing; @@ -350,6 +351,7 @@ struct RuntimeWorkers { openrouter_free_gate: Option>, router_runtime: Option>, dialog_worker_gauge: Option>, + memory_card_cache: Option>, readiness_probes: Vec, } @@ -796,7 +798,10 @@ fn static_web_routes_from_config( config.vip.chat_id, ))); } - let memory_store = PostgresMemoryStore::new(clients.postgres.clone()); + let mut memory_store = PostgresMemoryStore::new(clients.postgres.clone()); + if let Some(cache) = &runtime_workers.memory_card_cache { + memory_store = memory_store.with_card_cache(Arc::clone(cache)); + } routes.memory_store = Some(memory_store.clone()); if config.memory.enabled && let ( @@ -10535,6 +10540,7 @@ async fn start_runtime_workers( openrouter_free_gate: None, router_runtime: None, dialog_worker_gauge: None, + memory_card_cache: None, readiness_probes: Vec::new(), }; let readiness_general_postgres = service_clients.postgres.clone(); @@ -11237,7 +11243,52 @@ async fn start_runtime_workers( .with_run_buffer(llm_run_buffer.clone()), ); openplotva_llm::trace::set_observer(llm_observer); - let memory_store = PostgresMemoryStore::new(service_clients.postgres.clone()); + let memory_card_cache = usize::try_from(config.memory.card_cache_entries) + .ok() + .filter(|entries| *entries > 0) + .map(|entries| { + Arc::new(openplotva_storage::MemoryCardCache::new( + entries, + openplotva_storage::MEMORY_CARD_CACHE_TTL, + )) + }); + workers.memory_card_cache = memory_card_cache.clone(); + let memory_store = match &memory_card_cache { + Some(cache) => PostgresMemoryStore::new(service_clients.postgres.clone()) + .with_card_cache(Arc::clone(cache)), + None => PostgresMemoryStore::new(service_clients.postgres.clone()), + }; + let memory_card_usage = memory_card_usage::MemoryCardUsageJournal::new(); + { + let journal = memory_card_usage.clone(); + let store = memory_store.clone(); + let cache = memory_card_cache.clone(); + let stop = worker_stops.subscribe(RuntimeWorkerPhase::Ancillary); + let worker = tokio::spawn(async move { + let report = memory_card_usage::run_memory_card_usage_flush_worker_until( + journal, + store, + cache, + memory_card_usage::MEMORY_CARD_USAGE_FLUSH_INTERVAL, + wait_for_runtime_stop(stop), + ) + .await; + tracing::info!(?report, "memory card usage flush worker stopped"); + }); + readiness_checks.push(ReadinessCheck::ok( + "memory_card_usage_flush", + format!( + "memory card uses flushed every {}s; visible-card cache {} entries", + memory_card_usage::MEMORY_CARD_USAGE_FLUSH_INTERVAL.as_secs(), + config.memory.card_cache_entries + ), + )); + worker_registry.register( + "memory-card-usage-flush", + RuntimeWorkerPhase::Ancillary, + worker, + ); + } let memory_restart_trigger = Arc::new(tokio::sync::Notify::new()); let virtual_dialog_store = PostgresRuntimeVirtualDialogStore::new(service_clients.postgres.clone()); @@ -12944,7 +12995,8 @@ async fn start_runtime_workers( ), ); if dialog_memory_context_enabled(config) { - dialog_materializer = dialog_materializer.with_memory_store(memory_store.clone()); + dialog_materializer = dialog_materializer + .with_memory_store(memory_store.clone(), memory_card_usage.clone()); if let Some(embedder) = memory_query_embedder.clone() { dialog_materializer = dialog_materializer .with_memory_embedder(embedder, config.memory.embedding_dim); @@ -14099,6 +14151,7 @@ async fn shutdown_runtime_workers(workers: RuntimeWorkers) { openrouter_free_gate: _, router_runtime: _, dialog_worker_gauge: _, + memory_card_cache: _, readiness_probes: _, } = workers; let started_at = Instant::now(); diff --git a/crates/openplotva-app/src/memory_card_usage.rs b/crates/openplotva-app/src/memory_card_usage.rs new file mode 100644 index 00000000..7ec293a7 --- /dev/null +++ b/crates/openplotva-app/src/memory_card_usage.rs @@ -0,0 +1,181 @@ +//! Write-behind journal of memory-card uses. +//! +//! Dialog turns record which cards they recalled without waiting for Postgres; +//! a worker applies the collected uses in one statement per interval and once +//! more on shutdown. + +use std::{ + collections::HashMap, + sync::{Arc, Mutex, MutexGuard}, + time::Duration, +}; + +use openplotva_storage::{MemoryCardCache, MemoryCardUsage, PostgresMemoryStore}; +use time::OffsetDateTime; + +pub const MEMORY_CARD_USAGE_FLUSH_INTERVAL: Duration = Duration::from_secs(30); +/// Cache statistics are logged every this many flush ticks (5 minutes). +const CACHE_STATS_EVERY_TICKS: u64 = 10; + +#[derive(Clone, Debug, Default)] +pub struct MemoryCardUsageJournal { + pending: Arc>>, +} + +impl MemoryCardUsageJournal { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Count one use of each card at `used_at`. + pub fn record(&self, card_ids: &[i64], used_at: OffsetDateTime) { + let mut pending = self.lock(); + for card_id in card_ids.iter().copied().filter(|id| *id != 0) { + let entry = pending.entry(card_id).or_insert((used_at, 0)); + entry.0 = entry.0.max(used_at); + entry.1 += 1; + } + } + + fn drain(&self) -> Vec { + self.lock() + .drain() + .map(|(card_id, (last_used_at, uses))| MemoryCardUsage { + card_id, + last_used_at, + uses, + }) + .collect() + } + + /// Put back uses whose flush failed, merging with uses recorded since. + fn restore(&self, usage: Vec) { + let mut pending = self.lock(); + for entry in usage { + let slot = pending + .entry(entry.card_id) + .or_insert((entry.last_used_at, 0)); + slot.0 = slot.0.max(entry.last_used_at); + slot.1 += entry.uses; + } + } + + fn lock(&self) -> MutexGuard<'_, HashMap> { + self.pending + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} + +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct MemoryCardUsageWorkerReport { + pub flushes: u64, + pub cards_updated: u64, + pub errors: u64, +} + +/// Flush card uses every `interval` until `stop`, then flush what is left. +pub async fn run_memory_card_usage_flush_worker_until( + journal: MemoryCardUsageJournal, + store: PostgresMemoryStore, + card_cache: Option>, + interval: Duration, + stop: impl std::future::Future, +) -> MemoryCardUsageWorkerReport { + let mut report = MemoryCardUsageWorkerReport::default(); + let mut ticker = tokio::time::interval(interval); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + ticker.tick().await; + tokio::pin!(stop); + let mut ticks = 0u64; + loop { + tokio::select! { + () = &mut stop => break, + _ = ticker.tick() => { + ticks += 1; + flush_card_usage(&journal, &store, &mut report).await; + if let Some(cache) = &card_cache + && ticks.is_multiple_of(CACHE_STATS_EVERY_TICKS) + { + let stats = cache.stats(); + tracing::info!( + entries = stats.entries, + bytes = stats.bytes, + hits = stats.hits, + misses = stats.misses, + expired = stats.expired, + evictions = stats.evictions, + invalidated = stats.invalidated, + "memory card cache stats" + ); + } + } + } + } + flush_card_usage(&journal, &store, &mut report).await; + report +} + +async fn flush_card_usage( + journal: &MemoryCardUsageJournal, + store: &PostgresMemoryStore, + report: &mut MemoryCardUsageWorkerReport, +) { + let usage = journal.drain(); + if usage.is_empty() { + return; + } + match store.record_card_usage(&usage).await { + Ok(updated) => { + report.flushes += 1; + report.cards_updated += updated; + } + Err(error) => { + report.errors += 1; + tracing::warn!(%error, cards = usage.len(), "failed to flush memory card usage; keeping it for the next flush"); + journal.restore(usage); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn at(seconds: i64) -> OffsetDateTime { + OffsetDateTime::from_unix_timestamp(1_779_193_800 + seconds).expect("valid time") + } + + fn by_card(usage: Vec) -> Vec<(i64, OffsetDateTime, i64)> { + let mut usage = usage + .into_iter() + .map(|entry| (entry.card_id, entry.last_used_at, entry.uses)) + .collect::>(); + usage.sort_unstable(); + usage + } + + #[test] + fn journal_counts_uses_and_keeps_the_latest_time_per_card() { + let journal = MemoryCardUsageJournal::new(); + journal.record(&[1, 2, 0], at(10)); + journal.record(&[2], at(5)); + journal.record(&[1], at(20)); + + assert_eq!(by_card(journal.drain()), [(1, at(20), 2), (2, at(10), 2)]); + assert!(journal.drain().is_empty()); + } + + #[test] + fn failed_flush_is_merged_with_uses_recorded_meanwhile() { + let journal = MemoryCardUsageJournal::new(); + journal.record(&[1], at(10)); + let failed = journal.drain(); + journal.record(&[1, 3], at(30)); + + journal.restore(failed); + + assert_eq!(by_card(journal.drain()), [(1, at(30), 2), (3, at(30), 1)]); + } +} diff --git a/crates/openplotva-app/src/memory_runtime.rs b/crates/openplotva-app/src/memory_runtime.rs index a0148e9f..aaf9ed33 100644 --- a/crates/openplotva-app/src/memory_runtime.rs +++ b/crates/openplotva-app/src/memory_runtime.rs @@ -200,6 +200,10 @@ pub trait MemoryWriteStore: std::fmt::Debug + Send + Sync { prompt_version: &'a str, source_messages: &'a [openplotva_memory::Message], ) -> MemoryWriteStoreFuture<'a, (i64, bool), Self::Error>; + + /// Drop cached visible-card lists of a chat and its users after their + /// cards were written, so the next batch sees the new cards. + fn invalidate_visible_cards(&self, _chat_id: i64, _user_ids: &[i64]) {} } /// Storage boundary needed to claim and finalize memory runs. @@ -394,6 +398,10 @@ impl MemoryWriteStore for PostgresMemoryStore { .await }) } + + fn invalidate_visible_cards(&self, chat_id: i64, user_ids: &[i64]) { + self.invalidate_cached_cards(chat_id, user_ids); + } } impl MemoryRunStore for PostgresMemoryStore { @@ -441,7 +449,9 @@ impl MemoryRunStore for PostgresMemoryStore { scope: &'a openplotva_memory::RetrievalScope, limit: i32, ) -> MemoryWriteStoreFuture<'a, Vec, Self::Error> { - Box::pin(async move { PostgresMemoryStore::list_visible_cards(self, scope, limit).await }) + Box::pin( + async move { PostgresMemoryStore::list_visible_cards_cached(self, scope, limit).await }, + ) } fn retrieve_related_cards<'a>( @@ -3079,7 +3089,9 @@ where // batch. Log, put the failed group on the cooldown (so a persistently // failing group retries at the cooldown cadence instead of burning a model // call every tick and starving the others), and move on. - match process_one_merge_group(merger, store, embedder, cfg, now, &group).await { + let merged = process_one_merge_group(merger, store, embedder, cfg, now, &group).await; + store.invalidate_cached_cards(group.chat_id, &[group.user_id]); + match merged { Ok(Some(counts)) => { tick.cards_superseded += counts.superseded; tick.cards_demoted += counts.demoted; @@ -3169,6 +3181,31 @@ async fn write_memory_extraction_cards( output: &ExtractOutput, cfg: MemoryExtractionBatchConfig<'_>, ) -> Result> +where + ExtractorError: StdError + Send + Sync + 'static, + Store: MemoryWriteStore, + Embedder: EmbeddingProvider, +{ + let written = write_extraction_cards(store, embedder, input, output, cfg).await; + let mut user_ids = input + .messages + .iter() + .map(|message| message.user_id) + .filter(|user_id| *user_id != 0) + .collect::>(); + user_ids.sort_unstable(); + user_ids.dedup(); + store.invalidate_visible_cards(input.run.chat_id, &user_ids); + written +} + +async fn write_extraction_cards( + store: &Store, + embedder: Option<&Embedder>, + input: &ExtractInput, + output: &ExtractOutput, + cfg: MemoryExtractionBatchConfig<'_>, +) -> Result> where ExtractorError: StdError + Send + Sync + 'static, Store: MemoryWriteStore, diff --git a/crates/openplotva-config/src/lib.rs b/crates/openplotva-config/src/lib.rs index dff061d9..0ac213d3 100644 --- a/crates/openplotva-config/src/lib.rs +++ b/crates/openplotva-config/src/lib.rs @@ -383,6 +383,9 @@ pub const DEFAULT_MEMORY_MIN_MESSAGES_PER_RUN: i32 = 20; pub const DEFAULT_MEMORY_MAX_QUEUED_RUNS: i32 = 5_000; +/// Visible-card lists the memory pipeline keeps in RAM; `0` disables the cache. +pub const DEFAULT_MEMORY_CARD_CACHE_ENTRIES: i32 = 4_096; + pub const DEFAULT_MEMORY_TOKENIZER_MODEL: &str = "google/gemma-4-26B-A4B-it"; /// Discovery service name the embedder registers under on the AI Farm. @@ -1045,6 +1048,7 @@ pub struct MemoryConfig { pub max_messages_per_run: i32, pub min_messages_per_run: i32, pub max_queued_runs: i32, + pub card_cache_entries: i32, pub max_input_tokens: i32, pub tokenizer_model: String, pub tokenizer_file: String, @@ -1571,6 +1575,8 @@ pub struct RawConfig { pub memory_min_messages_per_run: Option, /// `MEMORY_MAX_QUEUED_RUNS`. pub memory_max_queued_runs: Option, + /// `MEMORY_CARD_CACHE_ENTRIES`. + pub memory_card_cache_entries: Option, /// `MEMORY_MAX_INPUT_TOKENS`. pub memory_max_input_tokens: Option, /// `MEMORY_TOKENIZER_MODEL`. @@ -2895,6 +2901,14 @@ impl AppConfig { DEFAULT_MEMORY_MAX_QUEUED_RUNS, )?, )?, + card_cache_entries: validate_non_negative_i32( + "MEMORY_CARD_CACHE_ENTRIES", + parse_i32( + "MEMORY_CARD_CACHE_ENTRIES", + raw.memory_card_cache_entries, + DEFAULT_MEMORY_CARD_CACHE_ENTRIES, + )?, + )?, max_input_tokens: parse_i32( "MEMORY_MAX_INPUT_TOKENS", raw.memory_max_input_tokens, @@ -3379,6 +3393,7 @@ impl RawConfig { memory_max_messages_per_run: env("MEMORY_MAX_MESSAGES_PER_RUN"), memory_min_messages_per_run: env("MEMORY_MIN_MESSAGES_PER_RUN"), memory_max_queued_runs: env("MEMORY_MAX_QUEUED_RUNS"), + memory_card_cache_entries: env("MEMORY_CARD_CACHE_ENTRIES"), memory_max_input_tokens: env("MEMORY_MAX_INPUT_TOKENS"), memory_tokenizer_model: env("MEMORY_TOKENIZER_MODEL"), memory_tokenizer_file: env("MEMORY_TOKENIZER_FILE"), @@ -4979,6 +4994,15 @@ mod tests { let config = AppConfig::from_raw(RawConfig::default())?; assert_eq!(config.memory.min_messages_per_run, 20); assert_eq!(config.memory.max_queued_runs, 5000); + assert_eq!( + config.memory.card_cache_entries, + super::DEFAULT_MEMORY_CARD_CACHE_ENTRIES + ); + let disabled = AppConfig::from_raw(RawConfig { + memory_card_cache_entries: Some("0".to_owned()), + ..RawConfig::default() + })?; + assert_eq!(disabled.memory.card_cache_entries, 0); let error = AppConfig::from_raw(RawConfig { memory_min_messages_per_run: Some("0".to_owned()), diff --git a/crates/openplotva-storage/src/lib.rs b/crates/openplotva-storage/src/lib.rs index 28b8fc40..ad1db210 100644 --- a/crates/openplotva-storage/src/lib.rs +++ b/crates/openplotva-storage/src/lib.rs @@ -43,11 +43,13 @@ pub mod gradius_ads; pub mod llm_routing; pub mod maintenance; pub mod maintenance_notifications; +mod memory_card_cache; mod task_queue_listing; pub mod telegram_delivery; pub mod telegram_outbox; pub mod telegram_projection; +pub use memory_card_cache::{MEMORY_CARD_CACHE_TTL, MemoryCardCache, MemoryCardCacheStats}; pub use task_queue_listing::{ TASK_QUEUE_ACTIVE_STATUSES, TaskQueueJobFilter, TaskQueueJobPage, TaskQueueSortColumn, TaskQueueTimeColumn, @@ -963,7 +965,7 @@ pub const SQL_UPDATE_MEMORY_CARD_TEXT: &str = "UPDATE memory_cards SET fact_text pub const SQL_REINFORCE_MEMORY_CARD: &str = "UPDATE memory_cards SET confidence = LEAST(1.0, confidence + $2), salience = LEAST(1.0, salience + $3), observation_count = observation_count + 1, last_observed_at = GREATEST(last_observed_at, $4), expires_at = NULL, updated_at = CURRENT_TIMESTAMP WHERE id = $1 AND status IN ('active', 'competing')"; -pub const SQL_RECORD_MEMORY_CARD_USAGE: &str = "UPDATE memory_cards SET last_used_at = CASE WHEN last_used_at IS NULL THEN $2 ELSE GREATEST(last_used_at, $2) END, use_count = LEAST(use_count::bigint + 1, 2147483647)::integer, updated_at = CURRENT_TIMESTAMP WHERE id = ANY($1::bigint[]) AND status IN ('active', 'competing')"; +pub const SQL_RECORD_MEMORY_CARD_USAGE: &str = "UPDATE memory_cards AS card SET last_used_at = CASE WHEN card.last_used_at IS NULL THEN usage.used_at ELSE GREATEST(card.last_used_at, usage.used_at) END, use_count = LEAST(card.use_count::bigint + usage.uses, 2147483647)::integer, updated_at = CURRENT_TIMESTAMP FROM unnest($1::bigint[], $2::timestamptz[], $3::bigint[]) AS usage(id, used_at, uses) WHERE card.id = usage.id AND card.status IN ('active', 'competing')"; pub const SQL_DEMOTE_MEMORY_CARD: &str = "UPDATE memory_cards SET confidence = GREATEST(0.0, confidence - $2), updated_at = CURRENT_TIMESTAMP WHERE id = $1 AND status IN ('active', 'competing')"; @@ -5168,6 +5170,15 @@ impl PostgresHistoryStore { #[derive(Clone, Debug)] pub struct PostgresMemoryStore { pool: PgPool, + card_cache: Option>, +} + +/// Uses of one card collected since the last flush. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub struct MemoryCardUsage { + pub card_id: i64, + pub last_used_at: OffsetDateTime, + pub uses: i64, } async fn merge_memory_card_into_on( @@ -5202,7 +5213,32 @@ impl PostgresMemoryStore { /// Build a memory store over an existing pool. #[must_use] pub fn new(pool: PgPool) -> Self { - Self { pool } + Self { + pool, + card_cache: None, + } + } + + /// Share one visible-card cache; stores built from the same cache see each + /// other's invalidations. + #[must_use] + pub fn with_card_cache(mut self, cache: Arc) -> Self { + self.card_cache = Some(cache); + self + } + + /// Drop cached lists for `chat_id` and for lists that include cards of + /// `user_ids`; call after writing cards of that scope. + pub fn invalidate_cached_cards(&self, chat_id: i64, user_ids: &[i64]) { + if let Some(cache) = &self.card_cache { + cache.invalidate_scope(chat_id, user_ids); + } + } + + fn clear_card_cache(&self) { + if let Some(cache) = &self.card_cache { + cache.clear(); + } } pub async fn get_dialog_memory_chat_meta( @@ -5484,11 +5520,15 @@ impl PostgresMemoryStore { /// Archive one batch of active cards whose TTL has elapsed (soft, restorable: /// status='expired', retracted_at stamped). Returns the number archived. pub async fn archive_expired_cards(&self, batch: i64) -> Result { - Ok(sqlx::query(SQL_ARCHIVE_EXPIRED_MEMORY_CARDS) + let archived = sqlx::query(SQL_ARCHIVE_EXPIRED_MEMORY_CARDS) .bind(batch.max(1)) .execute(&self.pool) .await? - .rows_affected()) + .rows_affected(); + if archived > 0 { + self.clear_card_cache(); + } + Ok(archived) } /// Give cold episodic cards a grace TTL so forgetting is bounded even when @@ -5547,6 +5587,9 @@ impl PostgresMemoryStore { tx.commit().await?; retired += dup_ids.len() as u64; } + if retired > 0 { + self.clear_card_cache(); + } Ok(retired) } @@ -6000,6 +6043,32 @@ impl PostgresMemoryStore { .collect::, _>>() } + /// `list_visible_cards` through the RAM cache, when this store has one. + pub async fn list_visible_cards_cached( + &self, + scope: &openplotva_memory::RetrievalScope, + limit: i32, + ) -> Result, StorageError> { + let Some(cache) = &self.card_cache else { + return self.list_visible_cards(scope, limit).await; + }; + let key = memory_card_cache::VisibleCardsKey { + chat_id: scope.chat_id, + thread_id: scope.thread_id, + user_id: scope.user_id, + own_dm: openplotva_memory::public_user_in_own_dm(scope), + limit: positive_or_default(limit, 100), + }; + match cache.get(&key) { + Ok(cards) => Ok(cards), + Err(generation) => { + let cards = self.list_visible_cards(scope, key.limit).await?; + cache.put(key, &cards, generation); + Ok(cards) + } + } + } + pub async fn soft_delete_visible_card( &self, id: i64, @@ -6018,6 +6087,7 @@ impl PostgresMemoryStore { .bind(openplotva_memory::public_user_in_own_dm(scope)) .execute(&self.pool) .await?; + self.invalidate_cached_cards(scope.chat_id, &[scope.user_id]); Ok(result.rows_affected() > 0) } @@ -6235,6 +6305,7 @@ impl PostgresMemoryStore { .bind(id) .execute(&self.pool) .await?; + self.clear_card_cache(); Ok(()) } @@ -6256,6 +6327,7 @@ impl PostgresMemoryStore { .bind(portable) .execute(&self.pool) .await?; + self.clear_card_cache(); Ok(()) } @@ -6268,32 +6340,45 @@ impl PostgresMemoryStore { .bind(id) .execute(&self.pool) .await?; + self.clear_card_cache(); Ok(()) } - /// Record cards that were actually included in a dialog context. - pub async fn record_card_retrieval( - &self, - card_ids: &[i64], - used_at: OffsetDateTime, - ) -> Result<(), StorageError> { - let mut card_ids = card_ids + /// Apply collected card uses in one statement. `updated_at` moves with + /// usage because deferred reconsolidation selects recently touched cards. + pub async fn record_card_usage(&self, usage: &[MemoryCardUsage]) -> Result { + let mut usage = usage .iter() + .filter(|entry| entry.card_id != 0 && entry.uses > 0) .copied() - .filter(|id| *id != 0) - .collect::>() - .into_iter() .collect::>(); - card_ids.sort_unstable(); - if card_ids.is_empty() { - return Ok(()); + if usage.is_empty() { + return Ok(0); } - sqlx::query(SQL_RECORD_MEMORY_CARD_USAGE) - .bind(&card_ids) - .bind(used_at) + usage.sort_unstable_by_key(|entry| entry.card_id); + // One row per card: a repeated id would join several usage rows and + // apply only one of them. + usage.dedup_by(|next, kept| { + if kept.card_id != next.card_id { + return false; + } + kept.uses += next.uses; + kept.last_used_at = kept.last_used_at.max(next.last_used_at); + true + }); + let ids = usage.iter().map(|entry| entry.card_id).collect::>(); + let used_at = usage + .iter() + .map(|entry| entry.last_used_at) + .collect::>(); + let uses = usage.iter().map(|entry| entry.uses).collect::>(); + Ok(sqlx::query(SQL_RECORD_MEMORY_CARD_USAGE) + .bind(&ids) + .bind(&used_at) + .bind(&uses) .execute(&self.pool) - .await?; - Ok(()) + .await? + .rows_affected()) } pub async fn retrieve_lexical( @@ -11730,7 +11815,9 @@ mod tests { ); assert!(super::SQL_REINFORCE_MEMORY_CARD.contains("expires_at = NULL")); assert!(super::SQL_RECORD_MEMORY_CARD_USAGE.contains("last_used_at = CASE")); - assert!(super::SQL_RECORD_MEMORY_CARD_USAGE.contains("use_count::bigint + 1")); + assert!( + super::SQL_RECORD_MEMORY_CARD_USAGE.contains("card.use_count::bigint + usage.uses") + ); assert!(super::SQL_COPY_MEMORY_CARD_SOURCES.contains("source.card_id = $2")); assert!(super::SQL_COPY_MEMORY_CARD_LINKS.contains("ON CONFLICT")); assert!( @@ -12252,20 +12339,39 @@ mod tests { assert_eq!(retrieved.episodes[0].id, episode_id); let retrieved_card_ids = retrieved.cards.iter().map(|card| card.id).collect::>(); - store - .record_card_retrieval( - &retrieved_card_ids, - observed + time::Duration::hours(1), - ) - .await?; + let mut usage = retrieved_card_ids + .iter() + .map(|card_id| super::MemoryCardUsage { + card_id: *card_id, + last_used_at: observed + time::Duration::hours(1), + uses: 1, + }) + .collect::>(); + usage.push(super::MemoryCardUsage { + card_id: retrieved_card_ids[0], + last_used_at: observed + time::Duration::minutes(30), + uses: 2, + }); + store.record_card_usage(&usage).await?; let used_cards: i64 = sqlx::query_scalar( "SELECT count(*) FROM memory_cards WHERE id = ANY($1::bigint[]) AND use_count = 1 AND last_used_at = $2", ) - .bind(&retrieved_card_ids) + .bind(&retrieved_card_ids[1..]) .bind(observed + time::Duration::hours(1)) .fetch_one(&pool) .await?; - assert_eq!(used_cards, retrieved.cards.len() as i64); + assert_eq!(used_cards, retrieved.cards.len() as i64 - 1); + let repeated: (i32, time::OffsetDateTime) = sqlx::query_as( + "SELECT use_count, last_used_at FROM memory_cards WHERE id = $1", + ) + .bind(retrieved_card_ids[0]) + .fetch_one(&pool) + .await?; + assert_eq!( + repeated, + (3, observed + time::Duration::hours(1)), + "a repeated card sums its uses and keeps its latest use" + ); let listed = store .list_cards(&openplotva_memory::CardFilter { diff --git a/crates/openplotva-storage/src/memory_card_cache.rs b/crates/openplotva-storage/src/memory_card_cache.rs new file mode 100644 index 00000000..983d2358 --- /dev/null +++ b/crates/openplotva-storage/src/memory_card_cache.rs @@ -0,0 +1,306 @@ +//! Read-through RAM cache of visible memory-card lists. +//! +//! The memory pipeline lists the same chat and participant cards many times +//! while it sizes one run's batches. Entries are dropped when a write touches +//! their chat or user, and expire after a TTL that bounds staleness from +//! writers that do not invalidate. + +use std::{ + collections::{BTreeMap, HashMap}, + fmt, + sync::{Arc, Mutex, MutexGuard}, + time::{Duration, Instant}, +}; + +use openplotva_memory::Card; + +/// Bounds how long a list can miss writes made outside the invalidating paths. +pub const MEMORY_CARD_CACHE_TTL: Duration = Duration::from_secs(5 * 60); + +/// Rough per-card overhead on top of its text, for the size estimate. +const CARD_OVERHEAD_BYTES: usize = 256; + +#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)] +pub(crate) struct VisibleCardsKey { + pub(crate) chat_id: i64, + pub(crate) thread_id: i32, + pub(crate) user_id: i64, + pub(crate) own_dm: bool, + pub(crate) limit: i32, +} + +struct Entry { + cards: Arc>, + stored_at: Instant, + recency: u64, + bytes: usize, +} + +#[derive(Default)] +struct CacheState { + entries: HashMap, + by_recency: BTreeMap, + next_recency: u64, + generation: u64, + bytes: usize, + stats: MemoryCardCacheStats, +} + +/// Counters for sizing the cache from production traffic. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub struct MemoryCardCacheStats { + pub hits: u64, + pub misses: u64, + pub expired: u64, + pub evictions: u64, + pub invalidated: u64, + pub entries: usize, + pub bytes: usize, +} + +pub struct MemoryCardCache { + capacity: usize, + ttl: Duration, + state: Mutex, +} + +impl fmt::Debug for MemoryCardCache { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("MemoryCardCache") + .field("capacity", &self.capacity) + .field("ttl", &self.ttl) + .field("stats", &self.stats()) + .finish() + } +} + +impl MemoryCardCache { + #[must_use] + pub fn new(capacity: usize, ttl: Duration) -> Self { + Self { + capacity: capacity.max(1), + ttl, + state: Mutex::new(CacheState::default()), + } + } + + /// Return a fresh entry, or the generation a caller must pass to `put` + /// after it reads the database. + pub(crate) fn get(&self, key: &VisibleCardsKey) -> Result, u64> { + self.get_at(key, Instant::now()) + } + + fn get_at(&self, key: &VisibleCardsKey, now: Instant) -> Result, u64> { + let mut state = self.lock(); + let Some(entry) = state.entries.get(key) else { + state.stats.misses += 1; + return Err(state.generation); + }; + if now.saturating_duration_since(entry.stored_at) > self.ttl { + state.stats.expired += 1; + state.stats.misses += 1; + remove_entry(&mut state, key); + return Err(state.generation); + } + let cards = Arc::clone(&entry.cards); + let previous = entry.recency; + let recency = state.next_recency; + state.next_recency += 1; + state.by_recency.remove(&previous); + state.by_recency.insert(recency, *key); + if let Some(entry) = state.entries.get_mut(key) { + entry.recency = recency; + } + state.stats.hits += 1; + drop(state); + Ok(cards.as_ref().clone()) + } + + /// Store a list read at `generation`; a write that invalidated the cache + /// since then makes the list possibly stale, so it is dropped. + pub(crate) fn put(&self, key: VisibleCardsKey, cards: &[Card], generation: u64) { + self.put_at(key, cards, generation, Instant::now()); + } + + fn put_at(&self, key: VisibleCardsKey, cards: &[Card], generation: u64, now: Instant) { + let mut state = self.lock(); + if state.generation != generation { + return; + } + remove_entry(&mut state, &key); + let bytes = cards.iter().map(card_bytes).sum(); + let recency = state.next_recency; + state.next_recency += 1; + state.by_recency.insert(recency, key); + state.bytes += bytes; + state.entries.insert( + key, + Entry { + cards: Arc::new(cards.to_vec()), + stored_at: now, + recency, + bytes, + }, + ); + while state.entries.len() > self.capacity { + let Some((_, oldest)) = state.by_recency.pop_first() else { + break; + }; + if let Some(entry) = state.entries.remove(&oldest) { + state.bytes -= entry.bytes; + state.stats.evictions += 1; + } + } + } + + /// Drop lists of `chat_id` and lists that include cards of `user_ids`. + pub fn invalidate_scope(&self, chat_id: i64, user_ids: &[i64]) { + let mut state = self.lock(); + state.generation += 1; + let stale: Vec = state + .entries + .keys() + .filter(|key| { + key.chat_id == chat_id || (key.user_id != 0 && user_ids.contains(&key.user_id)) + }) + .copied() + .collect(); + state.stats.invalidated += stale.len() as u64; + for key in &stale { + remove_entry(&mut state, key); + } + } + + /// Drop every list, for writes whose chats and users are not known. + pub fn clear(&self) { + let mut state = self.lock(); + state.generation += 1; + state.stats.invalidated += state.entries.len() as u64; + state.entries.clear(); + state.by_recency.clear(); + state.bytes = 0; + } + + #[must_use] + pub fn stats(&self) -> MemoryCardCacheStats { + let state = self.lock(); + MemoryCardCacheStats { + entries: state.entries.len(), + bytes: state.bytes, + ..state.stats + } + } + + fn lock(&self) -> MutexGuard<'_, CacheState> { + self.state + .lock() + .unwrap_or_else(|poisoned| poisoned.into_inner()) + } +} + +fn remove_entry(state: &mut CacheState, key: &VisibleCardsKey) { + if let Some(entry) = state.entries.remove(key) { + state.by_recency.remove(&entry.recency); + state.bytes -= entry.bytes; + } +} + +fn card_bytes(card: &Card) -> usize { + CARD_OVERHEAD_BYTES + + card.fact_text.len() + + card.subject.len() + + card.predicate.len() + + card.object.len() +} + +#[cfg(test)] +mod tests { + use super::*; + + fn key(chat_id: i64, user_id: i64) -> VisibleCardsKey { + VisibleCardsKey { + chat_id, + thread_id: 0, + user_id, + own_dm: false, + limit: 80, + } + } + + fn card(id: i64) -> Card { + Card { + id, + fact_text: format!("fact {id}"), + ..Card::default() + } + } + + #[test] + fn hit_returns_the_stored_list_until_the_ttl_passes() { + let cache = MemoryCardCache::new(8, Duration::from_secs(60)); + let start = Instant::now(); + let generation = cache.get_at(&key(1, 0), start).expect_err("cold cache"); + cache.put_at(key(1, 0), &[card(10)], generation, start); + + let cards = cache + .get_at(&key(1, 0), start + Duration::from_secs(30)) + .expect("fresh entry"); + assert_eq!(cards.iter().map(|card| card.id).collect::>(), [10]); + assert!( + cache + .get_at(&key(1, 0), start + Duration::from_secs(61)) + .is_err() + ); + let stats = cache.stats(); + assert_eq!((stats.hits, stats.misses, stats.expired), (1, 2, 1)); + assert_eq!(stats.entries, 0); + } + + #[test] + fn capacity_evicts_the_least_recently_used_list() { + let cache = MemoryCardCache::new(2, Duration::from_secs(60)); + let now = Instant::now(); + for chat_id in [1, 2] { + let generation = cache.get_at(&key(chat_id, 0), now).expect_err("cold"); + cache.put_at(key(chat_id, 0), &[card(chat_id)], generation, now); + } + cache.get_at(&key(1, 0), now).expect("chat 1 cached"); + let generation = cache.get_at(&key(3, 0), now).expect_err("cold"); + cache.put_at(key(3, 0), &[card(3)], generation, now); + + assert!(cache.get_at(&key(1, 0), now).is_ok()); + assert!(cache.get_at(&key(2, 0), now).is_err()); + assert!(cache.get_at(&key(3, 0), now).is_ok()); + assert_eq!(cache.stats().evictions, 1); + } + + #[test] + fn invalidation_drops_the_chat_and_its_users_lists_in_other_chats() { + let cache = MemoryCardCache::new(8, Duration::from_secs(60)); + let now = Instant::now(); + for (chat_id, user_id) in [(1, 0), (1, 7), (2, 7), (2, 8)] { + let generation = cache.get_at(&key(chat_id, user_id), now).expect_err("cold"); + cache.put_at(key(chat_id, user_id), &[card(1)], generation, now); + } + + cache.invalidate_scope(1, &[7]); + + assert!(cache.get_at(&key(1, 0), now).is_err()); + assert!(cache.get_at(&key(1, 7), now).is_err()); + assert!(cache.get_at(&key(2, 7), now).is_err()); + assert!(cache.get_at(&key(2, 8), now).is_ok()); + assert_eq!(cache.stats().invalidated, 3); + } + + #[test] + fn a_list_read_before_an_invalidation_is_not_stored() { + let cache = MemoryCardCache::new(8, Duration::from_secs(60)); + let now = Instant::now(); + let generation = cache.get_at(&key(1, 0), now).expect_err("cold"); + cache.invalidate_scope(1, &[]); + cache.put_at(key(1, 0), &[card(1)], generation, now); + + assert!(cache.get_at(&key(1, 0), now).is_err()); + assert_eq!(cache.stats().entries, 0); + } +}