Add logging for RPCs; fixed warnings

Author: TanmayArya-1p

crates/grpc_server/build.rs

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

crates/grpc_server/src/config.rs

@@ -63,7 +63,7 @@ impl GRPCServerConfig {
             .map_err(|_| errors::ConfigError::InvalidDimension)?;
 
         // fetch data path; create tempdir if not specified
-        let mut data_path: PathBuf;
+        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)?;
@@ -85,9 +85,9 @@ impl GRPCServerConfig {
         let addr: SocketAddr = format!("{}:{}", host, port).parse().unwrap();
 
         // check if logging is enabled
-        let mut logging: bool = false; // default to no logging
+        let mut logging: bool = true; // default to no logging
         if let Ok(logging_str) = env::var(ENV_LOGGING) {
-            logging = logging_str.parse().unwrap_or(false);
+            logging = logging_str.parse().unwrap_or(true);
         }
 
         Ok(GRPCServerConfig {

crates/grpc_server/src/interceptors.rs

@@ -1,18 +1,6 @@
-use tonic::{GrpcMethod, Request, Status, service::Interceptor};
+use tonic::{Status, service::Interceptor};
 use tracing::{Level, event};
 
-pub fn logging_interceptor(req: Request<()>) -> Result<Request<()>, Status> {
-    if let Some(method) = req.extensions().get::<GrpcMethod>() {
-        println!(
-            "RPC to {}:{} : From {:?}",
-            method.service(),
-            method.method(),
-            req.remote_addr()
-        );
-    }
-    Ok(req)
-}
-
 #[derive(Clone)]
 pub struct AuthInterceptor {
     root_password: String,

crates/grpc_server/src/lib.rs

@@ -2,3 +2,4 @@ pub mod config;
 pub mod constants;
 pub mod errors;
 pub mod interceptors;
+pub mod utils;

crates/grpc_server/src/main.rs

@@ -1,15 +1,15 @@
 use api;
 use defs;
+use grpc_server::config::GRPCServerConfig;
 use grpc_server::constants::SIMILARITY_PROTOBUF_MAP;
 use grpc_server::interceptors;
-use grpc_server::{config::GRPCServerConfig, interceptors::logging_interceptor};
+use grpc_server::utils::log_rpc;
 use std::panic;
 use tonic::{Request, Response, Status, service::InterceptorLayer, transport::Server};
-use tracing_subscriber;
-
 use tracing::{Level, event};
+use tracing_subscriber;
 use vectordb::{
-    DenseVector, InsertVectorRequest, Point, PointId, SearchRequest, SearchResponse, Similarity,
+    DenseVector, InsertVectorRequest, Point, PointId, SearchRequest, SearchResponse,
     vector_db_server::{VectorDb, VectorDbServer},
 };
 pub mod vectordb {
@@ -18,6 +18,7 @@ pub mod vectordb {
 
 pub struct VectorDBService {
     pub vector_db: api::VectorDb,
+    pub logging: bool,
 }
 
 #[tonic::async_trait]
@@ -26,16 +27,19 @@ impl VectorDb for VectorDBService {
         &self,
         request: Request<InsertVectorRequest>,
     ) -> Result<Response<PointId>, Status> {
+        log_rpc("insert_vector", self.logging);
+
         let inner_request = request.into_inner();
 
         let dense_vector = inner_request.vector;
         if dense_vector.is_none() {
             return Err(Status::invalid_argument("dense_vector is empty"));
         }
-        // fetch payload and default to empty struct otherwise
-        let payload = inner_request.payload.unwrap_or_default();
 
         // TODO: Implement payload handling once its defined
+        // fetch payload and default to empty struct otherwise
+        // let payload = inner_request.payload.unwrap_or_default();
+
         let point_id = self
             .vector_db
             .insert(dense_vector.unwrap().values, defs::Payload {});
@@ -53,6 +57,8 @@ impl VectorDb for VectorDBService {
     }
 
     async fn get_point(&self, request: Request<PointId>) -> Result<Response<Point>, Status> {
+        log_rpc("get_point", self.logging);
+
         let inner_request = request.into_inner();
 
         let point_id = inner_request.id;
@@ -79,6 +85,8 @@ impl VectorDb for VectorDBService {
         &self,
         request: Request<SearchRequest>,
     ) -> Result<Response<SearchResponse>, Status> {
+        log_rpc("search_points", self.logging);
+
         let search_request = request.into_inner();
 
         //TODO: distance function panics if different dimensions; no dimension enforcement
@@ -105,12 +113,14 @@ impl VectorDb for VectorDBService {
     }
 
     async fn delete_point(&self, request: Request<PointId>) -> Result<Response<()>, Status> {
+        log_rpc("delete_point", self.logging);
+
         let point_id = request.into_inner().id;
 
         // TODO: delete call needs to return a boolean indicating if point is present or not
-        // if self.vector_db.delete(point_id)? {
-        //     return Err(Status::not_found("Point not found"));
-        // };
+        if self.vector_db.delete(point_id).is_err() {
+            return Err(Status::not_found("Error deleting point"));
+        };
 
         Ok(Response::new(()))
     }
@@ -129,6 +139,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let vector_db_service = VectorDBService {
         vector_db: vector_db_api,
+        logging: config.logging,
     };
 
     event!(
@@ -140,7 +151,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let auth_interceptor = interceptors::AuthInterceptor::new(config.root_password);
 
     let _ = Server::builder()
-        .layer(InterceptorLayer::new(logging_interceptor))
         .layer(InterceptorLayer::new(auth_interceptor))
         .add_service(VectorDbServer::new(vector_db_service))
         .serve(config.addr)

crates/grpc_server/src/utils.rs

@@ -0,0 +1,8 @@
+use tracing;
+
+#[inline(always)] // doing function inlining for optimization
+pub fn log_rpc(rpc: &str, logging: bool) {
+    if logging {
+        tracing::event!(tracing::Level::INFO, "{} RPC called", rpc);
+    }
+}