Federation foundation + HNSW trial system + Postgres streaming + PRD reframe
Four shipped features and a PRD realignment, all measured end-to-end:
HNSW trial system (Phase 15 horizon item → complete)
- vectord: EmbeddingCache, harness (eval sets + brute-force ground truth),
TrialJournal, parameterized HnswConfig on build_index_with_config
- /vectors/hnsw/trial, /hnsw/trials/{idx}, /hnsw/trials/{idx}/best,
/hnsw/evals/{name}/autogen, /hnsw/cache/stats
- Measured on resumes_100k_v2 (100K × 768d): brute-force 44ms -> HNSW 873us
at 100% recall@10. ec=80 es=30 locked as HnswConfig::default()
- Lower ec values trade recall for build time: 20/30 = 0.96 recall in 8s,
80/30 = 1.00 recall in 230s
Catalog manifest repair
- catalogd: resync_from_parquet reads parquet footers to restore row_count
and columns on drifted manifests
- POST /catalog/datasets/{name}/resync + POST /catalog/resync-missing
- All 7 staffing tables recovered to PRD-matching 2,469,278 rows
Federation foundation (ADR-017)
- shared::secrets: SecretsProvider trait + FileSecretsProvider (reads
/etc/lakehouse/secrets.toml, enforces 0600 perms)
- storaged::registry::BucketRegistry — multi-bucket resolution with
rescue_bucket read fallback and reachability probing
- storaged::error_journal — bucket op failures visible in one HTTP call
- storaged::append_log — write-once batched append pattern (fixes the RMW
anti-pattern llms3.com calls out; errors and trial journals both use it)
- /storage/buckets, /storage/errors, /storage/bucket-health,
/storage/errors/{flush,compact}
- Bucket-aware I/O at /storage/buckets/{bucket}/objects/{*key} with
X-Lakehouse-Rescue-Used observability headers on fallback
Postgres streaming ingest
- ingestd::pg_stream: DSN parser, batched ORDER BY + LIMIT/OFFSET pagination
into ArrowWriter, lineage redacts password
- POST /ingest/db — verified against live knowledge_base.team_runs
(586 rows × 13 cols, 6 batches, 196ms end-to-end)
PRD realignment (2026-04-16)
- Dual use case: staffing analytics + local LLM knowledge substrate
- Removed "multi-tenancy (single-owner system)" from non-goals
- Added invariants 8-11: indexes hot-swappable, per-reader profiles,
trials-as-data, operational failures findable in one HTTP call
- New phases 16 (hot-swap generations), 17 (model profiles + dataset
bindings), 18 (Lance vs Parquet+sidecar evaluation)
- Known ceilings table documents the 5M vector wall and escape hatches
- ADR-017 (federation), ADR-018 (append-log pattern) added
- EXECUTION_PLAN.md sequences phases B-E with success gates and
decision rules
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
root
@@ -15,5 +15,6 @@ bytes = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
arrow = { workspace = true }
|
||||
proto = { path = "../proto" }
|
||||
tonic = { workspace = true }
|
||||
|
||||
@@ -140,6 +140,108 @@ impl Registry {
|
||||
datasets.values().cloned().collect()
|
||||
}
|
||||
|
||||
/// Re-read the parquet footer(s) for a dataset and repopulate `row_count`
|
||||
/// and `columns` from reality. Use this to repair manifests whose
|
||||
/// metadata was lost (e.g. migrated from a pre-Phase 10 catalog).
|
||||
///
|
||||
/// Does NOT touch owner/description/sensitivity/lineage/tags — only
|
||||
/// the structural facts that parquet can tell us authoritatively.
|
||||
/// The existing `schema_fingerprint` is updated if the recomputed one
|
||||
/// differs; a warning is logged so drift is visible.
|
||||
pub async fn resync_from_parquet(&self, name: &str) -> Result<DatasetManifest, String> {
|
||||
use shared::arrow_helpers::{fingerprint_schema, parquet_to_record_batches};
|
||||
|
||||
// Snapshot the target manifest so we don't hold the write lock during IO.
|
||||
let (id, objects, old_fp) = {
|
||||
let datasets = self.datasets.read().await;
|
||||
let m = datasets
|
||||
.values()
|
||||
.find(|d| d.name == name)
|
||||
.ok_or_else(|| format!("dataset not found: {name}"))?;
|
||||
(m.id.clone(), m.objects.clone(), m.schema_fingerprint.clone())
|
||||
};
|
||||
|
||||
if objects.is_empty() {
|
||||
return Err(format!("dataset '{name}' has no object references to resync from"));
|
||||
}
|
||||
|
||||
let mut total_rows: u64 = 0;
|
||||
let mut first_schema: Option<arrow::datatypes::SchemaRef> = None;
|
||||
for obj in &objects {
|
||||
let data = ops::get(&self.store, &obj.key).await
|
||||
.map_err(|e| format!("read {}: {e}", obj.key))?;
|
||||
let (schema, batches) = parquet_to_record_batches(&data)
|
||||
.map_err(|e| format!("parse {}: {e}", obj.key))?;
|
||||
let rows: u64 = batches.iter().map(|b| b.num_rows() as u64).sum();
|
||||
total_rows += rows;
|
||||
if first_schema.is_none() {
|
||||
first_schema = Some(schema);
|
||||
}
|
||||
}
|
||||
|
||||
let schema = first_schema.ok_or("no schema recovered")?;
|
||||
let new_fp = fingerprint_schema(&schema);
|
||||
if new_fp != old_fp {
|
||||
tracing::warn!(
|
||||
"dataset '{}' schema fingerprint drift: {} -> {} (updating to match parquet reality)",
|
||||
name, old_fp.0, new_fp.0,
|
||||
);
|
||||
}
|
||||
|
||||
let columns: Vec<ColumnMeta> = schema
|
||||
.fields()
|
||||
.iter()
|
||||
.map(|f| ColumnMeta {
|
||||
name: f.name().clone(),
|
||||
data_type: f.data_type().to_string(),
|
||||
sensitivity: None,
|
||||
description: String::new(),
|
||||
is_pii: false,
|
||||
})
|
||||
.collect();
|
||||
|
||||
// Apply updates.
|
||||
let mut datasets = self.datasets.write().await;
|
||||
let manifest = datasets
|
||||
.get_mut(&id)
|
||||
.ok_or_else(|| format!("dataset disappeared during resync: {name}"))?;
|
||||
manifest.row_count = Some(total_rows);
|
||||
manifest.columns = columns;
|
||||
manifest.schema_fingerprint = new_fp;
|
||||
manifest.updated_at = chrono::Utc::now();
|
||||
|
||||
// Persist.
|
||||
let manifest_key = format!("{MANIFEST_PREFIX}/{}.json", manifest.id);
|
||||
let json = serde_json::to_vec_pretty(manifest).map_err(|e| e.to_string())?;
|
||||
ops::put(&self.store, &manifest_key, json.into()).await?;
|
||||
|
||||
tracing::info!("resynced '{name}': row_count={total_rows}, {} columns", manifest.columns.len());
|
||||
Ok(manifest.clone())
|
||||
}
|
||||
|
||||
/// Resync every dataset that currently has a null row_count.
|
||||
/// Returns (successes, failures) where each entry is (name, detail).
|
||||
pub async fn resync_missing(&self) -> (Vec<(String, u64)>, Vec<(String, String)>) {
|
||||
let names: Vec<String> = {
|
||||
let datasets = self.datasets.read().await;
|
||||
datasets
|
||||
.values()
|
||||
.filter(|d| d.row_count.is_none() || d.columns.is_empty())
|
||||
.map(|d| d.name.clone())
|
||||
.collect()
|
||||
};
|
||||
|
||||
let mut ok = Vec::new();
|
||||
let mut err = Vec::new();
|
||||
for name in names {
|
||||
match self.resync_from_parquet(&name).await {
|
||||
Ok(m) => ok.push((name, m.row_count.unwrap_or(0))),
|
||||
Err(e) => err.push((name, e)),
|
||||
}
|
||||
}
|
||||
(ok, err)
|
||||
}
|
||||
|
||||
/// Add objects to an existing dataset.
|
||||
pub async fn add_objects(
|
||||
&self,
|
||||
|
||||
@@ -19,6 +19,8 @@ pub fn router(registry: Registry) -> Router {
|
||||
.route("/datasets/{id}", get(get_dataset))
|
||||
.route("/datasets/by-name/{name}", get(get_dataset_by_name))
|
||||
.route("/datasets/by-name/{name}/metadata", post(update_metadata))
|
||||
.route("/datasets/by-name/{name}/resync", post(resync_dataset))
|
||||
.route("/resync-missing", post(resync_all_missing))
|
||||
.with_state(registry)
|
||||
}
|
||||
|
||||
@@ -152,3 +154,43 @@ async fn update_metadata(
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
/// Re-read parquet footers for a single dataset and repopulate row_count
|
||||
/// and columns from reality. Useful for repairing manifests whose metadata
|
||||
/// was lost or never backfilled.
|
||||
async fn resync_dataset(
|
||||
State(registry): State<Registry>,
|
||||
Path(name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match registry.resync_from_parquet(&name).await {
|
||||
Ok(manifest) => Ok(Json(DatasetResponse::from(&manifest))),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ResyncAllResponse {
|
||||
succeeded: Vec<ResyncOk>,
|
||||
failed: Vec<ResyncErr>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ResyncOk {
|
||||
name: String,
|
||||
row_count: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ResyncErr {
|
||||
name: String,
|
||||
error: String,
|
||||
}
|
||||
|
||||
/// Resync every dataset that currently has null row_count or empty columns.
|
||||
async fn resync_all_missing(State(registry): State<Registry>) -> impl IntoResponse {
|
||||
let (ok, err) = registry.resync_missing().await;
|
||||
Json(ResyncAllResponse {
|
||||
succeeded: ok.into_iter().map(|(name, row_count)| ResyncOk { name, row_count }).collect(),
|
||||
failed: err.into_iter().map(|(name, error)| ResyncErr { name, error }).collect(),
|
||||
})
|
||||
}
|
||||
|
||||
@@ -24,8 +24,30 @@ async fn main() {
|
||||
tracing::info!("config loaded: gateway={}:{}, storage={}",
|
||||
config.gateway.host, config.gateway.port, config.storage.root);
|
||||
|
||||
// Storage backend
|
||||
let store = storaged::backend::init_local(&config.storage.root);
|
||||
// Secrets provider
|
||||
let secrets = std::sync::Arc::new(shared::secrets::FileSecretsProvider::from_env());
|
||||
if let Err(e) = secrets.load().await {
|
||||
tracing::warn!("secrets load: {e} — bucket credentials may be missing");
|
||||
}
|
||||
|
||||
// Federation: bucket registry (primary + any [[storage.buckets]] entries + rescue)
|
||||
let bucket_registry = std::sync::Arc::new(
|
||||
storaged::registry::BucketRegistry::from_config(&config.storage, secrets.clone())
|
||||
.await
|
||||
.expect("bucket registry init failed")
|
||||
);
|
||||
let buckets = bucket_registry.list().await;
|
||||
tracing::info!(
|
||||
"bucket registry: {} buckets [{}], rescue={:?}",
|
||||
buckets.len(),
|
||||
buckets.iter().map(|b| format!("{}:{}", b.name, b.backend)).collect::<Vec<_>>().join(", "),
|
||||
bucket_registry.rescue_name(),
|
||||
);
|
||||
|
||||
// Back-compat: existing components that expect a single Arc<ObjectStore>
|
||||
// get the primary bucket's store. As we migrate call sites to be bucket-
|
||||
// aware, they'll call bucket_registry.get(name) directly.
|
||||
let store = bucket_registry.default_store();
|
||||
|
||||
// Catalog
|
||||
let registry = catalogd::registry::Registry::new(store.clone());
|
||||
@@ -56,7 +78,8 @@ async fn main() {
|
||||
// HTTP router
|
||||
let mut app = Router::new()
|
||||
.route("/health", get(health))
|
||||
.nest("/storage", storaged::service::router(store.clone()))
|
||||
.nest("/storage", storaged::service::router(store.clone())
|
||||
.merge(storaged::federation_service::router(bucket_registry.clone())))
|
||||
.nest("/catalog", catalogd::service::router(registry.clone()))
|
||||
.nest("/query", queryd::service::router(engine.clone()))
|
||||
.nest("/ai", aibridge::service::router(ai_client.clone()))
|
||||
@@ -73,6 +96,8 @@ async fn main() {
|
||||
job_tracker: vectord::jobs::JobTracker::new(),
|
||||
index_registry: index_reg,
|
||||
hnsw_store: vectord::hnsw::HnswStore::new(),
|
||||
embedding_cache: vectord::embedding_cache::EmbeddingCache::new(store.clone()),
|
||||
trial_journal: vectord::trial::TrialJournal::new(store.clone()),
|
||||
}
|
||||
}))
|
||||
.nest("/workspaces", queryd::workspace_service::router(workspace_mgr))
|
||||
|
||||
@@ -21,3 +21,4 @@ csv = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
tokio-postgres = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod db_ingest;
|
||||
pub mod pg_stream;
|
||||
pub mod detect;
|
||||
pub mod csv_ingest;
|
||||
pub mod json_ingest;
|
||||
|
||||
@@ -0,0 +1,308 @@
|
||||
/// Streaming PostgreSQL ingest.
|
||||
///
|
||||
/// The original `db_ingest::import_postgres_table` loads every row into
|
||||
/// memory before emitting Parquet — fine for small tables, blows up on
|
||||
/// 1M+ rows. This module paginates via `ORDER BY <pk> LIMIT N OFFSET M`
|
||||
/// and streams batches into an `ArrowWriter`, closing once exhausted.
|
||||
///
|
||||
/// Pagination is OFFSET-based (not keyset) — simpler, works on any table
|
||||
/// without needing to know the PK type. OFFSET N scans N rows per fetch,
|
||||
/// which becomes quadratic for very large tables. Upgrade to keyset when
|
||||
/// a user actually ingests something that hurts.
|
||||
|
||||
use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array, Int64Array, RecordBatch, StringArray};
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use parquet::arrow::ArrowWriter;
|
||||
use std::sync::Arc;
|
||||
use tokio_postgres::{Client, NoTls, types::Type as PgType};
|
||||
|
||||
use crate::db_ingest::DbConfig;
|
||||
|
||||
/// Request shape for streaming ingest — takes a DSN string and optional
|
||||
/// tuning knobs.
|
||||
#[derive(Debug, Clone, serde::Deserialize)]
|
||||
pub struct PgStreamRequest {
|
||||
/// postgresql://user:pass@host:port/db — alternative to DbConfig.
|
||||
pub dsn: String,
|
||||
pub table: String,
|
||||
#[serde(default)]
|
||||
pub dataset_name: Option<String>,
|
||||
/// Rows per fetch. Default 10_000.
|
||||
#[serde(default)]
|
||||
pub batch_size: Option<usize>,
|
||||
/// Column to ORDER BY for stable pagination. If omitted, the first
|
||||
/// column returned by the schema probe is used.
|
||||
#[serde(default)]
|
||||
pub order_by: Option<String>,
|
||||
/// Hard cap on total rows (for sampling / previews).
|
||||
#[serde(default)]
|
||||
pub limit: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct PgStreamResult {
|
||||
pub table: String,
|
||||
pub rows: usize,
|
||||
pub batches: usize,
|
||||
pub columns: usize,
|
||||
pub schema: Vec<String>,
|
||||
pub parquet_bytes: u64,
|
||||
pub duration_secs: f32,
|
||||
}
|
||||
|
||||
/// Parse a postgresql:// DSN into a DbConfig.
|
||||
/// Supports: postgresql://[user[:password]@]host[:port][/db]
|
||||
/// Does NOT support: query parameters (sslmode etc) — add if needed.
|
||||
pub fn parse_dsn(dsn: &str) -> Result<DbConfig, String> {
|
||||
let rest = dsn
|
||||
.strip_prefix("postgresql://")
|
||||
.or_else(|| dsn.strip_prefix("postgres://"))
|
||||
.ok_or_else(|| "DSN must start with postgresql:// or postgres://".to_string())?;
|
||||
|
||||
// Split off path (database) first.
|
||||
let (auth_host, database) = match rest.split_once('/') {
|
||||
Some((ah, db)) => (ah, db.split('?').next().unwrap_or(db).to_string()),
|
||||
None => (rest, "postgres".to_string()),
|
||||
};
|
||||
|
||||
// Split user[:password] from host[:port].
|
||||
let (userpass, hostport) = match auth_host.rsplit_once('@') {
|
||||
Some((up, hp)) => (Some(up), hp),
|
||||
None => (None, auth_host),
|
||||
};
|
||||
|
||||
let (user, password) = match userpass {
|
||||
Some(up) => match up.split_once(':') {
|
||||
Some((u, p)) => (u.to_string(), p.to_string()),
|
||||
None => (up.to_string(), String::new()),
|
||||
},
|
||||
None => ("postgres".to_string(), String::new()),
|
||||
};
|
||||
|
||||
let (host, port) = match hostport.rsplit_once(':') {
|
||||
Some((h, p)) => (
|
||||
h.to_string(),
|
||||
p.parse::<u16>().map_err(|_| format!("invalid port in DSN: {p}"))?,
|
||||
),
|
||||
None => (hostport.to_string(), 5432),
|
||||
};
|
||||
|
||||
if host.is_empty() {
|
||||
return Err("DSN has no host".into());
|
||||
}
|
||||
|
||||
Ok(DbConfig { host, port, database, user, password })
|
||||
}
|
||||
|
||||
/// Stream a Postgres table as Parquet bytes.
|
||||
///
|
||||
/// Returns the full Parquet payload + summary stats. The payload is
|
||||
/// written to an in-memory buffer via `ArrowWriter` so each batch gets
|
||||
/// columnar compression as it arrives — memory footprint stays roughly
|
||||
/// at one batch plus Parquet footer state.
|
||||
pub async fn stream_table_to_parquet(
|
||||
req: &PgStreamRequest,
|
||||
) -> Result<(bytes::Bytes, PgStreamResult), String> {
|
||||
let t0 = std::time::Instant::now();
|
||||
let config = parse_dsn(&req.dsn)?;
|
||||
let batch_size = req.batch_size.unwrap_or(10_000).max(1);
|
||||
|
||||
let (client, connection) = tokio_postgres::connect(&config.connection_string(), NoTls)
|
||||
.await
|
||||
.map_err(|e| format!("postgres connect: {e}"))?;
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = connection.await {
|
||||
tracing::error!("pg connection: {e}");
|
||||
}
|
||||
});
|
||||
|
||||
// Probe schema.
|
||||
let columns = probe_columns(&client, &req.table).await?;
|
||||
if columns.is_empty() {
|
||||
return Err(format!("table '{}' not found or has no columns", req.table));
|
||||
}
|
||||
|
||||
let arrow_fields: Vec<Field> = columns
|
||||
.iter()
|
||||
.map(|(name, pg)| Field::new(name, pg_type_to_arrow(pg), true))
|
||||
.collect();
|
||||
let schema = Arc::new(Schema::new(arrow_fields));
|
||||
let schema_report: Vec<String> = columns
|
||||
.iter()
|
||||
.map(|(n, t)| format!("{}:{}", n, t))
|
||||
.collect();
|
||||
|
||||
// Pagination key: user-specified or first column.
|
||||
let order_col = req.order_by.clone().unwrap_or_else(|| columns[0].0.clone());
|
||||
|
||||
// Stream batches into ArrowWriter.
|
||||
let mut buf: Vec<u8> = Vec::with_capacity(1024 * 1024);
|
||||
let mut writer = ArrowWriter::try_new(&mut buf, schema.clone(), None)
|
||||
.map_err(|e| format!("ArrowWriter init: {e}"))?;
|
||||
|
||||
let mut total_rows: usize = 0;
|
||||
let mut batch_count: usize = 0;
|
||||
let row_cap = req.limit.unwrap_or(usize::MAX);
|
||||
|
||||
loop {
|
||||
let remaining = row_cap.saturating_sub(total_rows);
|
||||
if remaining == 0 { break; }
|
||||
let fetch = remaining.min(batch_size);
|
||||
|
||||
let sql = format!(
|
||||
"SELECT * FROM \"{}\" ORDER BY \"{}\" LIMIT {} OFFSET {}",
|
||||
req.table, order_col, fetch, total_rows,
|
||||
);
|
||||
|
||||
let rows = client
|
||||
.query(&sql, &[])
|
||||
.await
|
||||
.map_err(|e| format!("fetch batch at offset {total_rows}: {e}"))?;
|
||||
|
||||
if rows.is_empty() { break; }
|
||||
|
||||
let n = rows.len();
|
||||
let arrays: Vec<ArrayRef> = columns
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(idx, (_, pg))| rows_to_column(&rows, idx, pg))
|
||||
.collect::<Result<_, _>>()?;
|
||||
let batch = RecordBatch::try_new(schema.clone(), arrays)
|
||||
.map_err(|e| format!("RecordBatch: {e}"))?;
|
||||
writer.write(&batch).map_err(|e| format!("ArrowWriter::write: {e}"))?;
|
||||
|
||||
total_rows += n;
|
||||
batch_count += 1;
|
||||
tracing::info!(
|
||||
"pg stream '{}': fetched batch {} ({} rows, total {})",
|
||||
req.table, batch_count, n, total_rows,
|
||||
);
|
||||
|
||||
if n < fetch { break; } // short read = end of table
|
||||
}
|
||||
|
||||
writer.close().map_err(|e| format!("ArrowWriter::close: {e}"))?;
|
||||
let parquet_bytes = buf.len() as u64;
|
||||
let duration = t0.elapsed().as_secs_f32();
|
||||
|
||||
let result = PgStreamResult {
|
||||
table: req.table.clone(),
|
||||
rows: total_rows,
|
||||
batches: batch_count,
|
||||
columns: columns.len(),
|
||||
schema: schema_report,
|
||||
parquet_bytes,
|
||||
duration_secs: duration,
|
||||
};
|
||||
|
||||
Ok((bytes::Bytes::from(buf), result))
|
||||
}
|
||||
|
||||
async fn probe_columns(client: &Client, table: &str) -> Result<Vec<(String, PgType)>, String> {
|
||||
let stmt = client
|
||||
.prepare(&format!("SELECT * FROM \"{}\" LIMIT 0", table))
|
||||
.await
|
||||
.map_err(|e| format!("prepare schema probe: {e}"))?;
|
||||
Ok(stmt
|
||||
.columns()
|
||||
.iter()
|
||||
.map(|c| (c.name().to_string(), c.type_().clone()))
|
||||
.collect())
|
||||
}
|
||||
|
||||
fn pg_type_to_arrow(pg: &PgType) -> DataType {
|
||||
match *pg {
|
||||
PgType::BOOL => DataType::Boolean,
|
||||
PgType::INT2 | PgType::INT4 => DataType::Int32,
|
||||
PgType::INT8 | PgType::OID => DataType::Int64,
|
||||
PgType::FLOAT4 | PgType::FLOAT8 | PgType::NUMERIC => DataType::Float64,
|
||||
_ => DataType::Utf8,
|
||||
}
|
||||
}
|
||||
|
||||
fn rows_to_column(
|
||||
rows: &[tokio_postgres::Row],
|
||||
idx: usize,
|
||||
pg: &PgType,
|
||||
) -> Result<ArrayRef, String> {
|
||||
match *pg {
|
||||
PgType::BOOL => {
|
||||
let v: Vec<Option<bool>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||
Ok(Arc::new(BooleanArray::from(v)))
|
||||
}
|
||||
PgType::INT2 => {
|
||||
let v: Vec<Option<i32>> = rows.iter()
|
||||
.map(|r| r.try_get::<_, i16>(idx).ok().map(|x| x as i32)).collect();
|
||||
Ok(Arc::new(Int32Array::from(v)))
|
||||
}
|
||||
PgType::INT4 => {
|
||||
let v: Vec<Option<i32>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||
Ok(Arc::new(Int32Array::from(v)))
|
||||
}
|
||||
PgType::INT8 | PgType::OID => {
|
||||
let v: Vec<Option<i64>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||
Ok(Arc::new(Int64Array::from(v)))
|
||||
}
|
||||
PgType::FLOAT4 => {
|
||||
let v: Vec<Option<f64>> = rows.iter()
|
||||
.map(|r| r.try_get::<_, f32>(idx).ok().map(|x| x as f64)).collect();
|
||||
Ok(Arc::new(Float64Array::from(v)))
|
||||
}
|
||||
PgType::FLOAT8 => {
|
||||
let v: Vec<Option<f64>> = rows.iter().map(|r| r.try_get(idx).ok()).collect();
|
||||
Ok(Arc::new(Float64Array::from(v)))
|
||||
}
|
||||
_ => {
|
||||
// Safe-default per ADR-010: everything else becomes string.
|
||||
let v: Vec<Option<String>> = rows.iter().map(|r| {
|
||||
r.try_get::<_, String>(idx).ok()
|
||||
.or_else(|| r.try_get::<_, serde_json::Value>(idx).ok().map(|j| j.to_string()))
|
||||
.or_else(|| {
|
||||
// Timestamps + uuid, rendered via Display.
|
||||
r.try_get::<_, chrono::DateTime<chrono::Utc>>(idx).ok()
|
||||
.map(|t| t.to_rfc3339())
|
||||
})
|
||||
.or_else(|| r.try_get::<_, uuid::Uuid>(idx).ok().map(|u| u.to_string()))
|
||||
.or(Some(String::new()))
|
||||
}).collect();
|
||||
Ok(Arc::new(StringArray::from(v)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_full() {
|
||||
let c = parse_dsn("postgresql://daisy:secret@db.example.com:6543/my_db").unwrap();
|
||||
assert_eq!(c.host, "db.example.com");
|
||||
assert_eq!(c.port, 6543);
|
||||
assert_eq!(c.user, "daisy");
|
||||
assert_eq!(c.password, "secret");
|
||||
assert_eq!(c.database, "my_db");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_minimal() {
|
||||
let c = parse_dsn("postgresql://localhost/knowledge_base").unwrap();
|
||||
assert_eq!(c.host, "localhost");
|
||||
assert_eq!(c.port, 5432);
|
||||
assert_eq!(c.user, "postgres");
|
||||
assert_eq!(c.password, "");
|
||||
assert_eq!(c.database, "knowledge_base");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_no_password() {
|
||||
let c = parse_dsn("postgres://postgres@127.0.0.1:5432/mydb").unwrap();
|
||||
assert_eq!(c.user, "postgres");
|
||||
assert_eq!(c.password, "");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn parse_dsn_rejects_non_pg() {
|
||||
assert!(parse_dsn("http://host/db").is_err());
|
||||
}
|
||||
}
|
||||
@@ -11,7 +11,7 @@ use serde::Deserialize;
|
||||
use std::sync::Arc;
|
||||
|
||||
use catalogd::registry::Registry;
|
||||
use crate::{db_ingest, pipeline};
|
||||
use crate::{db_ingest, pg_stream, pipeline};
|
||||
use shared::arrow_helpers::record_batch_to_parquet;
|
||||
use shared::types::{ObjectRef, SchemaFingerprint};
|
||||
use storaged::ops;
|
||||
@@ -28,6 +28,7 @@ pub fn router(state: IngestState) -> Router {
|
||||
.route("/file", post(ingest_file))
|
||||
.route("/postgres/tables", post(list_pg_tables))
|
||||
.route("/postgres/import", post(import_pg_table))
|
||||
.route("/db", post(ingest_db_stream))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
@@ -188,3 +189,112 @@ async fn import_pg_table(
|
||||
"size_bytes": parquet_size,
|
||||
}))))
|
||||
}
|
||||
|
||||
/// Streaming DSN-based ingest. Paginates via ORDER BY + LIMIT/OFFSET so
|
||||
/// large tables don't blow up memory. This is the path the task spec calls
|
||||
/// `POST /ingest/db` with `{dsn, table, dataset_name}`.
|
||||
async fn ingest_db_stream(
|
||||
State(state): State<IngestState>,
|
||||
Json(req): Json<pg_stream::PgStreamRequest>,
|
||||
) -> Result<(StatusCode, Json<serde_json::Value>), (StatusCode, String)> {
|
||||
tracing::info!(
|
||||
"pg stream ingest: table='{}' dataset='{:?}' batch_size={:?}",
|
||||
req.table, req.dataset_name, req.batch_size,
|
||||
);
|
||||
|
||||
// Stream from Postgres into Parquet bytes.
|
||||
let (parquet, stream_result) = pg_stream::stream_table_to_parquet(&req)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::BAD_GATEWAY, e))?;
|
||||
|
||||
if stream_result.rows == 0 {
|
||||
return Ok((StatusCode::OK, Json(serde_json::json!({
|
||||
"table": req.table,
|
||||
"rows": 0,
|
||||
"message": "table is empty",
|
||||
}))));
|
||||
}
|
||||
|
||||
// Recover schema from the parquet footer — keeps a single source of truth
|
||||
// for the schema fingerprint + column metadata.
|
||||
let (schema, _) = shared::arrow_helpers::parquet_to_record_batches(&parquet)
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("reparse parquet: {e}")))?;
|
||||
|
||||
let dataset_name = req.dataset_name.clone().unwrap_or_else(|| req.table.clone());
|
||||
let storage_key = format!("datasets/{}.parquet", dataset_name);
|
||||
let size_bytes = parquet.len() as u64;
|
||||
|
||||
ops::put(&state.store, &storage_key, parquet)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||
|
||||
let schema_fp = shared::arrow_helpers::fingerprint_schema(&schema);
|
||||
let now = chrono::Utc::now();
|
||||
state.registry.register(
|
||||
dataset_name.clone(),
|
||||
SchemaFingerprint(schema_fp.0),
|
||||
vec![ObjectRef {
|
||||
bucket: "primary".to_string(),
|
||||
key: storage_key.clone(),
|
||||
size_bytes,
|
||||
created_at: now,
|
||||
}],
|
||||
).await.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e))?;
|
||||
|
||||
// Rich metadata: PII auto-detect, lineage, row_count, columns.
|
||||
let col_names: Vec<&str> = schema.fields().iter().map(|f| f.name().as_str()).collect();
|
||||
let sensitivity = shared::pii::detect_dataset_sensitivity(&col_names);
|
||||
let columns: Vec<shared::types::ColumnMeta> = schema.fields().iter().map(|f| {
|
||||
let sens = shared::pii::detect_sensitivity(f.name());
|
||||
shared::types::ColumnMeta {
|
||||
name: f.name().clone(),
|
||||
data_type: f.data_type().to_string(),
|
||||
sensitivity: sens.clone(),
|
||||
description: String::new(),
|
||||
is_pii: matches!(sens, Some(shared::types::Sensitivity::Pii)),
|
||||
}
|
||||
}).collect();
|
||||
|
||||
let lineage = shared::types::Lineage {
|
||||
source_system: "postgresql".to_string(),
|
||||
source_file: format!("dsn: {}", redact_dsn(&req.dsn)),
|
||||
ingest_job: format!("pg-stream-{}", now.timestamp_millis()),
|
||||
ingest_timestamp: now,
|
||||
parent_datasets: vec![],
|
||||
};
|
||||
|
||||
let _ = state.registry.update_metadata(&dataset_name, catalogd::registry::MetadataUpdate {
|
||||
sensitivity,
|
||||
columns: Some(columns),
|
||||
lineage: Some(lineage),
|
||||
row_count: Some(stream_result.rows as u64),
|
||||
..Default::default()
|
||||
}).await;
|
||||
|
||||
Ok((StatusCode::CREATED, Json(serde_json::json!({
|
||||
"dataset_name": dataset_name,
|
||||
"table": stream_result.table,
|
||||
"rows": stream_result.rows,
|
||||
"batches": stream_result.batches,
|
||||
"columns": stream_result.columns,
|
||||
"schema": stream_result.schema,
|
||||
"storage_key": storage_key,
|
||||
"size_bytes": size_bytes,
|
||||
"duration_secs": stream_result.duration_secs,
|
||||
}))))
|
||||
}
|
||||
|
||||
/// Redact the password in a postgresql:// DSN for logging / lineage.
|
||||
fn redact_dsn(dsn: &str) -> String {
|
||||
// Simple approach: find @...: and replace password between : and @.
|
||||
if let Some(at_idx) = dsn.rfind('@') {
|
||||
let prefix = &dsn[..at_idx];
|
||||
if let Some(colon_idx) = prefix.rfind(':') {
|
||||
// But only if the colon is after the scheme colons (postgresql://)
|
||||
if colon_idx > "postgresql://".len() {
|
||||
return format!("{}:***{}", &dsn[..colon_idx], &dsn[at_idx..]);
|
||||
}
|
||||
}
|
||||
}
|
||||
dsn.to_string()
|
||||
}
|
||||
|
||||
@@ -15,3 +15,5 @@ bytes = { workspace = true }
|
||||
sha2 = { workspace = true }
|
||||
toml = { workspace = true }
|
||||
tracing = { workspace = true }
|
||||
tokio = { workspace = true }
|
||||
async-trait = "0.1"
|
||||
|
||||
@@ -28,8 +28,39 @@ pub struct GatewayConfig {
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct StorageConfig {
|
||||
/// Legacy single-backend root. If `buckets` is empty, this is used to
|
||||
/// create an implicit `primary` bucket at this path — preserves the
|
||||
/// pre-federation config shape.
|
||||
#[serde(default = "default_storage_root")]
|
||||
pub root: String,
|
||||
|
||||
/// Where profile buckets are rooted when auto-provisioned.
|
||||
#[serde(default = "default_profile_root")]
|
||||
pub profile_root: String,
|
||||
|
||||
/// Name of the bucket used for read fallback when a target bucket is
|
||||
/// unreachable. If `None`, no fallback — reads fail hard.
|
||||
#[serde(default)]
|
||||
pub rescue_bucket: Option<String>,
|
||||
|
||||
/// Explicitly configured buckets. Empty = backward-compat single-bucket
|
||||
/// mode driven by `root`.
|
||||
#[serde(default)]
|
||||
pub buckets: Vec<BucketConfig>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct BucketConfig {
|
||||
pub name: String,
|
||||
pub backend: String, // "local" | "s3"
|
||||
/// Local filesystem root (for backend = "local")
|
||||
pub root: Option<String>,
|
||||
/// S3 bucket name (for backend = "s3")
|
||||
pub bucket: Option<String>,
|
||||
pub region: Option<String>,
|
||||
pub endpoint: Option<String>,
|
||||
/// Handle for the secrets provider — never the literal credential.
|
||||
pub secret_ref: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Default)]
|
||||
@@ -78,6 +109,7 @@ pub struct ObservabilityConfig {
|
||||
fn default_host() -> String { "0.0.0.0".to_string() }
|
||||
fn default_gateway_port() -> u16 { 3100 }
|
||||
fn default_storage_root() -> String { "./data".to_string() }
|
||||
fn default_profile_root() -> String { "./data/_profiles".to_string() }
|
||||
fn default_manifest_prefix() -> String { "_catalog/manifests".to_string() }
|
||||
fn default_sidecar_url() -> String { "http://localhost:3200".to_string() }
|
||||
fn default_embed_model() -> String { "nomic-embed-text".to_string() }
|
||||
@@ -115,7 +147,12 @@ impl Default for Config {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
gateway: GatewayConfig { host: default_host(), port: default_gateway_port() },
|
||||
storage: StorageConfig { root: default_storage_root() },
|
||||
storage: StorageConfig {
|
||||
root: default_storage_root(),
|
||||
profile_root: default_profile_root(),
|
||||
rescue_bucket: None,
|
||||
buckets: Vec::new(),
|
||||
},
|
||||
catalog: CatalogConfig::default(),
|
||||
query: QueryConfig::default(),
|
||||
sidecar: SidecarConfig { url: default_sidecar_url() },
|
||||
|
||||
@@ -3,3 +3,4 @@ pub mod errors;
|
||||
pub mod arrow_helpers;
|
||||
pub mod config;
|
||||
pub mod pii;
|
||||
pub mod secrets;
|
||||
|
||||
@@ -0,0 +1,144 @@
|
||||
/// Secrets layer for bucket credentials.
|
||||
///
|
||||
/// Credentials never appear in `lakehouse.toml`. Each bucket entry has an
|
||||
/// opaque `secret_ref` handle, resolved at startup by a `SecretsProvider`.
|
||||
///
|
||||
/// MVP ships `FileSecretsProvider` — reads a TOML file at a path given by
|
||||
/// `LAKEHOUSE_SECRETS` env var, defaulting to `/etc/lakehouse/secrets.toml`.
|
||||
/// The file should be root-owned, mode 0600, never in git.
|
||||
///
|
||||
/// Future providers (Vault, SOPS, OS keyring) implement the same trait.
|
||||
|
||||
use async_trait::async_trait;
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::Arc;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BucketCredentials {
|
||||
pub access_key: String,
|
||||
pub secret_key: String,
|
||||
}
|
||||
|
||||
// Redact credentials when debug-printing — NEVER let these end up in logs.
|
||||
impl BucketCredentials {
|
||||
pub fn redacted(&self) -> String {
|
||||
let prefix: String = self.access_key.chars().take(4).collect();
|
||||
format!("access_key={}... (redacted)", prefix)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait SecretsProvider: Send + Sync {
|
||||
/// Resolve a handle to credentials. Returns Err if the handle is unknown.
|
||||
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String>;
|
||||
|
||||
/// List known handles (for diagnostics only, never values).
|
||||
async fn list_handles(&self) -> Vec<String>;
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct RawSecretEntry {
|
||||
access_key: String,
|
||||
secret_key: String,
|
||||
}
|
||||
|
||||
/// File-backed secrets provider.
|
||||
pub struct FileSecretsProvider {
|
||||
path: PathBuf,
|
||||
cache: tokio::sync::RwLock<HashMap<String, BucketCredentials>>,
|
||||
}
|
||||
|
||||
impl FileSecretsProvider {
|
||||
pub fn new(path: impl Into<PathBuf>) -> Self {
|
||||
Self {
|
||||
path: path.into(),
|
||||
cache: tokio::sync::RwLock::new(HashMap::new()),
|
||||
}
|
||||
}
|
||||
|
||||
/// Default path: env var `LAKEHOUSE_SECRETS` or `/etc/lakehouse/secrets.toml`.
|
||||
pub fn from_env() -> Self {
|
||||
let path = std::env::var("LAKEHOUSE_SECRETS")
|
||||
.unwrap_or_else(|_| "/etc/lakehouse/secrets.toml".to_string());
|
||||
Self::new(path)
|
||||
}
|
||||
|
||||
/// Arc wrapper for convenient injection.
|
||||
pub fn shared(path: impl Into<PathBuf>) -> Arc<dyn SecretsProvider> {
|
||||
Arc::new(Self::new(path))
|
||||
}
|
||||
|
||||
/// Load and populate the cache. Call once at startup.
|
||||
pub async fn load(&self) -> Result<usize, String> {
|
||||
if !self.path.exists() {
|
||||
tracing::info!(
|
||||
"secrets file {} not present — running with zero credentials configured",
|
||||
self.path.display()
|
||||
);
|
||||
return Ok(0);
|
||||
}
|
||||
check_permissions(&self.path)?;
|
||||
let content = std::fs::read_to_string(&self.path)
|
||||
.map_err(|e| format!("read secrets file: {e}"))?;
|
||||
let raw: HashMap<String, RawSecretEntry> = toml::from_str(&content)
|
||||
.map_err(|e| format!("parse secrets file: {e}"))?;
|
||||
|
||||
let mut cache = self.cache.write().await;
|
||||
for (handle, entry) in raw {
|
||||
cache.insert(handle.clone(), BucketCredentials {
|
||||
access_key: entry.access_key,
|
||||
secret_key: entry.secret_key,
|
||||
});
|
||||
}
|
||||
let count = cache.len();
|
||||
tracing::info!("secrets: loaded {count} handles from {}", self.path.display());
|
||||
Ok(count)
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl SecretsProvider for FileSecretsProvider {
|
||||
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String> {
|
||||
let cache = self.cache.read().await;
|
||||
cache
|
||||
.get(handle)
|
||||
.cloned()
|
||||
.ok_or_else(|| format!("secret handle not found: {handle}"))
|
||||
}
|
||||
|
||||
async fn list_handles(&self) -> Vec<String> {
|
||||
self.cache.read().await.keys().cloned().collect()
|
||||
}
|
||||
}
|
||||
|
||||
/// Reject world-readable secrets files. Fail closed.
|
||||
fn check_permissions(path: &Path) -> Result<(), String> {
|
||||
#[cfg(unix)]
|
||||
{
|
||||
use std::os::unix::fs::PermissionsExt;
|
||||
let meta = std::fs::metadata(path)
|
||||
.map_err(|e| format!("stat secrets file: {e}"))?;
|
||||
let mode = meta.permissions().mode() & 0o777;
|
||||
if mode & 0o044 != 0 {
|
||||
return Err(format!(
|
||||
"secrets file {} is group/world-readable (mode {:o}); chmod 600 to proceed",
|
||||
path.display(), mode
|
||||
));
|
||||
}
|
||||
}
|
||||
let _ = path;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// A zero-handles provider for tests / purely-local setups.
|
||||
pub struct EmptySecretsProvider;
|
||||
|
||||
#[async_trait]
|
||||
impl SecretsProvider for EmptySecretsProvider {
|
||||
async fn resolve(&self, handle: &str) -> Result<BucketCredentials, String> {
|
||||
Err(format!("no secrets provider configured (wanted handle '{handle}')"))
|
||||
}
|
||||
async fn list_handles(&self) -> Vec<String> { Vec::new() }
|
||||
}
|
||||
@@ -13,3 +13,4 @@ tracing = { workspace = true }
|
||||
object_store = { workspace = true }
|
||||
bytes = { workspace = true }
|
||||
futures = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
|
||||
@@ -0,0 +1,239 @@
|
||||
/// Write-once batched append log.
|
||||
///
|
||||
/// Problem we're fixing: the error journal and HNSW trial journal both
|
||||
/// previously did read-modify-write of their whole JSONL file on every
|
||||
/// event. That's O(N²) cumulative work and generates huge churn at scale.
|
||||
/// It's exactly the pattern llms3.com flags as the "small-file /
|
||||
/// rewrite-amplification" anti-pattern.
|
||||
///
|
||||
/// This helper implements the pattern object storage actually wants:
|
||||
///
|
||||
/// - Events accumulate in an **in-memory buffer** (reads see them immediately).
|
||||
/// - When the buffer hits a threshold, or `flush()` is called, the buffer is
|
||||
/// written **as one new object** with a timestamp-sorted key.
|
||||
/// - Existing objects are never rewritten.
|
||||
/// - Reads enumerate all batch files, sort by key, and concat in order.
|
||||
/// - An explicit `compact()` reads every batch file, writes one consolidated
|
||||
/// file, and deletes the originals — the LSM-tree compaction idea applied
|
||||
/// to small JSONL events.
|
||||
///
|
||||
/// Storage layout:
|
||||
/// ```
|
||||
/// {prefix}/
|
||||
/// batch_0001776319628000123.jsonl
|
||||
/// batch_0001776319745987654.jsonl
|
||||
/// batch_compacted_00001776319800000000.jsonl (after compact())
|
||||
/// ```
|
||||
/// Key format: the zero-padded epoch microsecond of the write, so lexical
|
||||
/// sort == chronological sort.
|
||||
|
||||
use bytes::Bytes;
|
||||
use chrono::Utc;
|
||||
use object_store::ObjectStore;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::ops;
|
||||
|
||||
const DEFAULT_FLUSH_THRESHOLD: usize = 32;
|
||||
|
||||
pub struct AppendLog {
|
||||
store: Arc<dyn ObjectStore>,
|
||||
prefix: String,
|
||||
buffer: Mutex<Vec<Vec<u8>>>,
|
||||
flush_threshold: usize,
|
||||
}
|
||||
|
||||
impl AppendLog {
|
||||
/// Create a new append log rooted at `prefix` in the given object store.
|
||||
/// Events auto-flush when buffer reaches `flush_threshold` (default 32).
|
||||
pub fn new(store: Arc<dyn ObjectStore>, prefix: impl Into<String>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
prefix: prefix.into(),
|
||||
buffer: Mutex::new(Vec::new()),
|
||||
flush_threshold: DEFAULT_FLUSH_THRESHOLD,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_flush_threshold(mut self, threshold: usize) -> Self {
|
||||
self.flush_threshold = threshold.max(1);
|
||||
self
|
||||
}
|
||||
|
||||
/// Add an event. The returned future completes either immediately
|
||||
/// (buffered) or after a flush, depending on whether the buffer hit the
|
||||
/// threshold. Callers don't need to care either way.
|
||||
pub async fn append(&self, line: Vec<u8>) -> Result<(), String> {
|
||||
let should_flush = {
|
||||
let mut buf = self.buffer.lock().await;
|
||||
buf.push(line);
|
||||
buf.len() >= self.flush_threshold
|
||||
};
|
||||
if should_flush {
|
||||
self.flush().await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Force-flush the in-memory buffer to object storage as a single new
|
||||
/// batch file. Safe to call anytime (idempotent no-op when buffer empty).
|
||||
pub async fn flush(&self) -> Result<(), String> {
|
||||
let batch = {
|
||||
let mut buf = self.buffer.lock().await;
|
||||
if buf.is_empty() { return Ok(()); }
|
||||
std::mem::take(&mut *buf)
|
||||
};
|
||||
|
||||
let ts_us = Utc::now().timestamp_micros().max(0) as u128;
|
||||
let key = format!("{}/batch_{:019}.jsonl", self.prefix, ts_us);
|
||||
|
||||
let mut body = Vec::with_capacity(batch.iter().map(|b| b.len() + 1).sum());
|
||||
for line in batch {
|
||||
body.extend_from_slice(&line);
|
||||
if !body.ends_with(b"\n") { body.push(b'\n'); }
|
||||
}
|
||||
ops::put(&self.store, &key, Bytes::from(body)).await
|
||||
}
|
||||
|
||||
/// Read every event across all batch files + unflushed in-memory buffer.
|
||||
/// Events are returned in chronological order.
|
||||
pub async fn read_all(&self) -> Result<Vec<Vec<u8>>, String> {
|
||||
let mut keys = self.list_batch_keys().await?;
|
||||
keys.sort();
|
||||
|
||||
let mut out = Vec::new();
|
||||
for key in keys {
|
||||
let bytes = ops::get(&self.store, &key).await?;
|
||||
for line in bytes.split(|b| *b == b'\n') {
|
||||
if !line.is_empty() {
|
||||
out.push(line.to_vec());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Include unflushed events so callers see the latest state
|
||||
// whether or not someone ran flush() recently.
|
||||
let buf = self.buffer.lock().await;
|
||||
for line in buf.iter() {
|
||||
out.push(line.clone());
|
||||
}
|
||||
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Consolidate all current batch files into one compacted file, then
|
||||
/// delete the originals. Safe to call while appends are in flight:
|
||||
/// new batches written during compaction get a higher timestamp and
|
||||
/// survive. Fails closed — if anything goes wrong mid-delete, the
|
||||
/// compacted file coexists with originals and next read sees duplicates
|
||||
/// (which the dedup caller must handle) rather than data loss.
|
||||
pub async fn compact(&self) -> Result<CompactStats, String> {
|
||||
// Snapshot which files to compact BEFORE we write the new one.
|
||||
let mut originals = self.list_batch_keys().await?;
|
||||
originals.sort();
|
||||
|
||||
if originals.len() < 2 {
|
||||
return Ok(CompactStats { merged_files: originals.len(), events: 0, new_key: None });
|
||||
}
|
||||
|
||||
// Gather all existing events.
|
||||
let mut events = Vec::new();
|
||||
for key in &originals {
|
||||
let bytes = ops::get(&self.store, key).await?;
|
||||
for line in bytes.split(|b| *b == b'\n') {
|
||||
if !line.is_empty() {
|
||||
events.push(line.to_vec());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let total_events = events.len();
|
||||
if total_events == 0 {
|
||||
// Clean up empty files without writing a new one.
|
||||
for key in &originals {
|
||||
let _ = ops::delete(&self.store, key).await;
|
||||
}
|
||||
return Ok(CompactStats { merged_files: originals.len(), events: 0, new_key: None });
|
||||
}
|
||||
|
||||
// Name the compacted file with the SAME `batch_{ts}` format so it
|
||||
// sorts chronologically with future batches. Using a distinct prefix
|
||||
// ("batch_compacted_") would break lex ordering: later `batch_N`
|
||||
// files would sort BEFORE the compacted file because 'c' > digits.
|
||||
// Timestamp = now, so any appends arriving during compaction (which
|
||||
// get the current wall-clock time) sort AFTER this file.
|
||||
let ts_us = Utc::now().timestamp_micros().max(0) as u128;
|
||||
let new_key = format!("{}/batch_{:019}.jsonl", self.prefix, ts_us);
|
||||
|
||||
let mut body = Vec::new();
|
||||
for line in &events {
|
||||
body.extend_from_slice(line);
|
||||
body.push(b'\n');
|
||||
}
|
||||
ops::put(&self.store, &new_key, Bytes::from(body)).await?;
|
||||
|
||||
// Only delete originals once the consolidated file is persisted.
|
||||
let mut failures = 0;
|
||||
for key in &originals {
|
||||
if ops::delete(&self.store, key).await.is_err() {
|
||||
failures += 1;
|
||||
}
|
||||
}
|
||||
if failures > 0 {
|
||||
tracing::warn!(
|
||||
"compact '{}': {} original files failed to delete — consolidated file {} has the data",
|
||||
self.prefix, failures, new_key,
|
||||
);
|
||||
}
|
||||
|
||||
Ok(CompactStats {
|
||||
merged_files: originals.len(),
|
||||
events: total_events,
|
||||
new_key: Some(new_key),
|
||||
})
|
||||
}
|
||||
|
||||
/// How many batch files exist on disk right now.
|
||||
pub async fn file_count(&self) -> Result<usize, String> {
|
||||
Ok(self.list_batch_keys().await?.len())
|
||||
}
|
||||
|
||||
async fn list_batch_keys(&self) -> Result<Vec<String>, String> {
|
||||
let prefix_with_slash = format!("{}/", self.prefix);
|
||||
// list the prefix then filter for keys that match our naming scheme;
|
||||
// unrelated files at the same prefix won't be touched.
|
||||
let raw = ops::list(&self.store, Some(&prefix_with_slash)).await?;
|
||||
Ok(raw
|
||||
.into_iter()
|
||||
.filter(|k| {
|
||||
let basename = k.rsplit('/').next().unwrap_or(k);
|
||||
basename.starts_with("batch_") && basename.ends_with(".jsonl")
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, serde::Serialize)]
|
||||
pub struct CompactStats {
|
||||
pub merged_files: usize,
|
||||
pub events: usize,
|
||||
pub new_key: Option<String>,
|
||||
}
|
||||
|
||||
// Proactively flush on drop, best-effort.
|
||||
// We can't `.await` in Drop; we spawn the flush on the tokio runtime if one
|
||||
// is available. If the runtime is already shutting down the flush is lost —
|
||||
// acceptable for these observability journals, which are hints not records.
|
||||
impl Drop for AppendLog {
|
||||
fn drop(&mut self) {
|
||||
let buf_len = self.buffer.try_lock().map(|b| b.len()).unwrap_or(0);
|
||||
if buf_len == 0 { return; }
|
||||
// Can't spawn from sync Drop on every runtime shape; log + move on.
|
||||
tracing::debug!(
|
||||
"append_log '{}' dropping with {} unflushed events",
|
||||
self.prefix, buf_len,
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,200 @@
|
||||
/// Bucket operation error journal.
|
||||
///
|
||||
/// Every bucket op failure and every rescue fallback lands here. Goal:
|
||||
/// answering "is anything broken?" with one HTTP call.
|
||||
///
|
||||
/// Storage: batched write-once files under `_errors/bucket_errors/` in the
|
||||
/// primary bucket. Uses the shared `AppendLog` helper so we never rewrite
|
||||
/// existing files — see `append_log.rs` for the full pattern rationale.
|
||||
///
|
||||
/// In-memory ring buffer holds the last N events for fast response; on
|
||||
/// startup `load_recent` hydrates it from all batch files.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::VecDeque;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::append_log::{AppendLog, CompactStats};
|
||||
|
||||
const JOURNAL_PREFIX: &str = "_errors/bucket_errors";
|
||||
const RING_CAPACITY: usize = 2000;
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum BucketOp {
|
||||
Read,
|
||||
Write,
|
||||
Delete,
|
||||
List,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct BucketErrorEvent {
|
||||
pub ts: DateTime<Utc>,
|
||||
pub op: BucketOp,
|
||||
pub target: String,
|
||||
pub key: String,
|
||||
pub error: String,
|
||||
#[serde(default)]
|
||||
pub rescued: bool,
|
||||
}
|
||||
|
||||
impl BucketErrorEvent {
|
||||
pub fn new_read(target: &str, key: &str, error: &str) -> Self {
|
||||
Self { ts: Utc::now(), op: BucketOp::Read, target: target.into(), key: key.into(), error: error.into(), rescued: false }
|
||||
}
|
||||
pub fn new_write(target: &str, key: &str, error: &str) -> Self {
|
||||
Self { ts: Utc::now(), op: BucketOp::Write, target: target.into(), key: key.into(), error: error.into(), rescued: false }
|
||||
}
|
||||
pub fn new_delete(target: &str, key: &str, error: &str) -> Self {
|
||||
Self { ts: Utc::now(), op: BucketOp::Delete, target: target.into(), key: key.into(), error: error.into(), rescued: false }
|
||||
}
|
||||
pub fn new_list(target: &str, prefix: &str, error: &str) -> Self {
|
||||
Self { ts: Utc::now(), op: BucketOp::List, target: target.into(), key: prefix.into(), error: error.into(), rescued: false }
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct ErrorJournal {
|
||||
log: Arc<AppendLog>,
|
||||
ring: Arc<RwLock<VecDeque<BucketErrorEvent>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct HealthReport {
|
||||
pub period_minutes: i64,
|
||||
pub total_errors: usize,
|
||||
pub per_bucket: std::collections::HashMap<String, usize>,
|
||||
pub unhealthy_buckets: Vec<String>,
|
||||
}
|
||||
|
||||
impl ErrorJournal {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
// Keep flush threshold lowish — operators checking /storage/errors
|
||||
// after a recent incident want to see fresh rows even without an
|
||||
// explicit flush.
|
||||
let log = Arc::new(
|
||||
AppendLog::new(store, JOURNAL_PREFIX).with_flush_threshold(8),
|
||||
);
|
||||
Self {
|
||||
log,
|
||||
ring: Arc::new(RwLock::new(VecDeque::with_capacity(RING_CAPACITY))),
|
||||
}
|
||||
}
|
||||
|
||||
/// Hydrate the ring buffer from existing batch files. Call once at
|
||||
/// startup. Tolerates malformed lines (skipped with a warning) and
|
||||
/// missing files (returns 0).
|
||||
pub async fn load_recent(&self) -> Result<usize, String> {
|
||||
let lines = self.log.read_all().await.unwrap_or_default();
|
||||
let mut ring = self.ring.write().await;
|
||||
for line in lines {
|
||||
match serde_json::from_slice::<BucketErrorEvent>(&line) {
|
||||
Ok(ev) => {
|
||||
if ring.len() >= RING_CAPACITY { ring.pop_front(); }
|
||||
ring.push_back(ev);
|
||||
}
|
||||
Err(e) => tracing::warn!("error journal: skip malformed line ({e})"),
|
||||
}
|
||||
}
|
||||
Ok(ring.len())
|
||||
}
|
||||
|
||||
/// Append an event. In-memory ring updated immediately; persistence
|
||||
/// happens in batches via the underlying AppendLog.
|
||||
pub async fn append(&self, event: BucketErrorEvent) {
|
||||
{
|
||||
let mut ring = self.ring.write().await;
|
||||
if ring.len() >= RING_CAPACITY { ring.pop_front(); }
|
||||
ring.push_back(event.clone());
|
||||
}
|
||||
|
||||
match serde_json::to_vec(&event) {
|
||||
Ok(line) => {
|
||||
if let Err(e) = self.log.append(line).await {
|
||||
tracing::error!("error journal persist failed: {e}");
|
||||
}
|
||||
}
|
||||
Err(e) => tracing::error!("error journal serialize failed: {e}"),
|
||||
}
|
||||
}
|
||||
|
||||
/// Mark the most recent matching in-memory event as rescued.
|
||||
/// Only updates the ring buffer — the JSONL line already persisted
|
||||
/// records the failure fact; rescue status travels in response headers.
|
||||
pub async fn mark_rescued_last(&self, target: &str, key: &str) {
|
||||
let mut ring = self.ring.write().await;
|
||||
for ev in ring.iter_mut().rev() {
|
||||
if matches!(ev.op, BucketOp::Read) && ev.target == target && ev.key == key && !ev.rescued {
|
||||
ev.rescued = true;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn recent(&self, limit: usize) -> Vec<BucketErrorEvent> {
|
||||
let ring = self.ring.read().await;
|
||||
let start = ring.len().saturating_sub(limit);
|
||||
ring.iter().skip(start).cloned().collect()
|
||||
}
|
||||
|
||||
pub async fn filter(
|
||||
&self,
|
||||
bucket: Option<&str>,
|
||||
since: Option<DateTime<Utc>>,
|
||||
limit: usize,
|
||||
) -> Vec<BucketErrorEvent> {
|
||||
let ring = self.ring.read().await;
|
||||
ring.iter()
|
||||
.rev()
|
||||
.filter(|ev| bucket.map_or(true, |b| ev.target == b))
|
||||
.filter(|ev| since.map_or(true, |s| ev.ts >= s))
|
||||
.take(limit)
|
||||
.cloned()
|
||||
.collect::<Vec<_>>()
|
||||
.into_iter()
|
||||
.rev()
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Summarize errors in the last N minutes.
|
||||
pub async fn health(&self, period_minutes: i64) -> HealthReport {
|
||||
use std::collections::HashMap;
|
||||
let cutoff = Utc::now() - chrono::Duration::minutes(period_minutes);
|
||||
let ring = self.ring.read().await;
|
||||
let recent: Vec<_> = ring.iter().filter(|ev| ev.ts >= cutoff).collect();
|
||||
let mut per_bucket: HashMap<String, usize> = HashMap::new();
|
||||
for ev in &recent {
|
||||
*per_bucket.entry(ev.target.clone()).or_insert(0) += 1;
|
||||
}
|
||||
let unhealthy_buckets: Vec<String> = per_bucket
|
||||
.iter()
|
||||
.filter(|(_, c)| **c >= 3)
|
||||
.map(|(k, _)| k.clone())
|
||||
.collect();
|
||||
HealthReport {
|
||||
period_minutes,
|
||||
total_errors: recent.len(),
|
||||
per_bucket,
|
||||
unhealthy_buckets,
|
||||
}
|
||||
}
|
||||
|
||||
/// Force an immediate flush of buffered events to object storage.
|
||||
pub async fn flush(&self) -> Result<(), String> {
|
||||
self.log.flush().await
|
||||
}
|
||||
|
||||
/// Consolidate all batch files into one. Operator cleanup.
|
||||
pub async fn compact(&self) -> Result<CompactStats, String> {
|
||||
self.log.compact().await
|
||||
}
|
||||
|
||||
/// How many JSONL batch files currently exist.
|
||||
pub async fn file_count(&self) -> Result<usize, String> {
|
||||
self.log.file_count().await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
/// HTTP surface for federation operator endpoints.
|
||||
///
|
||||
/// - `GET /buckets` — configured buckets with reachability status
|
||||
/// - `GET /errors` — recent bucket op failures (filterable)
|
||||
/// - `GET /bucket-health` — aggregated errors-per-bucket in the last 5 minutes
|
||||
///
|
||||
/// Mounted under `/storage` by the gateway. `/storage/health` is already
|
||||
/// claimed by the existing storage router for service liveness, so the
|
||||
/// aggregated bucket health lives at `/storage/bucket-health`.
|
||||
|
||||
use axum::{
|
||||
Json, Router,
|
||||
extract::{Path, Query, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{get, put},
|
||||
};
|
||||
use chrono::{DateTime, Utc};
|
||||
use serde::Deserialize;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::error_journal::BucketErrorEvent;
|
||||
use crate::registry::{BucketInfo, BucketRegistry};
|
||||
|
||||
pub fn router(registry: Arc<BucketRegistry>) -> Router {
|
||||
Router::new()
|
||||
.route("/buckets", get(list_buckets))
|
||||
.route("/errors", get(list_errors))
|
||||
.route("/errors/flush", axum::routing::post(flush_errors))
|
||||
.route("/errors/compact", axum::routing::post(compact_errors))
|
||||
.route("/bucket-health", get(get_health))
|
||||
.route("/buckets/{bucket}/objects/{*key}", put(put_bucket_object))
|
||||
.route("/buckets/{bucket}/objects/{*key}", get(get_bucket_object))
|
||||
.with_state(registry)
|
||||
}
|
||||
|
||||
async fn list_buckets(State(reg): State<Arc<BucketRegistry>>) -> Json<Vec<BucketInfo>> {
|
||||
Json(reg.list().await)
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ErrorQuery {
|
||||
bucket: Option<String>,
|
||||
since: Option<DateTime<Utc>>,
|
||||
#[serde(default = "default_limit")]
|
||||
limit: usize,
|
||||
}
|
||||
|
||||
fn default_limit() -> usize { 50 }
|
||||
|
||||
async fn list_errors(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Query(q): Query<ErrorQuery>,
|
||||
) -> Json<Vec<BucketErrorEvent>> {
|
||||
Json(reg.journal().filter(q.bucket.as_deref(), q.since, q.limit).await)
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct HealthQuery {
|
||||
#[serde(default = "default_period")]
|
||||
minutes: i64,
|
||||
}
|
||||
|
||||
fn default_period() -> i64 { 5 }
|
||||
|
||||
async fn get_health(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Query(q): Query<HealthQuery>,
|
||||
) -> impl IntoResponse {
|
||||
Json(reg.journal().health(q.minutes).await)
|
||||
}
|
||||
|
||||
/// Bucket-aware write. Hard-fails if the target bucket is unreachable;
|
||||
/// never falls back to rescue for writes.
|
||||
async fn put_bucket_object(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Path((bucket, key)): Path<(String, String)>,
|
||||
body: bytes::Bytes,
|
||||
) -> impl IntoResponse {
|
||||
match reg.write_smart(&bucket, &key, body).await {
|
||||
Ok(()) => (StatusCode::CREATED, format!("stored: {bucket}/{key}")).into_response(),
|
||||
Err(e) => (StatusCode::SERVICE_UNAVAILABLE, e).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Bucket-aware read. Falls through to the rescue bucket if the target
|
||||
/// is unreachable or the key is missing; emits observability headers so
|
||||
/// callers can detect the fallback.
|
||||
async fn flush_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
|
||||
match reg.journal().flush().await {
|
||||
Ok(()) => (StatusCode::OK, "flushed").into_response(),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn compact_errors(State(reg): State<Arc<BucketRegistry>>) -> impl IntoResponse {
|
||||
match reg.journal().compact().await {
|
||||
Ok(stats) => Json(stats).into_response(),
|
||||
Err(e) => (StatusCode::INTERNAL_SERVER_ERROR, e).into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_bucket_object(
|
||||
State(reg): State<Arc<BucketRegistry>>,
|
||||
Path((bucket, key)): Path<(String, String)>,
|
||||
) -> impl IntoResponse {
|
||||
match reg.read_smart(&bucket, &key).await {
|
||||
Ok(outcome) => {
|
||||
let mut resp = outcome.data.into_response();
|
||||
if outcome.rescued {
|
||||
let headers = resp.headers_mut();
|
||||
headers.insert("x-lakehouse-rescue-used", "true".parse().unwrap());
|
||||
headers.insert(
|
||||
"x-lakehouse-original-bucket",
|
||||
outcome.original_bucket.parse().unwrap(),
|
||||
);
|
||||
headers.insert(
|
||||
"x-lakehouse-served-by",
|
||||
outcome.served_by.parse().unwrap(),
|
||||
);
|
||||
}
|
||||
resp
|
||||
}
|
||||
Err(e) => (StatusCode::NOT_FOUND, e).into_response(),
|
||||
}
|
||||
}
|
||||
@@ -1,3 +1,7 @@
|
||||
pub mod append_log;
|
||||
pub mod backend;
|
||||
pub mod error_journal;
|
||||
pub mod federation_service;
|
||||
pub mod ops;
|
||||
pub mod registry;
|
||||
pub mod service;
|
||||
|
||||
@@ -0,0 +1,283 @@
|
||||
/// Multi-backend bucket registry — the federation foundation.
|
||||
///
|
||||
/// Federation rule: every `ObjectRef` belongs to exactly one named bucket.
|
||||
/// The registry resolves bucket names to `object_store` backends, handles
|
||||
/// rescue-bucket fallback on read failure, writes every failure to the
|
||||
/// error journal, and exposes a health summary for operators.
|
||||
///
|
||||
/// Existing call sites can keep using `ops::*` with `registry.get(name)`.
|
||||
/// New bucket-aware call sites use `registry.read_smart` / `write_smart`
|
||||
/// which handle fallback + journaling automatically.
|
||||
|
||||
use object_store::ObjectStore;
|
||||
use object_store::local::LocalFileSystem;
|
||||
use serde::Serialize;
|
||||
use shared::config::{BucketConfig, StorageConfig};
|
||||
use shared::secrets::{BucketCredentials, SecretsProvider};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use crate::error_journal::{BucketErrorEvent, ErrorJournal};
|
||||
|
||||
/// A registered bucket — the store handle + its configuration.
|
||||
pub struct BucketEntry {
|
||||
pub name: String,
|
||||
pub backend: String,
|
||||
pub store: Arc<dyn ObjectStore>,
|
||||
pub config: BucketConfig,
|
||||
}
|
||||
|
||||
/// Read outcome — may have been rescued.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct ReadOutcome {
|
||||
pub data: bytes::Bytes,
|
||||
pub rescued: bool,
|
||||
pub original_bucket: String,
|
||||
pub served_by: String,
|
||||
}
|
||||
|
||||
/// Summary entry for GET /storage/buckets.
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct BucketInfo {
|
||||
pub name: String,
|
||||
pub backend: String,
|
||||
pub reachable: bool,
|
||||
pub role: BucketRole,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, PartialEq)]
|
||||
#[serde(rename_all = "lowercase")]
|
||||
pub enum BucketRole {
|
||||
Primary,
|
||||
Rescue,
|
||||
Profile,
|
||||
Tenant,
|
||||
}
|
||||
|
||||
pub struct BucketRegistry {
|
||||
buckets: HashMap<String, Arc<BucketEntry>>,
|
||||
default: String,
|
||||
rescue: Option<String>,
|
||||
profile_root: String,
|
||||
journal: ErrorJournal,
|
||||
}
|
||||
|
||||
impl BucketRegistry {
|
||||
/// Build the registry from storage config + secrets provider.
|
||||
/// Back-compat: if `buckets` is empty, synthesize a `primary` bucket from
|
||||
/// the legacy `root` field so pre-federation configs keep working.
|
||||
pub async fn from_config(
|
||||
cfg: &StorageConfig,
|
||||
secrets: Arc<dyn SecretsProvider>,
|
||||
) -> Result<Self, String> {
|
||||
let mut buckets: HashMap<String, Arc<BucketEntry>> = HashMap::new();
|
||||
|
||||
let bucket_configs: Vec<BucketConfig> = if cfg.buckets.is_empty() {
|
||||
vec![BucketConfig {
|
||||
name: "primary".to_string(),
|
||||
backend: "local".to_string(),
|
||||
root: Some(cfg.root.clone()),
|
||||
bucket: None,
|
||||
region: None,
|
||||
endpoint: None,
|
||||
secret_ref: None,
|
||||
}]
|
||||
} else {
|
||||
cfg.buckets.clone()
|
||||
};
|
||||
|
||||
for bc in bucket_configs {
|
||||
let store = build_store(&bc, secrets.as_ref()).await?;
|
||||
let entry = Arc::new(BucketEntry {
|
||||
name: bc.name.clone(),
|
||||
backend: bc.backend.clone(),
|
||||
store,
|
||||
config: bc.clone(),
|
||||
});
|
||||
buckets.insert(bc.name.clone(), entry);
|
||||
}
|
||||
|
||||
// Ensure `primary` always exists — it's where error journals live.
|
||||
if !buckets.contains_key("primary") {
|
||||
return Err("no bucket named 'primary' configured — required as error-journal home".into());
|
||||
}
|
||||
|
||||
// Rescue bucket is optional but, if named, must exist.
|
||||
if let Some(r) = &cfg.rescue_bucket {
|
||||
if !buckets.contains_key(r) {
|
||||
return Err(format!("rescue_bucket '{r}' not found among configured buckets"));
|
||||
}
|
||||
}
|
||||
|
||||
let journal = ErrorJournal::new(buckets.get("primary").unwrap().store.clone());
|
||||
let _ = journal.load_recent().await;
|
||||
|
||||
Ok(Self {
|
||||
buckets,
|
||||
default: "primary".to_string(),
|
||||
rescue: cfg.rescue_bucket.clone(),
|
||||
profile_root: cfg.profile_root.clone(),
|
||||
journal,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn default_name(&self) -> &str { &self.default }
|
||||
|
||||
pub fn rescue_name(&self) -> Option<&str> { self.rescue.as_deref() }
|
||||
|
||||
pub fn journal(&self) -> &ErrorJournal { &self.journal }
|
||||
|
||||
/// Resolve a bucket name to its object store. Existing call sites use
|
||||
/// this as a drop-in replacement for the old single-store pattern.
|
||||
pub fn get(&self, bucket: &str) -> Result<Arc<dyn ObjectStore>, String> {
|
||||
self.buckets
|
||||
.get(bucket)
|
||||
.map(|e| e.store.clone())
|
||||
.ok_or_else(|| format!("unknown bucket: {bucket}"))
|
||||
}
|
||||
|
||||
/// The default bucket's store — use for code paths that don't yet know
|
||||
/// about buckets.
|
||||
pub fn default_store(&self) -> Arc<dyn ObjectStore> {
|
||||
self.buckets.get(&self.default).unwrap().store.clone()
|
||||
}
|
||||
|
||||
/// List all registered buckets. Checks reachability by doing a trivial
|
||||
/// `list` with limit 1 on each.
|
||||
pub async fn list(&self) -> Vec<BucketInfo> {
|
||||
let mut out = Vec::with_capacity(self.buckets.len());
|
||||
for (name, entry) in &self.buckets {
|
||||
let reachable = probe(&entry.store).await;
|
||||
let role = self.classify(name);
|
||||
out.push(BucketInfo {
|
||||
name: name.clone(),
|
||||
backend: entry.backend.clone(),
|
||||
reachable,
|
||||
role,
|
||||
});
|
||||
}
|
||||
out.sort_by(|a, b| a.name.cmp(&b.name));
|
||||
out
|
||||
}
|
||||
|
||||
fn classify(&self, name: &str) -> BucketRole {
|
||||
if name == self.default { BucketRole::Primary }
|
||||
else if Some(name) == self.rescue.as_deref() { BucketRole::Rescue }
|
||||
else if name.starts_with("profile:") { BucketRole::Profile }
|
||||
else { BucketRole::Tenant }
|
||||
}
|
||||
|
||||
/// Read with rescue-bucket fallback. If the target bucket fails and a
|
||||
/// rescue is configured, retries against rescue. Records every failure
|
||||
/// in the error journal.
|
||||
pub async fn read_smart(&self, bucket: &str, key: &str) -> Result<ReadOutcome, String> {
|
||||
let target = self.buckets.get(bucket)
|
||||
.ok_or_else(|| format!("unknown bucket: {bucket}"))?;
|
||||
|
||||
match crate::ops::get(&target.store, key).await {
|
||||
Ok(data) => Ok(ReadOutcome {
|
||||
data, rescued: false,
|
||||
original_bucket: bucket.to_string(),
|
||||
served_by: bucket.to_string(),
|
||||
}),
|
||||
Err(err) => {
|
||||
// Record failure regardless of what happens next.
|
||||
self.journal.append(BucketErrorEvent::new_read(bucket, key, &err)).await;
|
||||
|
||||
// Try rescue, if any.
|
||||
if let Some(rescue_name) = &self.rescue {
|
||||
if rescue_name != bucket {
|
||||
if let Some(rescue) = self.buckets.get(rescue_name) {
|
||||
match crate::ops::get(&rescue.store, key).await {
|
||||
Ok(data) => {
|
||||
self.journal.mark_rescued_last(bucket, key).await;
|
||||
return Ok(ReadOutcome {
|
||||
data, rescued: true,
|
||||
original_bucket: bucket.to_string(),
|
||||
served_by: rescue_name.clone(),
|
||||
});
|
||||
}
|
||||
Err(rescue_err) => {
|
||||
return Err(format!(
|
||||
"read '{key}' failed in '{bucket}' ({err}); rescue '{rescue_name}' also failed ({rescue_err})"
|
||||
));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(format!("read '{key}' failed in '{bucket}': {err}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Write always goes to target. No rescue fallback for writes — writes
|
||||
/// that silently vanish are the worst possible failure.
|
||||
pub async fn write_smart(
|
||||
&self,
|
||||
bucket: &str,
|
||||
key: &str,
|
||||
data: bytes::Bytes,
|
||||
) -> Result<(), String> {
|
||||
let target = self.buckets.get(bucket)
|
||||
.ok_or_else(|| format!("unknown bucket: {bucket}"))?;
|
||||
match crate::ops::put(&target.store, key, data).await {
|
||||
Ok(()) => Ok(()),
|
||||
Err(err) => {
|
||||
self.journal.append(BucketErrorEvent::new_write(bucket, key, &err)).await;
|
||||
Err(format!("write '{key}' failed in '{bucket}': {err}"))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Trivial reachability check — try to list with limit 0.
|
||||
async fn probe(store: &Arc<dyn ObjectStore>) -> bool {
|
||||
use futures::StreamExt;
|
||||
let mut stream = store.list(None);
|
||||
// Pulling the first item confirms the store responds. Empty bucket = ok.
|
||||
match stream.next().await {
|
||||
Some(Ok(_)) => true,
|
||||
None => true, // empty but reachable
|
||||
Some(Err(_)) => false,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a concrete ObjectStore from a BucketConfig.
|
||||
async fn build_store(
|
||||
bc: &BucketConfig,
|
||||
secrets: &dyn SecretsProvider,
|
||||
) -> Result<Arc<dyn ObjectStore>, String> {
|
||||
match bc.backend.as_str() {
|
||||
"local" => {
|
||||
let root = bc.root.as_deref()
|
||||
.ok_or_else(|| format!("bucket '{}' is backend=local but has no root", bc.name))?;
|
||||
std::fs::create_dir_all(root)
|
||||
.map_err(|e| format!("create bucket dir '{root}': {e}"))?;
|
||||
let fs = LocalFileSystem::new_with_prefix(root)
|
||||
.map_err(|e| format!("init local bucket '{}': {e}", bc.name))?;
|
||||
Ok(Arc::new(fs))
|
||||
}
|
||||
"s3" => {
|
||||
let handle = bc.secret_ref.as_deref()
|
||||
.ok_or_else(|| format!("s3 bucket '{}' has no secret_ref", bc.name))?;
|
||||
let creds: BucketCredentials = secrets.resolve(handle).await?;
|
||||
let s3_bucket = bc.bucket.as_deref()
|
||||
.ok_or_else(|| format!("s3 bucket '{}' has no `bucket` name", bc.name))?;
|
||||
let region = bc.region.as_deref().unwrap_or("us-east-1");
|
||||
|
||||
let mut builder = object_store::aws::AmazonS3Builder::new()
|
||||
.with_bucket_name(s3_bucket)
|
||||
.with_region(region)
|
||||
.with_access_key_id(&creds.access_key)
|
||||
.with_secret_access_key(&creds.secret_key);
|
||||
if let Some(endpoint) = &bc.endpoint {
|
||||
builder = builder.with_endpoint(endpoint);
|
||||
}
|
||||
let s3 = builder.build()
|
||||
.map_err(|e| format!("init s3 bucket '{}': {e}", bc.name))?;
|
||||
Ok(Arc::new(s3))
|
||||
}
|
||||
other => Err(format!("unknown backend '{other}' for bucket '{}'", bc.name)),
|
||||
}
|
||||
}
|
||||
@@ -18,3 +18,4 @@ parquet = { workspace = true }
|
||||
arrow = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
instant-distance = { workspace = true }
|
||||
uuid = { workspace = true }
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
/// In-memory cache for StoredEmbedding vectors.
|
||||
///
|
||||
/// Rationale: loading 100K embeddings from Parquet takes ~2-5s. When an AI agent
|
||||
/// iterates on HNSW parameters, each trial would repeat that cost. The cache
|
||||
/// pins embeddings in memory so trials reuse them.
|
||||
///
|
||||
/// This is a pure performance layer — the Parquet file is still the source of
|
||||
/// truth (ADR-008). Eviction is safe; worst case is one slow reload.
|
||||
|
||||
use object_store::ObjectStore;
|
||||
use serde::Serialize;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::store::{self, StoredEmbedding};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct EmbeddingCache {
|
||||
store: Arc<dyn ObjectStore>,
|
||||
cache: Arc<RwLock<HashMap<String, Arc<Vec<StoredEmbedding>>>>>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct CacheEntry {
|
||||
pub index_name: String,
|
||||
pub vectors: usize,
|
||||
pub dimensions: usize,
|
||||
pub memory_bytes: u64,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize)]
|
||||
pub struct CacheStats {
|
||||
pub entries: Vec<CacheEntry>,
|
||||
pub total_memory_bytes: u64,
|
||||
}
|
||||
|
||||
impl EmbeddingCache {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
cache: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
/// Return cached embeddings, loading from object storage on first request.
|
||||
pub async fn get_or_load(
|
||||
&self,
|
||||
index_name: &str,
|
||||
) -> Result<Arc<Vec<StoredEmbedding>>, String> {
|
||||
if let Some(cached) = self.cache.read().await.get(index_name) {
|
||||
return Ok(cached.clone());
|
||||
}
|
||||
|
||||
// Load under a write lock so concurrent callers only hit disk once.
|
||||
let mut guard = self.cache.write().await;
|
||||
if let Some(cached) = guard.get(index_name) {
|
||||
return Ok(cached.clone());
|
||||
}
|
||||
tracing::info!("embedding_cache: loading '{index_name}' from object storage");
|
||||
let t0 = std::time::Instant::now();
|
||||
let loaded = store::load_embeddings(&self.store, index_name).await?;
|
||||
let n = loaded.len();
|
||||
let arc = Arc::new(loaded);
|
||||
guard.insert(index_name.to_string(), arc.clone());
|
||||
tracing::info!(
|
||||
"embedding_cache: loaded '{index_name}' — {n} vectors in {:.2}s",
|
||||
t0.elapsed().as_secs_f32()
|
||||
);
|
||||
Ok(arc)
|
||||
}
|
||||
|
||||
pub async fn evict(&self, index_name: &str) -> bool {
|
||||
self.cache.write().await.remove(index_name).is_some()
|
||||
}
|
||||
|
||||
pub async fn stats(&self) -> CacheStats {
|
||||
let cache = self.cache.read().await;
|
||||
let mut entries = Vec::with_capacity(cache.len());
|
||||
let mut total: u64 = 0;
|
||||
for (name, embs) in cache.iter() {
|
||||
let dims = embs.first().map(|e| e.vector.len()).unwrap_or(0);
|
||||
// Rough estimate: vector data + chunk_text + metadata overhead.
|
||||
let vec_bytes = (embs.len() * dims * std::mem::size_of::<f32>()) as u64;
|
||||
let text_bytes: u64 = embs.iter().map(|e| e.chunk_text.len() as u64).sum();
|
||||
let overhead = (embs.len() * 128) as u64; // strings + struct overhead
|
||||
let mem = vec_bytes + text_bytes + overhead;
|
||||
total += mem;
|
||||
entries.push(CacheEntry {
|
||||
index_name: name.clone(),
|
||||
vectors: embs.len(),
|
||||
dimensions: dims,
|
||||
memory_bytes: mem,
|
||||
});
|
||||
}
|
||||
CacheStats { entries, total_memory_bytes: total }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,216 @@
|
||||
/// Eval harness for HNSW tuning.
|
||||
///
|
||||
/// An EvalSet is a named collection of queries against a specific vector index.
|
||||
/// Each query has a query_text and (optionally pre-computed) ground_truth
|
||||
/// doc_ids — the "correct" top-k results that brute-force cosine returns.
|
||||
/// HNSW trials are scored by how well they recreate that top-k (recall@k).
|
||||
///
|
||||
/// Storage: `_hnsw_evals/{name}.json` as a single JSON document. Small.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
use aibridge::client::{AiClient, EmbedRequest};
|
||||
use storaged::ops;
|
||||
|
||||
use crate::store::StoredEmbedding;
|
||||
|
||||
/// A single eval query with optional pre-computed ground truth.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct EvalQuery {
|
||||
pub id: String,
|
||||
pub query_text: String,
|
||||
/// Ordered list of doc_ids that brute-force returns as top-k.
|
||||
/// `None` means the ground truth hasn't been computed yet; `compute_ground_truth`
|
||||
/// fills it in against the current embedding set.
|
||||
#[serde(default)]
|
||||
pub ground_truth: Option<Vec<String>>,
|
||||
/// Optional pre-computed query embedding. Filling this avoids re-embedding
|
||||
/// the same query on every trial.
|
||||
#[serde(default)]
|
||||
pub query_embedding: Option<Vec<f32>>,
|
||||
}
|
||||
|
||||
/// A named eval set tied to a specific vector index.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct EvalSet {
|
||||
pub name: String,
|
||||
pub index_name: String,
|
||||
pub k: usize, // top-k used for recall calculation
|
||||
pub queries: Vec<EvalQuery>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub ground_truth_built: bool,
|
||||
}
|
||||
|
||||
impl EvalSet {
|
||||
pub fn new(name: &str, index_name: &str, k: usize) -> Self {
|
||||
Self {
|
||||
name: name.to_string(),
|
||||
index_name: index_name.to_string(),
|
||||
k,
|
||||
queries: Vec::new(),
|
||||
created_at: Utc::now(),
|
||||
ground_truth_built: false,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn storage_key(&self) -> String {
|
||||
format!("_hnsw_evals/{}.json", self.name)
|
||||
}
|
||||
|
||||
pub async fn save(&self, store: &Arc<dyn ObjectStore>) -> Result<(), String> {
|
||||
let json = serde_json::to_vec_pretty(self).map_err(|e| e.to_string())?;
|
||||
ops::put(store, &self.storage_key(), json.into()).await
|
||||
}
|
||||
|
||||
pub async fn load(store: &Arc<dyn ObjectStore>, name: &str) -> Result<Self, String> {
|
||||
let key = format!("_hnsw_evals/{}.json", name);
|
||||
let data = ops::get(store, &key).await?;
|
||||
serde_json::from_slice(&data).map_err(|e| format!("parse eval set: {e}"))
|
||||
}
|
||||
|
||||
pub async fn list(store: &Arc<dyn ObjectStore>) -> Result<Vec<String>, String> {
|
||||
let keys = ops::list(store, Some("_hnsw_evals/")).await?;
|
||||
Ok(keys
|
||||
.into_iter()
|
||||
.filter(|k| k.ends_with(".json"))
|
||||
.map(|k| {
|
||||
k.trim_start_matches("_hnsw_evals/")
|
||||
.trim_end_matches(".json")
|
||||
.to_string()
|
||||
})
|
||||
.collect())
|
||||
}
|
||||
}
|
||||
|
||||
/// Cosine similarity for two same-length f32 slices.
|
||||
fn cosine(a: &[f32], b: &[f32]) -> f32 {
|
||||
let mut dot = 0.0f32;
|
||||
let mut na = 0.0f32;
|
||||
let mut nb = 0.0f32;
|
||||
for i in 0..a.len() {
|
||||
dot += a[i] * b[i];
|
||||
na += a[i] * a[i];
|
||||
nb += b[i] * b[i];
|
||||
}
|
||||
if na == 0.0 || nb == 0.0 {
|
||||
return 0.0;
|
||||
}
|
||||
dot / (na.sqrt() * nb.sqrt())
|
||||
}
|
||||
|
||||
/// Brute-force top-k search for a single query against all embeddings.
|
||||
/// This is the ground-truth oracle that HNSW trials must approximate.
|
||||
pub fn brute_force_top_k(
|
||||
query: &[f32],
|
||||
embeddings: &[StoredEmbedding],
|
||||
k: usize,
|
||||
) -> Vec<String> {
|
||||
let mut scored: Vec<(f32, usize)> = embeddings
|
||||
.iter()
|
||||
.enumerate()
|
||||
.map(|(i, e)| (cosine(query, &e.vector), i))
|
||||
.collect();
|
||||
// Partial sort — we only need top-k.
|
||||
scored.sort_unstable_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(std::cmp::Ordering::Equal));
|
||||
scored
|
||||
.into_iter()
|
||||
.take(k)
|
||||
.map(|(_, i)| embeddings[i].doc_id.clone())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Recall@k — fraction of ground-truth doc_ids that appear in predicted.
|
||||
pub fn recall_at_k(predicted: &[String], ground_truth: &[String], k: usize) -> f32 {
|
||||
if ground_truth.is_empty() || k == 0 {
|
||||
return 0.0;
|
||||
}
|
||||
let gt_set: std::collections::HashSet<&String> =
|
||||
ground_truth.iter().take(k).collect();
|
||||
let hits = predicted
|
||||
.iter()
|
||||
.take(k)
|
||||
.filter(|d| gt_set.contains(d))
|
||||
.count();
|
||||
hits as f32 / gt_set.len() as f32
|
||||
}
|
||||
|
||||
/// Populate query_embedding and ground_truth for every query that lacks them.
|
||||
pub async fn compute_ground_truth(
|
||||
eval: &mut EvalSet,
|
||||
embeddings: &[StoredEmbedding],
|
||||
ai_client: &AiClient,
|
||||
) -> Result<(), String> {
|
||||
let need_embed: Vec<(usize, String)> = eval
|
||||
.queries
|
||||
.iter()
|
||||
.enumerate()
|
||||
.filter(|(_, q)| q.query_embedding.is_none())
|
||||
.map(|(i, q)| (i, q.query_text.clone()))
|
||||
.collect();
|
||||
|
||||
if !need_embed.is_empty() {
|
||||
// Embed in one batch to keep things simple; for very large eval sets
|
||||
// we'd batch this in chunks of 32.
|
||||
let texts: Vec<String> = need_embed.iter().map(|(_, t)| t.clone()).collect();
|
||||
let resp = ai_client
|
||||
.embed(EmbedRequest { texts, model: None })
|
||||
.await
|
||||
.map_err(|e| format!("embed queries: {e}"))?;
|
||||
for ((idx, _), vec) in need_embed.iter().zip(resp.embeddings.iter()) {
|
||||
let v: Vec<f32> = vec.iter().map(|&x| x as f32).collect();
|
||||
eval.queries[*idx].query_embedding = Some(v);
|
||||
}
|
||||
}
|
||||
|
||||
for q in eval.queries.iter_mut() {
|
||||
if q.ground_truth.is_some() {
|
||||
continue;
|
||||
}
|
||||
let emb = q.query_embedding.as_ref().ok_or("missing embedding")?;
|
||||
q.ground_truth = Some(brute_force_top_k(emb, embeddings, eval.k));
|
||||
}
|
||||
eval.ground_truth_built = true;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Auto-generate a synthetic eval set by sampling every Nth chunk's text as
|
||||
/// its own query. Useful for a quick-start eval when the user doesn't have
|
||||
/// real natural-language queries yet.
|
||||
pub fn synthetic_from_chunks(
|
||||
eval_name: &str,
|
||||
index_name: &str,
|
||||
embeddings: &[StoredEmbedding],
|
||||
sample_count: usize,
|
||||
k: usize,
|
||||
) -> EvalSet {
|
||||
let n = embeddings.len();
|
||||
let sample_count = sample_count.min(n);
|
||||
let stride = (n / sample_count.max(1)).max(1);
|
||||
|
||||
let mut queries = Vec::with_capacity(sample_count);
|
||||
for i in 0..sample_count {
|
||||
let idx = (i * stride).min(n - 1);
|
||||
let chunk = &embeddings[idx];
|
||||
// Use the first ~200 chars of the chunk as the "query" — it should find
|
||||
// itself and nearby chunks as top results.
|
||||
let query_text: String = chunk.chunk_text.chars().take(200).collect();
|
||||
queries.push(EvalQuery {
|
||||
id: format!("syn-{}", i),
|
||||
query_text,
|
||||
ground_truth: None,
|
||||
query_embedding: None,
|
||||
});
|
||||
}
|
||||
|
||||
EvalSet {
|
||||
name: eval_name.to_string(),
|
||||
index_name: index_name.to_string(),
|
||||
k,
|
||||
queries,
|
||||
created_at: Utc::now(),
|
||||
ground_truth_built: false,
|
||||
}
|
||||
}
|
||||
@@ -7,6 +7,7 @@ use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use crate::store::StoredEmbedding;
|
||||
use crate::trial::HnswConfig;
|
||||
|
||||
/// A vector point for HNSW — wraps f32 slice with cosine distance.
|
||||
#[derive(Clone)]
|
||||
@@ -62,18 +63,32 @@ impl HnswStore {
|
||||
}
|
||||
}
|
||||
|
||||
/// Build an HNSW index from stored embeddings.
|
||||
/// Build an HNSW index from stored embeddings with default config.
|
||||
pub async fn build_index(
|
||||
&self,
|
||||
index_name: &str,
|
||||
embeddings: Vec<StoredEmbedding>,
|
||||
) -> Result<BuildStats, String> {
|
||||
self.build_index_with_config(index_name, embeddings, &HnswConfig::default()).await
|
||||
}
|
||||
|
||||
/// Build an HNSW index from stored embeddings with explicit config.
|
||||
/// Used by the trial system — each trial calls this with different params.
|
||||
pub async fn build_index_with_config(
|
||||
&self,
|
||||
index_name: &str,
|
||||
embeddings: Vec<StoredEmbedding>,
|
||||
config: &HnswConfig,
|
||||
) -> Result<BuildStats, String> {
|
||||
let n = embeddings.len();
|
||||
if n == 0 {
|
||||
return Err("no embeddings to index".into());
|
||||
}
|
||||
|
||||
tracing::info!("building HNSW index '{}' from {} vectors", index_name, n);
|
||||
tracing::info!(
|
||||
"building HNSW '{}': {} vectors, ef_construction={} ef_search={} seed={:?}",
|
||||
index_name, n, config.ef_construction, config.ef_search, config.seed,
|
||||
);
|
||||
let start = std::time::Instant::now();
|
||||
|
||||
// Separate points and metadata
|
||||
@@ -92,14 +107,17 @@ impl HnswStore {
|
||||
values.push(i);
|
||||
}
|
||||
|
||||
// Build HNSW — this is the expensive part
|
||||
let map = Builder::default()
|
||||
.ef_construction(40) // balanced for 100K scale
|
||||
.ef_search(50) // fast search with good recall
|
||||
.build(points, values);
|
||||
// Build HNSW — the expensive part
|
||||
let mut builder = Builder::default()
|
||||
.ef_construction(config.ef_construction)
|
||||
.ef_search(config.ef_search);
|
||||
if let Some(seed) = config.seed {
|
||||
builder = builder.seed(seed);
|
||||
}
|
||||
let map = builder.build(points, values);
|
||||
|
||||
let build_time = start.elapsed().as_secs_f32();
|
||||
tracing::info!("HNSW index '{}' built: {} vectors in {:.1}s", index_name, n, build_time);
|
||||
tracing::info!("HNSW '{}' built: {} vectors in {:.1}s", index_name, n, build_time);
|
||||
|
||||
let index = Arc::new(HnswIndex { map, metadata });
|
||||
self.indexes.write().await.insert(index_name.to_string(), index);
|
||||
@@ -111,6 +129,43 @@ impl HnswStore {
|
||||
})
|
||||
}
|
||||
|
||||
/// Run a batch of search queries and return raw per-query latencies in microseconds.
|
||||
/// Also returns the retrieved doc_ids per query (for recall calculation).
|
||||
pub async fn bench_search(
|
||||
&self,
|
||||
index_name: &str,
|
||||
query_vectors: &[Vec<f32>],
|
||||
top_k: usize,
|
||||
) -> Result<BenchResult, String> {
|
||||
let indexes = self.indexes.read().await;
|
||||
let index = indexes
|
||||
.get(index_name)
|
||||
.ok_or_else(|| format!("HNSW index not found: {index_name}"))?
|
||||
.clone();
|
||||
drop(indexes);
|
||||
|
||||
let mut latencies_us = Vec::with_capacity(query_vectors.len());
|
||||
let mut retrieved: Vec<Vec<String>> = Vec::with_capacity(query_vectors.len());
|
||||
|
||||
for qv in query_vectors {
|
||||
let query_point = VecPoint(qv.clone());
|
||||
let t0 = std::time::Instant::now();
|
||||
let mut search = Search::default();
|
||||
let results = index.map.search(&query_point, &mut search);
|
||||
let ids: Vec<String> = results
|
||||
.take(top_k)
|
||||
.map(|item| {
|
||||
let meta_idx = *item.value;
|
||||
index.metadata[meta_idx].doc_id.clone()
|
||||
})
|
||||
.collect();
|
||||
latencies_us.push(t0.elapsed().as_micros() as f32);
|
||||
retrieved.push(ids);
|
||||
}
|
||||
|
||||
Ok(BenchResult { latencies_us, retrieved })
|
||||
}
|
||||
|
||||
/// Search an HNSW index. Returns approximate nearest neighbors.
|
||||
pub async fn search(
|
||||
&self,
|
||||
@@ -166,3 +221,9 @@ pub struct BuildStats {
|
||||
pub vectors: usize,
|
||||
pub build_time_secs: f32,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BenchResult {
|
||||
pub latencies_us: Vec<f32>,
|
||||
pub retrieved: Vec<Vec<String>>,
|
||||
}
|
||||
|
||||
@@ -1,4 +1,6 @@
|
||||
pub mod chunker;
|
||||
pub mod embedding_cache;
|
||||
pub mod harness;
|
||||
pub mod hnsw;
|
||||
pub mod index_registry;
|
||||
pub mod jobs;
|
||||
@@ -7,3 +9,4 @@ pub mod search;
|
||||
pub mod rag;
|
||||
pub mod supervisor;
|
||||
pub mod service;
|
||||
pub mod trial;
|
||||
|
||||
+304
-10
@@ -10,7 +10,7 @@ use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
|
||||
use aibridge::client::{AiClient, EmbedRequest};
|
||||
use crate::{chunker, hnsw, index_registry, jobs, rag, search, store, supervisor};
|
||||
use crate::{chunker, embedding_cache, harness, hnsw, index_registry, jobs, rag, search, store, supervisor, trial};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct VectorState {
|
||||
@@ -19,6 +19,8 @@ pub struct VectorState {
|
||||
pub job_tracker: jobs::JobTracker,
|
||||
pub index_registry: index_registry::IndexRegistry,
|
||||
pub hnsw_store: hnsw::HnswStore,
|
||||
pub embedding_cache: embedding_cache::EmbeddingCache,
|
||||
pub trial_journal: trial::TrialJournal,
|
||||
}
|
||||
|
||||
pub fn router(state: VectorState) -> Router {
|
||||
@@ -34,6 +36,17 @@ pub fn router(state: VectorState) -> Router {
|
||||
.route("/hnsw/build", post(build_hnsw))
|
||||
.route("/hnsw/search", post(search_hnsw))
|
||||
.route("/hnsw/list", get(list_hnsw))
|
||||
// Trial system — parameterized tuning loop
|
||||
.route("/hnsw/trial", post(run_trial))
|
||||
.route("/hnsw/trials/{index_name}", get(list_trials))
|
||||
.route("/hnsw/trials/{index_name}/best", get(best_trial))
|
||||
// Eval sets
|
||||
.route("/hnsw/evals", get(list_evals))
|
||||
.route("/hnsw/evals/{name}", get(get_eval).put(put_eval))
|
||||
.route("/hnsw/evals/{name}/autogen", post(autogen_eval))
|
||||
// Cache management
|
||||
.route("/hnsw/cache/stats", get(cache_stats))
|
||||
.route("/hnsw/cache/{index_name}", axum::routing::delete(cache_evict))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
@@ -307,26 +320,35 @@ async fn rag_query(
|
||||
struct BuildHnswRequest {
|
||||
/// Name of the stored vector index to build HNSW from
|
||||
index_name: String,
|
||||
/// Optional config override. Omit to use the production default
|
||||
/// (ec=80 es=30 — see HnswConfig::default docs for rationale).
|
||||
#[serde(default)]
|
||||
config: Option<trial::HnswConfig>,
|
||||
}
|
||||
|
||||
/// Build an HNSW index from an existing stored vector index.
|
||||
/// Loads embeddings from Parquet, builds HNSW in memory.
|
||||
/// Uses the embedding cache so repeated builds don't reload from Parquet.
|
||||
async fn build_hnsw(
|
||||
State(state): State<VectorState>,
|
||||
Json(req): Json<BuildHnswRequest>,
|
||||
) -> impl IntoResponse {
|
||||
tracing::info!("building HNSW for '{}'", req.index_name);
|
||||
let config = req.config.unwrap_or_default();
|
||||
tracing::info!(
|
||||
"building HNSW for '{}' ef_construction={} ef_search={}",
|
||||
req.index_name, config.ef_construction, config.ef_search,
|
||||
);
|
||||
|
||||
// Load embeddings from Parquet
|
||||
let embeddings = store::load_embeddings(&state.store, &req.index_name)
|
||||
let embeddings = state
|
||||
.embedding_cache
|
||||
.get_or_load(&req.index_name)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("index not found: {e}")))?;
|
||||
|
||||
let n = embeddings.len();
|
||||
tracing::info!("loaded {} embeddings, building HNSW...", n);
|
||||
|
||||
// Build HNSW
|
||||
match state.hnsw_store.build_index(&req.index_name, embeddings).await {
|
||||
match state
|
||||
.hnsw_store
|
||||
.build_index_with_config(&req.index_name, (*embeddings).clone(), &config)
|
||||
.await
|
||||
{
|
||||
Ok(stats) => Ok(Json(stats)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
@@ -372,3 +394,275 @@ async fn search_hnsw(
|
||||
async fn list_hnsw(State(state): State<VectorState>) -> impl IntoResponse {
|
||||
Json(state.hnsw_store.list().await)
|
||||
}
|
||||
|
||||
// --- Trial System: parameterized HNSW tuning loop ---
|
||||
//
|
||||
// Flow:
|
||||
// 1. Agent picks an HnswConfig
|
||||
// 2. POST /hnsw/trial builds HNSW with that config against cached embeddings,
|
||||
// runs every query in the harness, measures latency + recall vs the
|
||||
// harness's ground truth, appends a Trial record to _hnsw_trials/{idx}.jsonl
|
||||
// 3. Agent reads GET /hnsw/trials/{index}, sees history, decides next config
|
||||
// 4. Repeat until converged.
|
||||
//
|
||||
// The first trial triggers embedding load (slow). Every subsequent trial reuses
|
||||
// the cache — so the agent iterates in seconds, not minutes.
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TrialRequest {
|
||||
index_name: String,
|
||||
harness: String,
|
||||
#[serde(default)]
|
||||
config: trial::HnswConfig,
|
||||
#[serde(default)]
|
||||
note: Option<String>,
|
||||
}
|
||||
|
||||
async fn run_trial(
|
||||
State(state): State<VectorState>,
|
||||
Json(req): Json<TrialRequest>,
|
||||
) -> Result<Json<trial::Trial>, (StatusCode, String)> {
|
||||
let mut harness_set = harness::EvalSet::load(&state.store, &req.harness)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("harness not found: {e}")))?;
|
||||
|
||||
if harness_set.index_name != req.index_name {
|
||||
return Err((
|
||||
StatusCode::BAD_REQUEST,
|
||||
format!(
|
||||
"harness '{}' is for index '{}', not '{}'",
|
||||
req.harness, harness_set.index_name, req.index_name
|
||||
),
|
||||
));
|
||||
}
|
||||
if harness_set.queries.is_empty() {
|
||||
return Err((StatusCode::BAD_REQUEST, "harness has no queries".into()));
|
||||
}
|
||||
|
||||
let embeddings = state
|
||||
.embedding_cache
|
||||
.get_or_load(&req.index_name)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("load embeddings: {e}")))?;
|
||||
|
||||
if !harness_set.ground_truth_built {
|
||||
tracing::info!("trial: computing ground truth for harness '{}'", harness_set.name);
|
||||
let t0 = std::time::Instant::now();
|
||||
harness::compute_ground_truth(&mut harness_set, &embeddings, &state.ai_client)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("ground truth: {e}")))?;
|
||||
tracing::info!("trial: ground truth built in {:.1}s", t0.elapsed().as_secs_f32());
|
||||
harness_set
|
||||
.save(&state.store)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("save harness: {e}")))?;
|
||||
}
|
||||
|
||||
let trial_id = trial::Trial::new_id();
|
||||
let hnsw_slot = format!("{}__{}", req.index_name, trial_id);
|
||||
|
||||
let build_stats = state
|
||||
.hnsw_store
|
||||
.build_index_with_config(&hnsw_slot, (*embeddings).clone(), &req.config)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("build: {e}")))?;
|
||||
|
||||
let query_vectors: Vec<Vec<f32>> = harness_set
|
||||
.queries
|
||||
.iter()
|
||||
.filter_map(|q| q.query_embedding.clone())
|
||||
.collect();
|
||||
let bench = state
|
||||
.hnsw_store
|
||||
.bench_search(&hnsw_slot, &query_vectors, harness_set.k)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("search: {e}")))?;
|
||||
|
||||
let mut recalls = Vec::with_capacity(harness_set.queries.len());
|
||||
for (q, hits) in harness_set.queries.iter().zip(bench.retrieved.iter()) {
|
||||
if let Some(gt) = &q.ground_truth {
|
||||
recalls.push(harness::recall_at_k(hits, gt, harness_set.k));
|
||||
}
|
||||
}
|
||||
let mean_recall = if recalls.is_empty() {
|
||||
0.0
|
||||
} else {
|
||||
recalls.iter().sum::<f32>() / recalls.len() as f32
|
||||
};
|
||||
|
||||
let mut lats = bench.latencies_us.clone();
|
||||
lats.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
|
||||
let p = |pct: f32| -> f32 {
|
||||
if lats.is_empty() { return 0.0; }
|
||||
let idx = ((lats.len() as f32 - 1.0) * pct).round() as usize;
|
||||
lats[idx.min(lats.len() - 1)]
|
||||
};
|
||||
|
||||
// One brute-force reference latency — keeps the cost proportional to
|
||||
// whatever the agent is willing to pay per trial.
|
||||
let brute_latency_us = if let Some(qv) = query_vectors.first() {
|
||||
let t0 = std::time::Instant::now();
|
||||
let _ = harness::brute_force_top_k(qv, &embeddings, harness_set.k);
|
||||
t0.elapsed().as_micros() as f32
|
||||
} else {
|
||||
0.0
|
||||
};
|
||||
|
||||
let dims = embeddings.first().map(|e| e.vector.len()).unwrap_or(0);
|
||||
let memory_bytes =
|
||||
(embeddings.len() * dims * std::mem::size_of::<f32>() + embeddings.len() * 128) as u64;
|
||||
|
||||
let trial_record = trial::Trial {
|
||||
id: trial_id.clone(),
|
||||
index_name: req.index_name.clone(),
|
||||
eval_set: req.harness.clone(),
|
||||
config: req.config.clone(),
|
||||
metrics: trial::TrialMetrics {
|
||||
build_time_secs: build_stats.build_time_secs,
|
||||
search_latency_p50_us: p(0.50),
|
||||
search_latency_p95_us: p(0.95),
|
||||
search_latency_p99_us: p(0.99),
|
||||
recall_at_k: mean_recall,
|
||||
memory_bytes,
|
||||
vectors: build_stats.vectors,
|
||||
eval_queries: harness_set.queries.len(),
|
||||
brute_force_latency_us: brute_latency_us,
|
||||
},
|
||||
created_at: chrono::Utc::now(),
|
||||
note: req.note,
|
||||
};
|
||||
|
||||
state
|
||||
.trial_journal
|
||||
.append(&trial_record)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("journal: {e}")))?;
|
||||
|
||||
state.hnsw_store.drop(&hnsw_slot).await;
|
||||
|
||||
Ok(Json(trial_record))
|
||||
}
|
||||
|
||||
async fn list_trials(
|
||||
State(state): State<VectorState>,
|
||||
Path(index_name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match state.trial_journal.list(&index_name).await {
|
||||
Ok(trials) => Ok(Json(trials)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct BestTrialQuery {
|
||||
#[serde(default = "default_metric")]
|
||||
metric: String,
|
||||
}
|
||||
|
||||
fn default_metric() -> String {
|
||||
"pareto".to_string()
|
||||
}
|
||||
|
||||
async fn best_trial(
|
||||
State(state): State<VectorState>,
|
||||
Path(index_name): Path<String>,
|
||||
Query(q): Query<BestTrialQuery>,
|
||||
) -> impl IntoResponse {
|
||||
match state.trial_journal.best(&index_name, &q.metric).await {
|
||||
Ok(Some(t)) => Ok(Json(t)),
|
||||
Ok(None) => Err((StatusCode::NOT_FOUND, "no trials yet".to_string())),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
// --- Harness management ---
|
||||
|
||||
async fn list_evals(State(state): State<VectorState>) -> impl IntoResponse {
|
||||
match harness::EvalSet::list(&state.store).await {
|
||||
Ok(names) => Ok(Json(names)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_eval(
|
||||
State(state): State<VectorState>,
|
||||
Path(name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
match harness::EvalSet::load(&state.store, &name).await {
|
||||
Ok(e) => Ok(Json(e)),
|
||||
Err(err) => Err((StatusCode::NOT_FOUND, err)),
|
||||
}
|
||||
}
|
||||
|
||||
async fn put_eval(
|
||||
State(state): State<VectorState>,
|
||||
Path(name): Path<String>,
|
||||
Json(mut harness_set): Json<harness::EvalSet>,
|
||||
) -> impl IntoResponse {
|
||||
harness_set.name = name;
|
||||
harness_set.ground_truth_built = harness_set
|
||||
.queries
|
||||
.iter()
|
||||
.all(|q| q.ground_truth.is_some());
|
||||
match harness_set.save(&state.store).await {
|
||||
Ok(()) => Ok(Json(harness_set)),
|
||||
Err(e) => Err((StatusCode::INTERNAL_SERVER_ERROR, e)),
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct AutogenRequest {
|
||||
index_name: String,
|
||||
#[serde(default = "default_sample_count")]
|
||||
sample_count: usize,
|
||||
#[serde(default = "default_k")]
|
||||
k: usize,
|
||||
}
|
||||
|
||||
fn default_sample_count() -> usize { 100 }
|
||||
fn default_k() -> usize { 10 }
|
||||
|
||||
async fn autogen_eval(
|
||||
State(state): State<VectorState>,
|
||||
Path(name): Path<String>,
|
||||
Json(req): Json<AutogenRequest>,
|
||||
) -> Result<Json<harness::EvalSet>, (StatusCode, String)> {
|
||||
let embeddings = state
|
||||
.embedding_cache
|
||||
.get_or_load(&req.index_name)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::NOT_FOUND, format!("load embeddings: {e}")))?;
|
||||
|
||||
let mut harness_set = harness::synthetic_from_chunks(
|
||||
&name,
|
||||
&req.index_name,
|
||||
&embeddings,
|
||||
req.sample_count,
|
||||
req.k,
|
||||
);
|
||||
|
||||
harness::compute_ground_truth(&mut harness_set, &embeddings, &state.ai_client)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("ground truth: {e}")))?;
|
||||
|
||||
harness_set
|
||||
.save(&state.store)
|
||||
.await
|
||||
.map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("save: {e}")))?;
|
||||
|
||||
Ok(Json(harness_set))
|
||||
}
|
||||
|
||||
// --- Embedding cache management ---
|
||||
|
||||
async fn cache_stats(State(state): State<VectorState>) -> impl IntoResponse {
|
||||
Json(state.embedding_cache.stats().await)
|
||||
}
|
||||
|
||||
async fn cache_evict(
|
||||
State(state): State<VectorState>,
|
||||
Path(index_name): Path<String>,
|
||||
) -> impl IntoResponse {
|
||||
let ok = state.embedding_cache.evict(&index_name).await;
|
||||
Json(serde_json::json!({ "evicted": ok, "index_name": index_name }))
|
||||
}
|
||||
|
||||
@@ -0,0 +1,213 @@
|
||||
/// Trial journal for HNSW parameter tuning.
|
||||
///
|
||||
/// Every HNSW build+eval is recorded as a Trial. The journal is append-only
|
||||
/// and stored under `_hnsw_trials/{index_name}/` as batched JSONL files —
|
||||
/// an AI agent iterating on configs reads prior trials to decide what to
|
||||
/// try next, and writes a new trial on each attempt.
|
||||
///
|
||||
/// Storage uses the shared `storaged::append_log::AppendLog` so appends are
|
||||
/// write-once (new file per batch) rather than rewriting a single growing
|
||||
/// JSONL on every event. See `append_log.rs` for the full rationale.
|
||||
|
||||
use chrono::{DateTime, Utc};
|
||||
use object_store::ObjectStore;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
|
||||
use storaged::append_log::AppendLog;
|
||||
|
||||
/// HNSW build/search parameters the agent can tune.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct HnswConfig {
|
||||
pub ef_construction: usize,
|
||||
pub ef_search: usize,
|
||||
#[serde(default)]
|
||||
pub seed: Option<u64>,
|
||||
}
|
||||
|
||||
impl Default for HnswConfig {
|
||||
/// Production default, locked in 2026-04-16 based on trial grid against
|
||||
/// resumes_100k_v2 (100K vectors, 20 queries, recall@10):
|
||||
/// ec=20 es=30 → recall 0.960, p50 509us, build 8s
|
||||
/// ec=80 es=30 → recall 1.000, p50 873us, build 230s ← sweet spot
|
||||
/// ec=200 es=30 → recall 1.000, p50 874us, build 106s (no recall gain)
|
||||
///
|
||||
/// `ec=80` is the smallest value that reaches 100% recall. Higher values
|
||||
/// waste build time. `es=30` gives faster search than `es=100` with no
|
||||
/// recall penalty at this scale.
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
ef_construction: 80,
|
||||
ef_search: 30,
|
||||
seed: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Metrics collected on every trial. All latencies in microseconds.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct TrialMetrics {
|
||||
pub build_time_secs: f32,
|
||||
pub search_latency_p50_us: f32,
|
||||
pub search_latency_p95_us: f32,
|
||||
pub search_latency_p99_us: f32,
|
||||
pub recall_at_k: f32,
|
||||
pub memory_bytes: u64,
|
||||
pub vectors: usize,
|
||||
pub eval_queries: usize,
|
||||
/// Brute-force latency for comparison — how much speedup did HNSW buy us?
|
||||
pub brute_force_latency_us: f32,
|
||||
}
|
||||
|
||||
/// A single tuning attempt.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Trial {
|
||||
pub id: String,
|
||||
pub index_name: String,
|
||||
pub eval_set: String,
|
||||
pub config: HnswConfig,
|
||||
pub metrics: TrialMetrics,
|
||||
pub created_at: DateTime<Utc>,
|
||||
/// Free-form note — the agent can record why it tried this config.
|
||||
#[serde(default)]
|
||||
pub note: Option<String>,
|
||||
}
|
||||
|
||||
impl Trial {
|
||||
pub fn new_id() -> String {
|
||||
format!(
|
||||
"trial-{}-{}",
|
||||
Utc::now().timestamp_millis(),
|
||||
&uuid::Uuid::new_v4().to_string()[..8]
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Per-index append log, lazy-created on first write.
|
||||
#[derive(Clone)]
|
||||
pub struct TrialJournal {
|
||||
store: Arc<dyn ObjectStore>,
|
||||
/// Cache per-index AppendLog instances so the in-memory buffer persists
|
||||
/// across calls.
|
||||
logs: Arc<RwLock<HashMap<String, Arc<AppendLog>>>>,
|
||||
}
|
||||
|
||||
impl TrialJournal {
|
||||
pub fn new(store: Arc<dyn ObjectStore>) -> Self {
|
||||
Self {
|
||||
store,
|
||||
logs: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
fn prefix(index_name: &str) -> String {
|
||||
format!("_hnsw_trials/{}", index_name)
|
||||
}
|
||||
|
||||
async fn log_for(&self, index_name: &str) -> Arc<AppendLog> {
|
||||
if let Some(log) = self.logs.read().await.get(index_name) {
|
||||
return log.clone();
|
||||
}
|
||||
let mut guard = self.logs.write().await;
|
||||
if let Some(log) = guard.get(index_name) {
|
||||
return log.clone();
|
||||
}
|
||||
// Trials arrive one at a time during human/agent iteration — a low
|
||||
// threshold gives "hit /trials and see my latest attempt" immediacy
|
||||
// without creating one file per event.
|
||||
let log = Arc::new(
|
||||
AppendLog::new(self.store.clone(), Self::prefix(index_name))
|
||||
.with_flush_threshold(4),
|
||||
);
|
||||
guard.insert(index_name.to_string(), log.clone());
|
||||
log
|
||||
}
|
||||
|
||||
/// Append a trial record. In-memory buffered; persisted in batches.
|
||||
pub async fn append(&self, trial: &Trial) -> Result<(), String> {
|
||||
let line = serde_json::to_vec(trial).map_err(|e| e.to_string())?;
|
||||
let log = self.log_for(&trial.index_name).await;
|
||||
log.append(line).await
|
||||
}
|
||||
|
||||
/// Read all trials for an index (flushed batches + unflushed buffer).
|
||||
pub async fn list(&self, index_name: &str) -> Result<Vec<Trial>, String> {
|
||||
let log = self.log_for(index_name).await;
|
||||
let lines = log.read_all().await?;
|
||||
let mut trials = Vec::with_capacity(lines.len());
|
||||
for line in lines {
|
||||
match serde_json::from_slice::<Trial>(&line) {
|
||||
Ok(t) => trials.push(t),
|
||||
Err(e) => tracing::warn!("trial journal: skip malformed line: {e}"),
|
||||
}
|
||||
}
|
||||
Ok(trials)
|
||||
}
|
||||
|
||||
/// Explicit flush for callers that want write-through semantics
|
||||
/// (e.g. an agent that wants to commit a trial before querying stats).
|
||||
pub async fn flush(&self, index_name: &str) -> Result<(), String> {
|
||||
let log = self.log_for(index_name).await;
|
||||
log.flush().await
|
||||
}
|
||||
|
||||
/// Compact all batch files for an index into one.
|
||||
pub async fn compact(&self, index_name: &str) -> Result<storaged::append_log::CompactStats, String> {
|
||||
let log = self.log_for(index_name).await;
|
||||
log.compact().await
|
||||
}
|
||||
|
||||
/// Current champion for an index by the named metric.
|
||||
/// Valid metrics: `recall`, `latency`, `pareto`.
|
||||
///
|
||||
/// The `pareto` strategy is a placeholder — J should tune the scoring
|
||||
/// function to match what matters in production. Right now it's a simple
|
||||
/// weighted sum.
|
||||
pub async fn best(
|
||||
&self,
|
||||
index_name: &str,
|
||||
metric: &str,
|
||||
) -> Result<Option<Trial>, String> {
|
||||
let trials = self.list(index_name).await?;
|
||||
if trials.is_empty() {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
let best = match metric {
|
||||
"recall" => trials
|
||||
.into_iter()
|
||||
.max_by(|a, b| {
|
||||
a.metrics
|
||||
.recall_at_k
|
||||
.partial_cmp(&b.metrics.recall_at_k)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.unwrap(),
|
||||
"latency" => trials
|
||||
.into_iter()
|
||||
.min_by(|a, b| {
|
||||
a.metrics
|
||||
.search_latency_p95_us
|
||||
.partial_cmp(&b.metrics.search_latency_p95_us)
|
||||
.unwrap_or(std::cmp::Ordering::Equal)
|
||||
})
|
||||
.unwrap(),
|
||||
"pareto" | _ => trials
|
||||
.into_iter()
|
||||
.max_by(|a, b| pareto_score(a).partial_cmp(&pareto_score(b)).unwrap())
|
||||
.unwrap(),
|
||||
};
|
||||
Ok(Some(best))
|
||||
}
|
||||
}
|
||||
|
||||
/// Simple Pareto-style score: reward recall, penalize p95 latency.
|
||||
/// Tunable — J should swap this in production to match what matters.
|
||||
fn pareto_score(t: &Trial) -> f32 {
|
||||
// Recall is [0, 1]. Latency is us — assume 100us baseline.
|
||||
let recall = t.metrics.recall_at_k;
|
||||
let latency_penalty = (t.metrics.search_latency_p95_us / 1000.0).min(1.0); // cap at 1ms
|
||||
recall - 0.2 * latency_penalty
|
||||
}
|
||||
Reference in New Issue
Block a user