feat: add grpc RPC context extractor

Author: rsfreitas

Cargo.toml

@@ -17,4 +17,6 @@ tokio = { version = "1.41.1", features = ["full"] }
 toml = "0.8.19"
 tower = "0.5.1"
 tower-service = "0.3.3"
-validator = { version = "0.19.0", features = ["derive"] }
+validator = { version = "0.19.0", features = ["derive"] }
+futures = "0.3.31"
+prost = "0.13.3"

examples/services/grpc/src/main.rs

@@ -2,7 +2,8 @@ pub mod helloworld {
     include!("generated/helloworld.rs"); // Include generated code
 }
 
-use mikros::service::builder::ServiceBuilder;
+use mikros::features;
+use mikros::service::{builder::ServiceBuilder, context};
 use tonic::{Request, Response, Status};
 
 use helloworld::greeter_server::{Greeter, GreeterServer};
@@ -14,6 +15,13 @@ pub struct MyGreeter {}
 #[tonic::async_trait]
 impl Greeter for MyGreeter {
     async fn say_hello(&self, request: Request<HelloRequest>) -> Result<Response<HelloReply>, Status> {
+        let ctx = context::from_request(&request);
+        ctx.logger().info("say hello RPC called");
+
+        let _ = features::example::retrieve(&ctx, |api| {
+            api.do_something();
+        });
+
         let reply = HelloReply {
             message: format!("Hello, {}!", request.into_inner().name),
         };
@@ -26,7 +34,7 @@ impl Greeter for MyGreeter {
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let greeter = MyGreeter::default();
     let greeter_service = GreeterServer::new(greeter);
-    
+
     let svc = ServiceBuilder::default()
         .grpc(greeter_service)
         .build();
@@ -39,6 +47,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         },
         Err(e) => panic!("{}", e.to_string())
     }
-    
+
     Ok(())
 }

src/grpc/mod.rs

@@ -0,0 +1,65 @@
+use std::sync::Arc;
+use std::task::{Context, Poll};
+
+use tower::{Layer, Service};
+
+use crate::service::context;
+
+#[derive(Clone)]
+pub(crate) struct ContextExtractor {
+    ctx: Arc<context::Context>,
+}
+
+impl ContextExtractor {
+    pub(crate) fn new(ctx: &context::Context) -> Self {
+        ContextExtractor {
+            ctx: Arc::new(ctx.clone()),
+        }
+    }
+}
+
+impl<S> Layer<S> for ContextExtractor {
+    type Service = ContextExtractorMiddleware<S>;
+
+    fn layer(&self, service: S) -> Self::Service {
+        ContextExtractorMiddleware {
+            inner: service,
+            ctx: self.ctx.clone(),
+        }
+    }
+}
+
+#[derive(Clone)]
+pub(crate) struct ContextExtractorMiddleware<S> {
+    inner: S,
+    ctx: Arc<context::Context>,
+}
+
+impl<S, B> Service<http::Request<B>> for ContextExtractorMiddleware<S>
+where
+    S: Service<http::Request<B>, Response = http::Response<B>>
+        + Clone
+        + Send
+        + 'static,
+    S::Future: Send,
+    B: Send + 'static,
+{
+    type Response = S::Response;
+    type Error = S::Error;
+    type Future = futures::future::BoxFuture<'static, Result<Self::Response, Self::Error>>;
+
+    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.inner.poll_ready(cx)
+    }
+
+    fn call(&mut self, mut req: http::Request<B>) -> Self::Future {
+        let clone = self.inner.clone();
+        let mut inner = std::mem::replace(&mut self.inner, clone);
+
+        req.extensions_mut().insert(self.ctx.clone());
+        Box::pin(async move {
+            let response = inner.call(req).await?;
+            Ok(response)
+        })
+    }
+}

src/lib.rs

@@ -6,3 +6,4 @@ pub mod plugin;
 pub mod service;
 
 mod args;
+mod grpc;

src/service/context.rs

@@ -46,4 +46,9 @@ impl Context {
     pub fn service_name(&self) -> String {
         self.definitions.name.clone()
     }
-}
+}
+
+/// Retrieves the mikros Context from an RPC request argument.
+pub fn from_request<B: prost::Message>(request: &tonic::Request<B>) -> Arc<Context> {
+    request.extensions().get::<Arc<Context>>().unwrap().clone()
+}

src/service/grpc/mod.rs

@@ -1,6 +1,7 @@
 use std::collections::HashMap;
 use std::convert::Infallible;
 use std::sync::Arc;
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 
 use http::{request::Request, response::Response};
 use logger::fields::FieldValue;
@@ -12,6 +13,7 @@ use tonic::transport::Server;
 use crate::{definition, env, plugin};
 use crate::service::context::Context;
 use crate::errors as merrors;
+use crate::grpc;
 
 #[derive(Clone)]
 pub(crate) struct Grpc<S> {
@@ -69,16 +71,25 @@ where
         ]
     }
 
-    async fn run(&self, _ctx: &Context, shutdown_rx: watch::Receiver<()>) -> merrors::Result<()> {
-        let addr = format!("127.0.0.1:{}", self.port).parse().unwrap();
+    async fn run(&self, ctx: &Context, shutdown_rx: watch::Receiver<()>) -> merrors::Result<()> {
+        let addr = SocketAddr::new(
+            IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
+            self.port as u16,
+        );
+
         let shutdown_signal = async {
             let mut shutdown_rx = shutdown_rx.clone();
 
             // Wait until the receiver sees the shutdown signal
             shutdown_rx.changed().await.ok();
         };
 
+        let layer = tower::ServiceBuilder::new()
+            .layer(grpc::ContextExtractor::new(&ctx))
+            .into_inner();
+
         if let Err(e) = Server::builder()
+            .layer(layer)
             .add_service(self.server.clone())
             .serve_with_shutdown(addr, shutdown_signal)
             .await