use async_trait::async_trait;
use std::fmt::Debug;
use std::sync::Arc;
use std::time::Duration;
use tracing::{debug, info, warn};
use crate::error::StoreError;
use crate::traits::VectorStore;
use crate::types::StoreStats;
#[derive(Debug, Clone)]
pub struct IndexThrottleConfig {
pub cooldown: Duration,
pub count_threshold: usize,
}
impl Default for IndexThrottleConfig {
fn default() -> Self {
Self {
cooldown: Duration::from_secs(60),
count_threshold: 1000,
}
}
}
#[derive(Debug)]
pub struct IndexBuilder {
store: Arc<dyn VectorStore>,
config: IndexThrottleConfig,
}
impl IndexBuilder {
pub fn new(store: Arc<dyn VectorStore>, config: IndexThrottleConfig) -> Self {
Self { store, config }
}
pub async fn rebuild(&self) -> Result<(), StoreError> {
info!(target: "xz_embed", "starting index rebuild");
self.store.rebuild_index().await?;
info!(target: "xz_embed", "index rebuild completed");
Ok(())
}
pub async fn rebuild_if_needed(&self) -> Result<(), StoreError> {
let count = self.store.count().await?;
if count >= self.config.count_threshold {
self.rebuild().await?;
} else {
debug!(target: "xz_embed", count, threshold = self.config.count_threshold, "skipping rebuild");
}
Ok(())
}
pub async fn start_background_rebuild(self: Arc<Self>) {
tokio::spawn(async move {
loop {
tokio::time::sleep(self.config.cooldown).await;
match self.rebuild_if_needed().await {
Ok(_) => {}
Err(e) => warn!(target: "xz_embed", error = %e, "background rebuild failed"),
}
}
});
}
}