Add a serving multiplexer to server both web and grpc requests on the…

… same port.
dinowernli · May 12, 2024 · 0759875 · 0759875
1 parent 32b8b4a
commit 0759875
Show file tree

Hide file tree

Showing 7 changed files with 246 additions and 27 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -7,9 +7,12 @@ edition = "2021"
 [dependencies]
 async-std = {version = "1.12", features = ["attributes"]}
 async-trait = "0.1.80"
+axum = "0.6.12"
 bytes = "1.6"
 chrono = "0.4"
 futures = "0.3"
+hyper="0.14.25"
+multiplex-tonic-hyper = "0.1"
 im = "15"
 prost = "0.12"
 rand = "0.8"
@@ -22,6 +25,7 @@ tracing = "0.1.40"
 tracing-subscriber = {version = "0.3.1", features = ["env-filter"]}
 tower = "0.4.13"
 pin-project = "1.1"
+querystring = "1.1"
 
 [build-dependencies]
 tonic-build = "0.11"
diff --git a/src/keyvalue/http.rs b/src/keyvalue/http.rs
@@ -0,0 +1,117 @@
+use axum::routing::get;
+use axum::Router;
+use hyper::{Body, StatusCode};
+use std::sync::Arc;
+use tonic::Request;
+
+use crate::keyvalue::keyvalue_proto::key_value_server::KeyValue;
+use crate::keyvalue::keyvalue_proto::GetRequest;
+
+#[derive(Clone)]
+pub struct HttpHandler {
+    service: Arc<dyn KeyValue>,
+}
+
+impl HttpHandler {
+    pub fn new(service: Arc<dyn KeyValue>) -> Self {
+        Self { service }
+    }
+
+    /**
+     * Installs routes that allow interacting with the keyvalue store. Example queries
+     * include /get?key=foo.
+     */
+    pub fn routes(self: Arc<Self>) -> Router {
+        let hello = move |req: hyper::Request<Body>| async move { self.handle_get(req).await };
+        Router::new().route("/get", get(hello))
+    }
+
+    fn invalid(message: String) -> hyper::Response<Body> {
+        hyper::Response::builder()
+            .status(StatusCode::BAD_REQUEST)
+            .body(Body::from(format!("{}\n", message)))
+            .unwrap()
+    }
+
+    async fn handle_get(&self, request: hyper::Request<Body>) -> hyper::Response<Body> {
+        let query = match request.uri().query() {
+            Some(q) => q,
+            None => return HttpHandler::invalid("must pass query".to_string()),
+        };
+
+        let parsed = querystring::querify(query);
+        let key = match parsed
+            .iter()
+            .find(|(a, _)| a.eq(&"key"))
+            .map(|(_, b)| b.to_string())
+        {
+            Some(value) => value,
+            _ => return HttpHandler::invalid("must pass key parameter".to_string()),
+        };
+
+        let request = Request::new(GetRequest {
+            key: key.clone().as_bytes().to_vec(),
+            version: -1,
+        });
+
+        let output = match self.service.get(request).await {
+            Ok(p) => {
+                let proto = p.into_inner();
+                match proto.entry {
+                    Some(e) => match String::from_utf8(e.value) {
+                        Ok(value) => format!("{}={}, version={}", key, value, proto.version),
+                        _ => format!("failed to parse value as utf8 for key {}", key),
+                    },
+                    None => format!("no value for key {}", key),
+                }
+            }
+            Err(status) => format!("Failed to query keyvalue store {}", status.to_string()),
+        };
+
+        hyper::Response::builder()
+            .status(StatusCode::OK)
+            .body(Body::from(output))
+            .unwrap()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::keyvalue::grpc::PutRequest;
+    use crate::keyvalue::keyvalue_proto::{Entry, GetResponse, PutResponse};
+    use crate::root;
+    use tonic::{Response, Status};
+
+    struct FakeKeyValue {}
+
+    #[tonic::async_trait]
+    impl KeyValue for FakeKeyValue {
+        async fn get(&self, request: Request<GetRequest>) -> Result<Response<GetResponse>, Status> {
+            todo!()
+            // Ok(Response::new(GetResponse { entry: Some(Entry {
+            //     key: request.into_inner().key.clone(),
+            //     value: vec![],
+            // }), version: 0 }))
+        }
+
+        async fn put(&self, request: Request<PutRequest>) -> Result<Response<PutResponse>, Status> {
+            todo!()
+        }
+    }
+
+    #[tokio::test]
+    async fn test_http() {
+        assert!(false);
+
+        let kv = Arc::new(FakeKeyValue {});
+        let http = Arc::new(HttpHandler {
+            service: kv.clone(),
+        });
+
+        let web_service = Router::new()
+            .route("/", get(root))
+            .nest("/keyvalue", http.routes())
+            .into_make_service();
+    }
+}
diff --git a/src/keyvalue/mod.rs b/src/keyvalue/mod.rs
@@ -2,15 +2,17 @@
 // cluster. Users of this module are expected to run the service in their grpc
 // server and/or use the generated grpc client code to make requests.
 
-pub use crate::keyvalue::store::{MapStore};
+pub use crate::keyvalue::store::{MapStore, Store};
+pub use http::HttpHandler;
 pub use service::KeyValueService;
 
 pub mod grpc {
     pub use crate::keyvalue::keyvalue_proto::key_value_client::KeyValueClient;
     pub use crate::keyvalue::keyvalue_proto::key_value_server::KeyValueServer;
-    pub use crate::keyvalue::keyvalue_proto::{PutRequest};
+    pub use crate::keyvalue::keyvalue_proto::PutRequest;
 }
 
+pub(in crate::keyvalue) mod http;
 #[path = "generated/keyvalue_proto.rs"]
 pub(in crate::keyvalue) mod keyvalue_proto;
 pub(in crate::keyvalue) mod service;

diff --git a/src/keyvalue/store.rs b/src/keyvalue/store.rs
@@ -1,11 +1,11 @@
 extern crate bytes;
 extern crate im;
 
-use std::collections::VecDeque;
 use bytes::Bytes;
 use im::HashMap;
-use prost::Message;
 use keyvalue_proto::{Entry, Operation, Snapshot};
+use prost::Message;
+use std::collections::VecDeque;
 
 use crate::keyvalue::keyvalue_proto;
 use crate::keyvalue::keyvalue_proto::operation::Op::Set;