Merge pull request 'docs(clawhdf5): document DType variants, fix unresolved doc links' (#17) from sdlc-docs/clawhdf5-types-20260514-165210 into main
This commit is contained in:
redclawsystems
@@ -0,0 +1,923 @@
|
||||
//! Async wrapper for [`HDF5Memory`] with background flush support.
|
||||
//!
|
||||
//! Provides non-blocking access to the synchronous HDF5Memory core by
|
||||
//! offloading all I/O and CPU-bound work to `spawn_blocking`. Includes
|
||||
//! a background flush task that:
|
||||
//!
|
||||
//! - **Batches saves** through an mpsc channel (avoids per-entry disk writes)
|
||||
//! - **Auto-flushes** on a configurable interval (default 5s)
|
||||
//! - **Threshold-flushes** when pending WAL entries exceed a limit
|
||||
//! - **Tracks dirty state** to skip no-op flushes
|
||||
//!
|
||||
//! # Architecture
|
||||
//!
|
||||
//! ```text
|
||||
//! ┌───────────────┐ ┌──────────────┐
|
||||
//! │ AsyncHDF5 │ spawn_blocking │ HDF5Memory │
|
||||
//! │ Memory │ ─────────────────── │ (sync core) │
|
||||
//! └───────┬───────┘ └──────────────┘
|
||||
//! │
|
||||
//! │ mpsc channel (saves + commands)
|
||||
//! ▼
|
||||
//! ┌───────────────┐
|
||||
//! │ Background │ interval timer + threshold check
|
||||
//! │ Writer Task │ batch save → WAL → periodic .h5 merge
|
||||
//! └───────────────┘
|
||||
//! ```
|
||||
//!
|
||||
//! # Usage
|
||||
//!
|
||||
//! ```ignore
|
||||
//! use clawhdf5_agent::async_memory::{AsyncHDF5Memory, AsyncConfig};
|
||||
//!
|
||||
//! let config = AsyncConfig {
|
||||
//! flush_interval: Duration::from_secs(10),
|
||||
//! flush_threshold: 100,
|
||||
//! };
|
||||
//! let mem = AsyncHDF5Memory::open_with(path, config).await?;
|
||||
//! mem.save(entry).await?; // buffered → background writer
|
||||
//! mem.save_batch(entries).await?; // also buffered
|
||||
//! let results = mem.hybrid_search(emb, "query".into(), 0.7, 0.3, 5).await;
|
||||
//! mem.shutdown().await?; // final flush + stop
|
||||
//! ```
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use tokio::sync::{Mutex, mpsc, oneshot};
|
||||
use tokio::task::spawn_blocking;
|
||||
|
||||
use crate::memory_strategy::{Exchange, StrategyOutput};
|
||||
use crate::{
|
||||
AgentMemory, HDF5Memory, MemoryConfig, MemoryEntry, MemoryError, Result, SearchResult,
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Configuration
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Configuration for the async background writer.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct AsyncConfig {
|
||||
/// How often the background task auto-flushes WAL → .h5.
|
||||
/// Set to `Duration::ZERO` to disable periodic flush (threshold-only).
|
||||
pub flush_interval: Duration,
|
||||
|
||||
/// Flush WAL → .h5 when pending WAL entries reach this count.
|
||||
/// Set to `0` to disable threshold-based flush (interval-only).
|
||||
pub flush_threshold: usize,
|
||||
|
||||
/// Channel capacity for buffered save commands.
|
||||
/// Higher = more batching, more memory. Default 256.
|
||||
pub channel_capacity: usize,
|
||||
}
|
||||
|
||||
impl Default for AsyncConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
flush_interval: Duration::from_secs(5),
|
||||
flush_threshold: 200,
|
||||
channel_capacity: 256,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Background writer commands
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
enum WriteCmd {
|
||||
/// Buffer one or more entries for saving.
|
||||
Save {
|
||||
entries: Vec<MemoryEntry>,
|
||||
reply: oneshot::Sender<Result<Vec<usize>>>,
|
||||
},
|
||||
/// Flush WAL → .h5 now.
|
||||
FlushNow(oneshot::Sender<Result<()>>),
|
||||
/// Tick session (decay + flush).
|
||||
TickSession(oneshot::Sender<Result<()>>),
|
||||
/// Shut down the background task.
|
||||
Shutdown(oneshot::Sender<()>),
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// AsyncHDF5Memory
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Async wrapper around [`HDF5Memory`].
|
||||
///
|
||||
/// Saves are buffered through a channel and batched by the background
|
||||
/// writer task. Reads/searches use `spawn_blocking` directly (they need
|
||||
/// the latest state, so they acquire the lock and run immediately).
|
||||
pub struct AsyncHDF5Memory {
|
||||
inner: Arc<Mutex<HDF5Memory>>,
|
||||
write_tx: mpsc::Sender<WriteCmd>,
|
||||
}
|
||||
|
||||
impl AsyncHDF5Memory {
|
||||
// -- Construction -------------------------------------------------------
|
||||
|
||||
/// Create a new HDF5 memory file with default async config.
|
||||
pub async fn create(config: MemoryConfig) -> Result<Self> {
|
||||
Self::create_with(config, AsyncConfig::default()).await
|
||||
}
|
||||
|
||||
/// Create a new HDF5 memory file with custom async config.
|
||||
pub async fn create_with(config: MemoryConfig, async_config: AsyncConfig) -> Result<Self> {
|
||||
let mem = spawn_blocking(move || HDF5Memory::create(config))
|
||||
.await
|
||||
.map_err(join_err)??;
|
||||
Ok(Self::wrap_with(mem, async_config))
|
||||
}
|
||||
|
||||
/// Open an existing HDF5 memory file with default async config.
|
||||
pub async fn open(path: impl AsRef<Path>) -> Result<Self> {
|
||||
Self::open_with(path, AsyncConfig::default()).await
|
||||
}
|
||||
|
||||
/// Open an existing HDF5 memory file with custom async config.
|
||||
pub async fn open_with(path: impl AsRef<Path>, async_config: AsyncConfig) -> Result<Self> {
|
||||
let path = path.as_ref().to_path_buf();
|
||||
let mem = spawn_blocking(move || HDF5Memory::open(&path))
|
||||
.await
|
||||
.map_err(join_err)??;
|
||||
Ok(Self::wrap_with(mem, async_config))
|
||||
}
|
||||
|
||||
/// Wrap a sync `HDF5Memory` with default async config.
|
||||
pub fn wrap(mem: HDF5Memory) -> Self {
|
||||
Self::wrap_with(mem, AsyncConfig::default())
|
||||
}
|
||||
|
||||
/// Wrap a sync `HDF5Memory` with custom async config.
|
||||
pub fn wrap_with(mem: HDF5Memory, async_config: AsyncConfig) -> Self {
|
||||
let inner = Arc::new(Mutex::new(mem));
|
||||
let (write_tx, write_rx) = mpsc::channel(async_config.channel_capacity);
|
||||
let bg_inner = Arc::clone(&inner);
|
||||
tokio::spawn(background_writer(bg_inner, write_rx, async_config));
|
||||
Self { inner, write_tx }
|
||||
}
|
||||
|
||||
// -- Buffered save ops --------------------------------------------------
|
||||
|
||||
/// Save a single memory entry. Buffered through the background writer.
|
||||
///
|
||||
/// Returns the entry index once the background writer has applied it
|
||||
/// to the in-memory cache (does NOT wait for .h5 flush).
|
||||
pub async fn save(&self, entry: MemoryEntry) -> Result<usize> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.write_tx
|
||||
.send(WriteCmd::Save {
|
||||
entries: vec![entry],
|
||||
reply: tx,
|
||||
})
|
||||
.await
|
||||
.map_err(|_| channel_gone())?;
|
||||
let indices = rx.await.map_err(|_| channel_gone())??;
|
||||
Ok(indices[0])
|
||||
}
|
||||
|
||||
/// Save a batch of entries. Buffered through the background writer.
|
||||
pub async fn save_batch(&self, entries: Vec<MemoryEntry>) -> Result<Vec<usize>> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.write_tx
|
||||
.send(WriteCmd::Save { entries, reply: tx })
|
||||
.await
|
||||
.map_err(|_| channel_gone())?;
|
||||
rx.await.map_err(|_| channel_gone())?
|
||||
}
|
||||
|
||||
// -- Direct mutating ops (not buffered — need immediate consistency) ----
|
||||
|
||||
/// Delete a memory entry by index.
|
||||
pub async fn delete(&self, id: usize) -> Result<()> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mut mem = inner.blocking_lock();
|
||||
mem.delete(id)
|
||||
})
|
||||
.await
|
||||
.map_err(join_err)?
|
||||
}
|
||||
|
||||
/// Compact tombstoned entries.
|
||||
pub async fn compact(&self) -> Result<usize> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mut mem = inner.blocking_lock();
|
||||
mem.compact()
|
||||
})
|
||||
.await
|
||||
.map_err(join_err)?
|
||||
}
|
||||
|
||||
/// Record an exchange using the configured memory strategy.
|
||||
pub async fn record(&self, exchange: Exchange) -> Result<StrategyOutput> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mut mem = inner.blocking_lock();
|
||||
mem.record(exchange)
|
||||
})
|
||||
.await
|
||||
.map_err(join_err)?
|
||||
}
|
||||
|
||||
// -- Search ops ---------------------------------------------------------
|
||||
|
||||
/// Hybrid vector + BM25 search.
|
||||
pub async fn hybrid_search(
|
||||
&self,
|
||||
query_embedding: Vec<f32>,
|
||||
query_text: String,
|
||||
vector_weight: f32,
|
||||
keyword_weight: f32,
|
||||
k: usize,
|
||||
) -> Vec<SearchResult> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mut mem = inner.blocking_lock();
|
||||
mem.hybrid_search(
|
||||
&query_embedding,
|
||||
&query_text,
|
||||
vector_weight,
|
||||
keyword_weight,
|
||||
k,
|
||||
)
|
||||
})
|
||||
.await
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
// -- Read ops -----------------------------------------------------------
|
||||
|
||||
/// Number of entries (including tombstoned).
|
||||
pub async fn count(&self) -> usize {
|
||||
self.inner.lock().await.count()
|
||||
}
|
||||
|
||||
/// Number of active (non-tombstoned) entries.
|
||||
pub async fn count_active(&self) -> usize {
|
||||
self.inner.lock().await.count_active()
|
||||
}
|
||||
|
||||
/// Get a chunk by index.
|
||||
pub async fn get_chunk(&self, index: usize) -> Option<String> {
|
||||
self.inner.lock().await.get_chunk(index).map(String::from)
|
||||
}
|
||||
|
||||
/// Number of pending WAL entries.
|
||||
pub async fn wal_pending_count(&self) -> usize {
|
||||
self.inner.lock().await.wal_pending_count()
|
||||
}
|
||||
|
||||
/// Get a clone of the config.
|
||||
pub async fn config(&self) -> MemoryConfig {
|
||||
self.inner.lock().await.config().clone()
|
||||
}
|
||||
|
||||
/// Generate AGENTS.md content.
|
||||
pub async fn generate_agents_md(&self) -> String {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mem = inner.blocking_lock();
|
||||
mem.generate_agents_md()
|
||||
})
|
||||
.await
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
// -- Session ops --------------------------------------------------------
|
||||
|
||||
/// Add a session record.
|
||||
pub async fn add_session(
|
||||
&self,
|
||||
id: String,
|
||||
start: usize,
|
||||
end: usize,
|
||||
channel: String,
|
||||
summary: String,
|
||||
) -> Result<()> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mut mem = inner.blocking_lock();
|
||||
mem.add_session(&id, start, end, &channel, &summary)
|
||||
})
|
||||
.await
|
||||
.map_err(join_err)?
|
||||
}
|
||||
|
||||
/// Get a session summary by ID.
|
||||
pub async fn get_session_summary(&self, session_id: String) -> Result<Option<String>> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mem = inner.blocking_lock();
|
||||
mem.get_session_summary(&session_id)
|
||||
})
|
||||
.await
|
||||
.map_err(join_err)?
|
||||
}
|
||||
|
||||
// -- Knowledge graph ops ------------------------------------------------
|
||||
|
||||
/// Add an entity to the knowledge graph.
|
||||
pub async fn add_entity(
|
||||
&self,
|
||||
name: String,
|
||||
entity_type: String,
|
||||
embedding_idx: i64,
|
||||
) -> Result<u64> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mut mem = inner.blocking_lock();
|
||||
mem.add_entity(&name, &entity_type, embedding_idx)
|
||||
})
|
||||
.await
|
||||
.map_err(join_err)?
|
||||
}
|
||||
|
||||
/// Add an entity alias.
|
||||
pub async fn add_entity_alias(&self, alias: String, entity_id: i64) -> Result<()> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mut mem = inner.blocking_lock();
|
||||
mem.add_entity_alias(&alias, entity_id)
|
||||
})
|
||||
.await
|
||||
.map_err(join_err)?
|
||||
}
|
||||
|
||||
/// Add a relation to the knowledge graph.
|
||||
pub async fn add_relation(
|
||||
&self,
|
||||
src: u64,
|
||||
tgt: u64,
|
||||
relation: String,
|
||||
weight: f32,
|
||||
) -> Result<()> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mut mem = inner.blocking_lock();
|
||||
mem.add_relation(src, tgt, &relation, weight)
|
||||
})
|
||||
.await
|
||||
.map_err(join_err)?
|
||||
}
|
||||
|
||||
// -- Snapshot -----------------------------------------------------------
|
||||
|
||||
/// Snapshot the memory file to a destination path.
|
||||
pub async fn snapshot(&self, dest: PathBuf) -> Result<PathBuf> {
|
||||
let inner = Arc::clone(&self.inner);
|
||||
spawn_blocking(move || {
|
||||
let mem = inner.blocking_lock();
|
||||
mem.snapshot(&dest)
|
||||
})
|
||||
.await
|
||||
.map_err(join_err)?
|
||||
}
|
||||
|
||||
// -- Flush / lifecycle --------------------------------------------------
|
||||
|
||||
/// Request an immediate WAL → .h5 flush.
|
||||
pub async fn flush(&self) -> Result<()> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.write_tx
|
||||
.send(WriteCmd::FlushNow(tx))
|
||||
.await
|
||||
.map_err(|_| channel_gone())?;
|
||||
rx.await.map_err(|_| channel_gone())?
|
||||
}
|
||||
|
||||
/// Tick the session (decay activations + flush).
|
||||
pub async fn tick_session(&self) -> Result<()> {
|
||||
let (tx, rx) = oneshot::channel();
|
||||
self.write_tx
|
||||
.send(WriteCmd::TickSession(tx))
|
||||
.await
|
||||
.map_err(|_| channel_gone())?;
|
||||
rx.await.map_err(|_| channel_gone())?
|
||||
}
|
||||
|
||||
/// Gracefully shut down the background writer.
|
||||
///
|
||||
/// Performs a final flush before stopping. Call before drop to
|
||||
/// ensure all buffered data is persisted.
|
||||
pub async fn shutdown(&self) -> Result<()> {
|
||||
self.flush().await?;
|
||||
let (tx, rx) = oneshot::channel();
|
||||
let _ = self.write_tx.send(WriteCmd::Shutdown(tx)).await;
|
||||
let _ = rx.await;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Background writer task
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// The background writer loop. Handles:
|
||||
/// 1. Batched saves from the channel
|
||||
/// 2. Periodic auto-flush on a timer
|
||||
/// 3. Threshold-based flush when WAL grows too large
|
||||
async fn background_writer(
|
||||
inner: Arc<Mutex<HDF5Memory>>,
|
||||
mut rx: mpsc::Receiver<WriteCmd>,
|
||||
config: AsyncConfig,
|
||||
) {
|
||||
let use_interval = config.flush_interval > Duration::ZERO;
|
||||
let use_threshold = config.flush_threshold > 0;
|
||||
|
||||
// Dirty flag: true when cache has unsaved changes that haven't been
|
||||
// flushed to .h5 yet. Saves always go through WAL first, so data
|
||||
// is durable — this just tracks whether we need a full .h5 rewrite.
|
||||
let mut dirty = false;
|
||||
|
||||
let mut interval = tokio::time::interval(if use_interval {
|
||||
config.flush_interval
|
||||
} else {
|
||||
// If disabled, set a very long interval so it never fires
|
||||
Duration::from_secs(86400)
|
||||
});
|
||||
// Don't fire immediately on creation
|
||||
interval.tick().await;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
// --- Channel commands ---
|
||||
cmd = rx.recv() => {
|
||||
match cmd {
|
||||
Some(WriteCmd::Save { entries, reply }) => {
|
||||
let mem = Arc::clone(&inner);
|
||||
let result = spawn_blocking(move || {
|
||||
let mut m = mem.blocking_lock();
|
||||
let mut indices = Vec::with_capacity(entries.len());
|
||||
for entry in entries {
|
||||
// Push to cache + WAL only (no .h5 rewrite).
|
||||
// We use the existing save() which handles
|
||||
// WAL append + auto-merge at wal_max_entries.
|
||||
match m.save(entry) {
|
||||
Ok(idx) => indices.push(idx),
|
||||
Err(e) => return Err(e),
|
||||
}
|
||||
}
|
||||
Ok(indices)
|
||||
})
|
||||
.await
|
||||
.unwrap_or_else(|e| Err(MemoryError::Io(std::io::Error::new(
|
||||
std::io::ErrorKind::Other, e,
|
||||
))));
|
||||
dirty = result.is_ok();
|
||||
let _ = reply.send(result);
|
||||
|
||||
// Check threshold
|
||||
if use_threshold && dirty {
|
||||
let mem = Arc::clone(&inner);
|
||||
let threshold = config.flush_threshold;
|
||||
let pending = spawn_blocking(move || {
|
||||
let m = mem.blocking_lock();
|
||||
m.wal_pending_count()
|
||||
}).await.unwrap_or(0);
|
||||
|
||||
if pending >= threshold {
|
||||
let mem = Arc::clone(&inner);
|
||||
let _ = spawn_blocking(move || {
|
||||
let mut m = mem.blocking_lock();
|
||||
m.flush_wal()
|
||||
}).await;
|
||||
dirty = false;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Some(WriteCmd::FlushNow(reply)) => {
|
||||
if dirty {
|
||||
let mem = Arc::clone(&inner);
|
||||
let result = spawn_blocking(move || {
|
||||
let mut m = mem.blocking_lock();
|
||||
m.flush_wal()
|
||||
})
|
||||
.await
|
||||
.unwrap_or_else(|e| Err(MemoryError::Io(std::io::Error::new(
|
||||
std::io::ErrorKind::Other, e,
|
||||
))));
|
||||
if result.is_ok() { dirty = false; }
|
||||
let _ = reply.send(result);
|
||||
} else {
|
||||
let _ = reply.send(Ok(()));
|
||||
}
|
||||
}
|
||||
|
||||
Some(WriteCmd::TickSession(reply)) => {
|
||||
let mem = Arc::clone(&inner);
|
||||
let result = spawn_blocking(move || {
|
||||
let mut m = mem.blocking_lock();
|
||||
m.tick_session()
|
||||
})
|
||||
.await
|
||||
.unwrap_or_else(|e| Err(MemoryError::Io(std::io::Error::new(
|
||||
std::io::ErrorKind::Other, e,
|
||||
))));
|
||||
if result.is_ok() { dirty = false; }
|
||||
let _ = reply.send(result);
|
||||
}
|
||||
|
||||
Some(WriteCmd::Shutdown(reply)) => {
|
||||
let _ = reply.send(());
|
||||
break;
|
||||
}
|
||||
|
||||
None => break, // channel closed
|
||||
}
|
||||
}
|
||||
|
||||
// --- Periodic auto-flush ---
|
||||
_ = interval.tick(), if use_interval && dirty => {
|
||||
let mem = Arc::clone(&inner);
|
||||
let ok = spawn_blocking(move || {
|
||||
let mut m = mem.blocking_lock();
|
||||
m.flush_wal()
|
||||
}).await;
|
||||
if ok.is_ok() { dirty = false; }
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn join_err(e: tokio::task::JoinError) -> MemoryError {
|
||||
MemoryError::Io(std::io::Error::new(std::io::ErrorKind::Other, e))
|
||||
}
|
||||
|
||||
fn channel_gone() -> MemoryError {
|
||||
MemoryError::Io(std::io::Error::new(
|
||||
std::io::ErrorKind::BrokenPipe,
|
||||
"background writer task gone",
|
||||
))
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Tests
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn make_config(dir: &tempfile::TempDir) -> MemoryConfig {
|
||||
let mut c = MemoryConfig::new(dir.path().join("async_test.h5"), "async-agent", 4);
|
||||
c.wal_enabled = true;
|
||||
c
|
||||
}
|
||||
|
||||
fn fast_async_config() -> AsyncConfig {
|
||||
AsyncConfig {
|
||||
flush_interval: Duration::from_millis(100),
|
||||
flush_threshold: 50,
|
||||
channel_capacity: 64,
|
||||
}
|
||||
}
|
||||
|
||||
fn make_entry(chunk: &str, embedding: &[f32]) -> MemoryEntry {
|
||||
MemoryEntry {
|
||||
chunk: chunk.to_string(),
|
||||
embedding: embedding.to_vec(),
|
||||
source_channel: "test".to_string(),
|
||||
timestamp: 1000000.0,
|
||||
session_id: "session-1".to_string(),
|
||||
tags: "".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn create_and_count() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(mem.count().await, 0);
|
||||
assert_eq!(mem.count_active().await, 0);
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn save_and_search() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let idx = mem
|
||||
.save(make_entry("hello async world", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(idx, 0);
|
||||
assert_eq!(mem.count().await, 1);
|
||||
|
||||
let results = mem
|
||||
.hybrid_search(vec![1.0, 0.0, 0.0, 0.0], String::new(), 1.0, 0.0, 5)
|
||||
.await;
|
||||
assert!(!results.is_empty());
|
||||
assert_eq!(results[0].index, 0);
|
||||
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn save_batch_async() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let entries = vec![
|
||||
make_entry("a", &[1.0, 0.0, 0.0, 0.0]),
|
||||
make_entry("b", &[0.0, 1.0, 0.0, 0.0]),
|
||||
make_entry("c", &[0.0, 0.0, 1.0, 0.0]),
|
||||
];
|
||||
let indices = mem.save_batch(entries).await.unwrap();
|
||||
assert_eq!(indices, vec![0, 1, 2]);
|
||||
assert_eq!(mem.count().await, 3);
|
||||
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn delete_and_compact() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let mut config = make_config(&dir);
|
||||
config.compact_threshold = 0.0;
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mem.save(make_entry("a", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
mem.save(make_entry("b", &[0.0, 1.0, 0.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
mem.save(make_entry("c", &[0.0, 0.0, 1.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mem.delete(1).await.unwrap();
|
||||
assert_eq!(mem.count_active().await, 2);
|
||||
|
||||
let removed = mem.compact().await.unwrap();
|
||||
assert_eq!(removed, 1);
|
||||
assert_eq!(mem.count().await, 2);
|
||||
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn flush_and_reopen() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let path = config.path.clone();
|
||||
|
||||
{
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
mem.save(make_entry("persist me", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
let mem = AsyncHDF5Memory::open_with(&path, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(mem.count().await, 1);
|
||||
let chunk = mem.get_chunk(0).await;
|
||||
assert_eq!(chunk.as_deref(), Some("persist me"));
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn session_tracking_async() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mem.add_session(
|
||||
"s1".into(),
|
||||
0,
|
||||
5,
|
||||
"discord".into(),
|
||||
"talked about rust".into(),
|
||||
)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let summary = mem.get_session_summary("s1".into()).await.unwrap();
|
||||
assert_eq!(summary.as_deref(), Some("talked about rust"));
|
||||
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn knowledge_graph_async() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let id1 = mem
|
||||
.add_entity("Rust".into(), "language".into(), -1)
|
||||
.await
|
||||
.unwrap();
|
||||
let id2 = mem
|
||||
.add_entity("HDF5".into(), "format".into(), -1)
|
||||
.await
|
||||
.unwrap();
|
||||
mem.add_relation(id1, id2, "uses".into(), 1.0)
|
||||
.await
|
||||
.unwrap();
|
||||
mem.add_entity_alias("rustlang".into(), id1 as i64)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn tick_session_async() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mem.save(make_entry("decay test", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mem.tick_session().await.unwrap();
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn snapshot_async() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mem.save(make_entry("snap", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
let snap_dest = dir.path().join("snapshot.h5");
|
||||
let snap_path = mem.snapshot(snap_dest).await.unwrap();
|
||||
assert!(snap_path.exists());
|
||||
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn concurrent_saves() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let mem = Arc::new(
|
||||
AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap(),
|
||||
);
|
||||
|
||||
let mut handles = Vec::new();
|
||||
for i in 0..20 {
|
||||
let m = Arc::clone(&mem);
|
||||
handles.push(tokio::spawn(async move {
|
||||
m.save(make_entry(
|
||||
&format!("concurrent-{i}"),
|
||||
&[i as f32, 0.0, 0.0, 0.0],
|
||||
))
|
||||
.await
|
||||
.unwrap()
|
||||
}));
|
||||
}
|
||||
|
||||
for h in handles {
|
||||
h.await.unwrap();
|
||||
}
|
||||
|
||||
assert_eq!(mem.count().await, 20);
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn periodic_auto_flush() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let path = config.path.clone();
|
||||
|
||||
let async_config = AsyncConfig {
|
||||
flush_interval: Duration::from_millis(50),
|
||||
flush_threshold: 0, // disable threshold
|
||||
channel_capacity: 64,
|
||||
};
|
||||
let mem = AsyncHDF5Memory::create_with(config, async_config)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
mem.save(make_entry("auto-flush", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Wait for the periodic flush to fire
|
||||
tokio::time::sleep(Duration::from_millis(150)).await;
|
||||
|
||||
// Verify data is on disk by reopening without explicit flush
|
||||
drop(mem);
|
||||
let mem2 = AsyncHDF5Memory::open(&path).await.unwrap();
|
||||
assert_eq!(mem2.count().await, 1);
|
||||
mem2.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn threshold_flush() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let mut config = make_config(&dir);
|
||||
config.wal_max_entries = 1000; // high so sync auto-merge doesn't trigger
|
||||
let path = config.path.clone();
|
||||
|
||||
let async_config = AsyncConfig {
|
||||
flush_interval: Duration::ZERO, // disable periodic
|
||||
flush_threshold: 5,
|
||||
channel_capacity: 64,
|
||||
};
|
||||
let mem = AsyncHDF5Memory::create_with(config, async_config)
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Save enough entries to cross the threshold
|
||||
for i in 0..6 {
|
||||
mem.save(make_entry(
|
||||
&format!("thresh-{i}"),
|
||||
&[i as f32, 0.0, 0.0, 0.0],
|
||||
))
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Give background task a moment to process the threshold flush
|
||||
tokio::time::sleep(Duration::from_millis(50)).await;
|
||||
|
||||
assert_eq!(mem.count().await, 6);
|
||||
mem.shutdown().await.unwrap();
|
||||
|
||||
// Verify persistence
|
||||
let mem2 = AsyncHDF5Memory::open(&path).await.unwrap();
|
||||
assert_eq!(mem2.count().await, 6);
|
||||
mem2.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn dirty_flag_skips_noop_flush() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
let mem = AsyncHDF5Memory::create_with(config, fast_async_config())
|
||||
.await
|
||||
.unwrap();
|
||||
|
||||
// Flush with nothing dirty — should be instant no-op
|
||||
mem.flush().await.unwrap();
|
||||
mem.flush().await.unwrap();
|
||||
mem.flush().await.unwrap();
|
||||
|
||||
// Save something, flush, then flush again (second should be no-op)
|
||||
mem.save(make_entry("dirty", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
mem.flush().await.unwrap();
|
||||
mem.flush().await.unwrap(); // no-op
|
||||
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn default_config_works() {
|
||||
let dir = tempfile::TempDir::new().unwrap();
|
||||
let config = make_config(&dir);
|
||||
// Use default AsyncConfig (no _with variant)
|
||||
let mem = AsyncHDF5Memory::create(config).await.unwrap();
|
||||
mem.save(make_entry("default", &[1.0, 0.0, 0.0, 0.0]))
|
||||
.await
|
||||
.unwrap();
|
||||
assert_eq!(mem.count().await, 1);
|
||||
mem.shutdown().await.unwrap();
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user