Implement grpc server

Author: TanmayArya-1p

Cargo.lock

@@ -5,6 +5,7 @@ members = [
     "crates/storage",
     "crates/index",
     "crates/server",
+    "crates/grpc_server",
 ]
 
 # You can define shared dependencies for all crates here

Cargo.toml

@@ -19,15 +19,27 @@ fn generate_point_id() -> u64 {
 pub struct VectorDb {
     storage: Arc<dyn StorageEngine>,
     index: Arc<RwLock<dyn VectorIndex>>, // Using a RwLock instead of Mutex to improve concurrency
+    dimension: usize,
 }
 
 impl VectorDb {
-    fn _new(storage: Arc<dyn StorageEngine>, index: Arc<RwLock<dyn VectorIndex>>) -> Self {
-        Self { storage, index }
+    fn _new(
+        storage: Arc<dyn StorageEngine>,
+        index: Arc<RwLock<dyn VectorIndex>>,
+        dimension: usize,
+    ) -> Self {
+        Self {
+            storage,
+            index,
+            dimension,
+        }
     }
 
     //TODO: Make this an atomic operation
     pub fn insert(&self, vector: DenseVector, payload: Payload) -> Result<PointId, DbError> {
+        if vector.len() != self.dimension {
+            return Err(DbError::DimensionMismatch);
+        }
         // Generate a new point id
         let point_id = generate_point_id();
         self.storage
@@ -44,13 +56,13 @@ impl VectorDb {
     }
 
     //TODO: Make this an atomic operation
-    pub fn delete(&self, id: PointId) -> Result<(), DbError> {
+    pub fn delete(&self, id: PointId) -> Result<bool, DbError> {
         // Remove from storage
         self.storage.delete_point(id)?;
         // Remove from index
         let mut index = self.index.write().map_err(|_| DbError::LockError)?;
-        index.delete(id)?;
-        Ok(())
+        let point_found = index.delete(id)?;
+        Ok(point_found)
     }
 
     pub fn get(&self, id: PointId) -> Result<Option<Point>, DbError> {
@@ -105,7 +117,7 @@ pub fn init_api(config: DbConfig) -> Result<VectorDb, DbError> {
     };
 
     // Init the db
-    let db = VectorDb::_new(storage, index);
+    let db = VectorDb::_new(storage, index, config.dimension);
 
     Ok(db)
 }
@@ -147,6 +159,22 @@ mod tests {
         assert_eq!(point.payload.as_ref().unwrap(), &payload);
     }
 
+    #[test]
+    fn test_dimension_mismatch() {
+        let db = create_test_db();
+        let v1 = vec![1.0, 2.0, 3.0];
+        let v2 = vec![1.0, 2.0];
+        let payload = Payload {};
+
+        let res1 = db.insert(v1, payload);
+        assert!(res1.is_ok());
+
+        // Insert vector of dimension 2 != 3
+        let res2 = db.insert(v2, payload);
+        assert!(res2.is_err());
+        assert_eq!(res2.unwrap_err(), DbError::DimensionMismatch);
+    }
+
     #[test]
     fn test_delete() {
         let db = create_test_db();
@@ -156,6 +184,12 @@ mod tests {
         // Insert a point
         let id = db.insert(vector, payload).unwrap();
 
+        // try deleting a point that does not exist
+        let found = db.delete(id + 1);
+        assert!(found.is_ok());
+        assert!(!found.unwrap());
+
+        // delete the point
         assert!(db.get(id).unwrap().is_some());
         db.delete(id).unwrap();
         assert!(db.get(id).unwrap().is_none());

crates/api/src/lib.rs

@@ -1,9 +1,18 @@
-#[derive(Debug)]
+#[derive(Debug, PartialEq, Eq)]
 pub enum DbError {
     ParseError,
     StorageError(String),
     SerializationError(String),
     DeserializationError,
     IndexError(String),
     LockError,
+    DimensionMismatch,
 }
+
+impl std::fmt::Display for DbError {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "{:?}", self)
+    }
+}
+
+impl std::error::Error for DbError {}

crates/defs/src/error.rs

@@ -0,0 +1,9 @@
+GRPC_SERVER_ROOT_PASSWORD=123 # required
+GRPC_SERVER_DIMENSION=3 # required
+
+GRPC_SERVER_HOST=localhost # defaults to 127.0.0.1 aka localhost
+GRPC_SERVER_PORT=8080 # defaults to 8080
+GRPC_SERVER_STORAGE_TYPE=inmemory # (inmemory/rocksdb) defaults to 'inmemory'
+GRPC_SERVER_INDEX_TYPE=flat # defaults to flat
+GRPC_SERVER_DATA_PATH=data # defaults to a temporary directory
+GRPC_SERVER_LOGGING=true # defaults to true

crates/grpc_server/.sample.env

@@ -0,0 +1,27 @@
+[package]
+name = "grpc_server"
+version = "0.1.0"
+edition = "2024"
+
+[dependencies]
+prost = "0.14.1"
+tokio = { version = "1.47.1", features = ["macros", "rt-multi-thread"] }
+tonic = "0.14.2"
+tonic-prost-build = "0.14.2"
+
+api = { path = "../api" }
+storage = { path = "../storage" }
+index = { path = "../index" }
+defs = { path = "../defs" }
+
+tonic-prost = "0.14.2"
+prost-types = "0.14.1"
+dotenv = "0.15.0"
+tempfile = "3.23.0"
+tracing = "0.1.41"
+tracing-subscriber = "0.3.20"
+tokio-stream = "0.1.17"
+
+[build-dependencies]
+tonic-build = "0.14.2"
+tonic-prost-build = "0.14.2"

crates/grpc_server/Cargo.toml

@@ -0,0 +1,32 @@
+### Build
+
+Clone the repository and run `cargo build --bin grpc_server` to build the binary of the gRPC server crate.
+
+You can than then start the gRPC server by running:
+
+```bash
+cargo run --bin grpc_server
+```
+
+### Configuration
+
+Use the [.sample.env](.sample.env) shown below as a reference to set your environment variables in a `.env` file.
+
+```bash
+GRPC_SERVER_ROOT_PASSWORD=123 # required
+GRPC_SERVER_DIMENSION=3 # required
+
+GRPC_SERVER_HOST=localhost # defaults to 127.0.0.1 aka localhost
+GRPC_SERVER_PORT=8080 # defaults to 8080
+GRPC_SERVER_STORAGE_TYPE=inmemory # (inmemory/rocksdb) defaults to 'inmemory'
+GRPC_SERVER_INDEX_TYPE=flat # defaults to flat
+GRPC_SERVER_DATA_PATH=data # defaults to a temporary directory
+GRPC_SERVER_LOGGING=true # defaults to true
+
+
+```
+
+
+### Testing
+
+The [vector-db.proto](proto/vector-db.proto) can be imported into any gRPC client.

crates/grpc_server/README.md

@@ -0,0 +1,4 @@
+fn main() -> Result<(), Box<dyn std::error::Error>> {
+    tonic_prost_build::compile_protos("proto/vector-db.proto")?;
+    Ok(())
+}

crates/grpc_server/build.rs

@@ -0,0 +1,60 @@
+syntax = "proto3";
+
+package vectordb;
+
+import "google/protobuf/struct.proto";
+import "google/protobuf/empty.proto";
+
+
+service VectorDB {
+  //Insert a vector with a payload and return the assigned PointID
+  rpc InsertVector(InsertVectorRequest) returns (PointID) {}
+
+  //Delete a vector by its PointID
+  rpc DeletePoint(PointID) returns (google.protobuf.Empty) {}
+
+  //Get a vector and its payload by PointID
+  rpc GetPoint(PointID) returns (Point) {}
+
+  //Search for the k nearest vectors to a target vector given a distance function
+  rpc SearchPoints(SearchRequest) returns (SearchResponse) {}
+}
+
+
+message InsertVectorRequest {
+  DenseVector vector = 1;
+  optional google.protobuf.Struct payload = 2;
+}
+
+
+message SearchRequest {
+  DenseVector query_vector = 1;
+  Similarity similarity = 2;
+  int64 limit = 3;
+}
+
+
+message SearchResponse {
+  repeated PointID result_point_ids = 1;
+}
+
+message DenseVector {
+  repeated float values = 1;
+}
+
+message Point {
+  PointID id = 1;
+  optional google.protobuf.Struct payload = 2;
+  DenseVector vector = 3;
+}
+
+message PointID {
+  uint64 id = 1;
+}
+
+enum Similarity{
+  Euclidean = 0;
+  Manhattan = 1;
+  Hamming = 2;
+  Cosine = 3;
+}

crates/grpc_server/proto/vector-db.proto

@@ -0,0 +1,125 @@
+use crate::constants::{
+    self, DEFAULT_PORT, ENV_DATA_PATH, ENV_DIMENSION, ENV_INDEX_TYPE, ENV_LOGGING, ENV_PORT,
+    ENV_ROOT_PASSWORD, ENV_STORAGE_TYPE,
+};
+use crate::errors;
+use api;
+use dotenv::dotenv;
+use index::IndexType;
+use std::net::SocketAddr;
+use std::path::PathBuf;
+use std::{env, fs};
+use storage;
+use tempfile::tempdir;
+use tracing::{Level, event};
+
+pub struct GRPCServerConfig {
+    pub addr: SocketAddr,
+    pub root_password: String,
+    pub db_config: api::DbConfig,
+    pub logging: bool,
+}
+
+impl GRPCServerConfig {
+    pub fn load_config() -> Result<GRPCServerConfig, Box<dyn std::error::Error>> {
+        dotenv().ok();
+
+        // fetch server host; default to localhost if not defined
+        let host = env::var(constants::ENV_HOST)
+            .inspect_err(|_| {
+                event!(Level::WARN, "Host not defined, defaulting to 'localhost'");
+            })
+            .unwrap_or("127.0.0.1".to_string());
+
+        // fetch server port; default to 8080 if not defined
+        let port: u32 = env::var(ENV_PORT)
+            .inspect_err(|_| {
+                event!(
+                    Level::WARN,
+                    "Port not defined, defaulting to {}",
+                    DEFAULT_PORT
+                );
+            })
+            .unwrap_or(DEFAULT_PORT.to_string())
+            .parse()
+            .unwrap_or(DEFAULT_PORT.parse::<u32>().unwrap());
+
+        // fetch server root password; return err if not defined
+        let root_password = env::var(ENV_ROOT_PASSWORD).map_err(|_| {
+            errors::ConfigError::MissingRequiredEnvVar(ENV_ROOT_PASSWORD.to_string())
+        })?;
+
+        // fetch server storage type
+        let storage_type_str = env::var(ENV_STORAGE_TYPE)
+            .inspect_err(|_| {
+                event!(
+                    Level::WARN,
+                    "Storage Type not defined, defaulting to InMemory"
+                )
+            })
+            .unwrap_or_default();
+        let storage_type = match storage_type_str.as_str() {
+            "inmemory" => storage::StorageType::InMemory,
+            "rocksdb" => storage::StorageType::RocksDb,
+            _ => storage::StorageType::InMemory, // default to InMemory if not specified
+        };
+
+        // fetch server index type
+        let index_type_str = env::var(ENV_INDEX_TYPE)
+            .inspect_err(|_| event!(Level::WARN, "Index Type not defined, defaulting to flat"))
+            .unwrap_or("flat".to_string())
+            .to_lowercase();
+        let index_type = match index_type_str.as_str() {
+            "flat" => IndexType::Flat,
+            "kdtree" => IndexType::KDTree,
+            "hnsw" => IndexType::HNSW,
+            _ => IndexType::Flat, // default to Flat if not specified
+        };
+
+        // fetch dimension size
+        let dimension: usize = env::var(ENV_DIMENSION)
+            .map_err(|_| errors::ConfigError::MissingRequiredEnvVar(ENV_DIMENSION.to_string()))?
+            .parse()
+            .map_err(|_| errors::ConfigError::InvalidDimension)?;
+
+        // fetch data path; create tempdir if not specified
+        let data_path: PathBuf;
+        if let Ok(data_path_str) = env::var(ENV_DATA_PATH) {
+            data_path = PathBuf::from(data_path_str);
+            fs::create_dir_all(&data_path).map_err(|_| errors::ConfigError::InvalidDataPath)?;
+        } else {
+            let tempbuf = tempdir().unwrap().path().to_path_buf().join("vectordb");
+            fs::create_dir_all(&tempbuf)?;
+            event!(
+                Level::WARN,
+                "Data Path not specified, using temporary directory: {:?}",
+                tempbuf.clone()
+            );
+            data_path = tempbuf;
+        }
+
+        // create db config for api
+        let db_config = api::DbConfig {
+            storage_type,
+            index_type,
+            data_path,
+            dimension,
+        };
+
+        // create socket address for grpc server
+        let addr: SocketAddr = format!("{}:{}", host, port).parse()?;
+
+        // check if logging is enabled
+        let mut logging: bool = true; // default to logging enabled
+        if let Ok(logging_str) = env::var(ENV_LOGGING) {
+            logging = logging_str.parse().unwrap_or(true);
+        }
+
+        Ok(GRPCServerConfig {
+            addr,
+            root_password,
+            db_config,
+            logging,
+        })
+    }
+}