Skip to content

improve #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
andrewgazelka wants to merge 2 commits into
base: main
Choose a base branch
from
Open

improve #10

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9add56b
use correct build
andrewgazelka Nov 17, 2025
9cc0bca
much more idiomatic
andrewgazelka Nov 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
731 changes: 242 additions & 489 deletions Cargo.lock
View file
Open in desktop

Large diffs are not rendered by default.

13 changes: 6 additions & 7 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,21 +3,20 @@ name = "haystackdb"
version = "0.1.0"
edition = "2021"

[build]
target = "x86_64-unknown-linux-gnu"
flags = ["-C", "target-cpu=native"]

[dependencies]
warp = "0.3"
axum = "0.7"
tokio = { version = "1", features = ["full"] }
tower = "0.4"
tower-http = { version = "0.5", features = ["cors"] }
serde = { version = "1.0", features = ["derive"] }
rayon = "1.10.0"
uuid = { version = "1.8.0", features = ["v4", "serde"] }
memmap = "0.7.0"
log = "0.4.14"
tracing = "0.1"
fs2 = "0.4.0"
env_logger = "0.11.3"
tracing-subscriber = { version = "0.3", features = ["env-filter", "fmt"] }
serde_json = "1.0.68"
parking_lot = "0.12"

[profile.release]
opt-level = 3
Expand Down
3 changes: 2 additions & 1 deletion src/constants.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
pub const VECTOR_SIZE: usize = 1024;
pub const QUANTIZED_VECTOR_SIZE: usize = VECTOR_SIZE / 8;
pub const BITS_PER_ELEMENT: usize = 8;
pub const QUANTIZED_VECTOR_SIZE: usize = VECTOR_SIZE / BITS_PER_ELEMENT;
225 changes: 98 additions & 127 deletions src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,169 +1,140 @@
use env_logger::Builder;
use haystackdb::constants::VECTOR_SIZE;
use haystackdb::services::CommitService;
use haystackdb::services::QueryService;
use haystackdb::structures::filters::Filter as QueryFilter;
use haystackdb::structures::metadata_index::KVPair;
use log::info;
use log::LevelFilter;
use std::io::Write;
use std::sync::{Arc, Mutex};
use parking_lot::Mutex;
use std::sync::Arc;
use std::{self, path::PathBuf};
use tokio::time::{interval, Duration};
use tracing::info;

use axum::{
extract::{Path, State},
routing::{get, post},
Json, Router,
};
use std::collections::HashMap;
use tokio::sync::OnceCell;
use warp::Filter;

static ACTIVE_NAMESPACES: OnceCell<Arc<Mutex<HashMap<String, tokio::task::JoinHandle<()>>>>> =
OnceCell::const_new();
type AppState = Arc<Mutex<HashMap<String, tokio::task::JoinHandle<()>>>>;

#[tokio::main]
async fn main() {
let mut builder = Builder::new();
builder
.format(|buf, record| writeln!(buf, "{}: {}", record.level(), record.args()))
.filter(None, LevelFilter::Info)
tracing_subscriber::fmt()
.with_env_filter(
tracing_subscriber::EnvFilter::try_from_default_env()
.unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")),
)
.init();

let active_namespaces = ACTIVE_NAMESPACES
.get_or_init(|| async { Arc::new(Mutex::new(HashMap::new())) })
.await;

let search_route = warp::path!("query" / String)
.and(warp::post())
.and(warp::body::json())
.and(with_active_namespaces(active_namespaces.clone()))
.then(
|namespace_id: String, body: (Vec<f64>, QueryFilter, usize), active_namespaces| async move {
let base_path = PathBuf::from(format!("/workspace/data/{}/current", namespace_id.clone()));
ensure_namespace_initialized(&namespace_id, &active_namespaces, base_path.clone())
.await;

let mut query_service = QueryService::new(base_path, namespace_id.clone()).unwrap();
let fvec = &body.0;
let metadata = &body.1;
let top_k = body.2;

let mut vec: [f32; VECTOR_SIZE] = [0.0; VECTOR_SIZE];
fvec.iter()
.enumerate()
.for_each(|(i, &val)| vec[i] = val as f32);
let active_namespaces = Arc::new(Mutex::new(HashMap::new()));

let start = std::time::Instant::now();
let app = Router::new()
.route("/query/:namespace_id", post(query_handler))
.route("/addVector/:namespace_id", post(add_vector_handler))
.route("/pitr/:namespace_id/:timestamp", get(pitr_handler))
.with_state(active_namespaces);

let search_result = query_service
.query(&vec, metadata, top_k)
.expect("Failed to query");
let listener = tokio::net::TcpListener::bind("0.0.0.0:8080").await.unwrap();

let duration = start.elapsed();
info!("Server listening on 0.0.0.0:8080");
axum::serve(listener, app).await.unwrap();
}

async fn query_handler(
Path(namespace_id): Path<String>,
State(active_namespaces): State<AppState>,
Json(body): Json<(Vec<f64>, QueryFilter, usize)>,
) -> Json<serde_json::Value> {
let base_path = PathBuf::from(format!("/workspace/data/{}/current", namespace_id));
ensure_namespace_initialized(&namespace_id, &active_namespaces, base_path.clone()).await;

let mut query_service = QueryService::new(base_path, namespace_id.clone()).unwrap();
let fvec = &body.0;
let metadata = &body.1;
let top_k = body.2;

let mut vec: [f32; VECTOR_SIZE] = [0.0; VECTOR_SIZE];
fvec.iter()
.enumerate()
.for_each(|(i, &val)| vec[i] = val as f32);

let start = std::time::Instant::now();

let search_result = query_service
.query(&vec, metadata, top_k)
.expect("Failed to query");

let duration = start.elapsed();

info!(?duration, "Query completed");
Json(serde_json::to_value(search_result).unwrap())
}

async fn add_vector_handler(
Path(namespace_id): Path<String>,
State(active_namespaces): State<AppState>,
Json(body): Json<(Vec<f64>, Vec<KVPair>, String)>,
) -> Json<&'static str> {
let base_path = PathBuf::from(format!("/workspace/data/{}/current", namespace_id));

ensure_namespace_initialized(&namespace_id, &active_namespaces, base_path.clone()).await;

let mut commit_service = CommitService::new(base_path, namespace_id.clone()).unwrap();
let fvec = &body.0;
let metadata = &body.1;

let mut vec: [f32; VECTOR_SIZE] = [0.0; VECTOR_SIZE];
fvec.iter()
.enumerate()
.for_each(|(i, &val)| vec[i] = val as f32);

println!("Query took {:?} to complete", duration);
warp::reply::json(&search_result)
},
);

let add_vector_route =
warp::path!("addVector" / String)
.and(warp::post())
.and(warp::body::json())
.and(with_active_namespaces(active_namespaces.clone()))
.then(
|namespace_id: String,
body: (Vec<f64>, Vec<KVPair>, String),
active_namespaces| async move {
let base_path = PathBuf::from(format!(
"/workspace/data/{}/current",
namespace_id.clone()
));

ensure_namespace_initialized(
&namespace_id,
&active_namespaces,
base_path.clone(),
)
.await;

let mut commit_service =
CommitService::new(base_path, namespace_id.clone()).unwrap();
let fvec = &body.0;
let metadata = &body.1;

let mut vec: [f32; VECTOR_SIZE] = [0.0; VECTOR_SIZE];
fvec.iter()
.enumerate()
.for_each(|(i, &val)| vec[i] = val as f32);

// let id = uuid::Uuid::from_str(id_str).unwrap();
commit_service.add_to_wal(vec![vec], vec![metadata.clone()]).expect("Failed to add to WAL");
warp::reply::json(&"Success")
},
);

// add a PITR route
let pitr_route = warp::path!("pitr" / String / String)
.and(warp::get())
.and(with_active_namespaces(active_namespaces.clone()))
.then(
|namespace_id: String, timestamp: String, active_namespaces| async move {
println!("PITR for namespace: {}", namespace_id);
let base_path =
PathBuf::from(format!("/workspace/data/{}/current", namespace_id.clone()));

ensure_namespace_initialized(&namespace_id, &active_namespaces, base_path.clone())
.await;

let mut commit_service =
CommitService::new(base_path, namespace_id.clone()).unwrap();

let timestamp = timestamp.parse::<u64>().unwrap();
commit_service
.recover_point_in_time(timestamp)
.expect("Failed to PITR");
warp::reply::json(&"Success")
},
);

let routes = search_route
.or(add_vector_route)
.or(pitr_route)
.with(warp::cors().allow_any_origin());
warp::serve(routes).run(([0, 0, 0, 0], 8080)).await;
commit_service
.add_to_wal(vec![vec], vec![metadata.clone()])
.expect("Failed to add to WAL");

Json("Success")
}

fn with_active_namespaces(
active_namespaces: Arc<Mutex<HashMap<String, tokio::task::JoinHandle<()>>>>,
) -> impl Filter<
Extract = (Arc<Mutex<HashMap<String, tokio::task::JoinHandle<()>>>>,),
Error = std::convert::Infallible,
> + Clone {
warp::any().map(move || active_namespaces.clone())
async fn pitr_handler(
Path((namespace_id, timestamp)): Path<(String, String)>,
State(active_namespaces): State<AppState>,
) -> Json<&'static str> {
info!(namespace_id = %namespace_id, "Executing PITR");
let base_path = PathBuf::from(format!("/workspace/data/{}/current", namespace_id));

ensure_namespace_initialized(&namespace_id, &active_namespaces, base_path.clone()).await;

let mut commit_service = CommitService::new(base_path, namespace_id.clone()).unwrap();

let timestamp = timestamp.parse::<u64>().unwrap();
commit_service
.recover_point_in_time(timestamp)
.expect("Failed to PITR");

Json("Success")
}

async fn ensure_namespace_initialized(
namespace_id: &String,
active_namespaces: &Arc<Mutex<HashMap<String, tokio::task::JoinHandle<()>>>>,
base_path_for_async: PathBuf,
) {
let mut namespaces = active_namespaces.lock().unwrap();
let mut namespaces = active_namespaces.lock();
if !namespaces.contains_key(namespace_id) {
let namespace_id_cloned = namespace_id.clone();
let handle = tokio::spawn(async move {
let mut interval = interval(Duration::from_secs(10));
loop {
interval.tick().await;
println!("Committing for namespace {}", namespace_id_cloned);
info!(namespace_id = %namespace_id_cloned, "Starting commit");
let start = std::time::Instant::now();
let commit_worker = std::sync::Arc::new(std::sync::Mutex::new(
let commit_worker = Arc::new(Mutex::new(
CommitService::new(base_path_for_async.clone(), namespace_id_cloned.clone())
.unwrap(),
));

commit_worker
.lock()
.unwrap()
.commit()
.expect("Failed to commit");
commit_worker.lock().commit().expect("Failed to commit");
let duration = start.elapsed();
info!("Commit worker took {:?} to complete", duration);
}
Expand Down
Loading