Callback-based gRPC client with C interface #963
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,126 @@ | ||
| //! This module implements support for callback-based gRPC service that has a callback invoked for | ||
| //! every gRPC call instead of using network. | ||
|
|
||
| use anyhow::anyhow; | ||
| use bytes::{BufMut, BytesMut}; | ||
| use futures_util::future::BoxFuture; | ||
| use futures_util::stream; | ||
| use http::{HeaderMap, Request, Response}; | ||
| use http_body_util::{BodyExt, StreamBody, combinators::BoxBody}; | ||
| use hyper::body::{Bytes, Frame}; | ||
| use std::{ | ||
| sync::Arc, | ||
| task::{Context, Poll}, | ||
| }; | ||
| use tonic::{Status, metadata::GRPC_CONTENT_TYPE}; | ||
| use tower::Service; | ||
|
|
||
| /// gRPC request for use by a callback. | ||
| pub struct GrpcRequest<'a> { | ||
| /// Reference to the gRPC service name. | ||
| pub service: &'a str, | ||
|
|
||
| /// Reference to the gRPC RPC name. | ||
| pub rpc: &'a str, | ||
|
|
||
| /// Reference to the gRPC request headers. | ||
| pub headers: &'a HeaderMap, | ||
|
|
||
| /// Reference to the gRPC protobuf bytes of the request. | ||
cretz marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| pub proto: Bytes, | ||
| } | ||
|
|
||
| /// Successful gRPC response returned by a callback. | ||
| pub struct GrpcSuccessResponse { | ||
| /// Response headers. | ||
| pub headers: HeaderMap, | ||
|
|
||
| /// Response proto bytes. | ||
| pub proto: Vec<u8>, | ||
| } | ||
|
|
||
| /// gRPC service that invokes the given callback on each call. | ||
| #[derive(Clone)] | ||
| pub struct CallbackBasedGrpcService { | ||
| /// Callback to invoke on each RPC call. | ||
| #[allow(clippy::type_complexity)] // Signature is not that complex | ||
| pub callback: Arc< | ||
| dyn for<'a> Fn(GrpcRequest<'a>) -> BoxFuture<'a, Result<GrpcSuccessResponse, Status>> | ||
| + Send | ||
| + Sync, | ||
| >, | ||
| } | ||
|
|
||
| impl Service<Request<tonic::body::Body>> for CallbackBasedGrpcService { | ||
| type Response = http::Response<tonic::body::Body>; | ||
| type Error = anyhow::Error; | ||
| type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>; | ||
|
|
||
| fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> { | ||
| Poll::Ready(Ok(())) | ||
| } | ||
|
|
||
| fn call(&mut self, req: Request<tonic::body::Body>) -> Self::Future { | ||
| let callback = self.callback.clone(); | ||
|
|
||
| Box::pin(async move { | ||
| // Build req | ||
| let (parts, body) = req.into_parts(); | ||
| let mut path_parts = parts.uri.path().trim_start_matches('/').split('/'); | ||
| let req_body = body.collect().await.map_err(|e| anyhow!(e))?.to_bytes(); | ||
| // Body is flag saying whether compressed (we do not support that), then 32-bit length, | ||
| // then the actual proto. | ||
| if req_body.len() < 5 { | ||
| return Err(anyhow!("Too few request bytes: {}", req_body.len())); | ||
| } else if req_body[0] != 0 { | ||
| return Err(anyhow!("Compression not supported")); | ||
| } | ||
| let req_proto_len = | ||
| u32::from_be_bytes([req_body[1], req_body[2], req_body[3], req_body[4]]) as usize; | ||
| if req_body.len() < 5 + req_proto_len { | ||
| return Err(anyhow!( | ||
| "Expected request body length at least {}, got {}", | ||
| 5 + req_proto_len, | ||
| req_body.len() | ||
| )); | ||
| } | ||
| let req = GrpcRequest { | ||
| service: path_parts.next().unwrap_or_default(), | ||
| rpc: path_parts.next().unwrap_or_default(), | ||
| headers: &parts.headers, | ||
|
There was a problem hiding this comment. Why manipulate the body at all? If we pass through the compression flag and length, we can document that and allow the callback implementer to handle compression if they want There was a problem hiding this comment. Because most gRPC clients deal in protos, not full bodies with the extra 5 bytes in front. For callback implementers that are, say, delegating to their own in-language clients, those clients operate on protobuf bodies, not raw HTTP ones. I don't think we should ask them to carve up the body to get the proto out (or put it back). If there is a use case for needing to know the compression byte, we can add it, but we are in control of the client call so it will always be what we want anyways. There was a problem hiding this comment. Right, what I'm saying is the callback client isn't really a pure gRPC client, because it may want to delegate to something that implements compression too. So seems to me we should just document that and leave it to the caller. There was a problem hiding this comment. They can't enable compression, we'd have to with our Tonic client which is what builds this body. But there's no value for us to do so in this case since it's all in memory. If they want to delegate to something that implements compression, they can/should. They can wrap the whole proto bytes in pre-negotiated compression if they'd like. But from our in-memory perspective, we give them proto bytes, how that's represented on the in-memory wire via the few bits of Rust code between Tonic and this callback should have no user effect. It's up to us and it should always be 0 (no compression) IMO. The compression byte is about pre-negotiated compression algorithm with the server, which doesn't apply to in-memory nor should it. Overall, this compression is just a boolean used by gRPC HTTP, but has no value for in-memory representation and will always be 0. It doesn't compress the bytes or even tell you the algorithm, it's just a note saying the negotiated compression with the server is in effect (there is no server here). This is unrelated to whether a user wants to use compression to their upstream implementation. |
||
| proto: req_body.slice(5..5 + req_proto_len), | ||
| }; | ||
|
|
||
| // Invoke and handle response | ||
| match (callback)(req).await { | ||
| Ok(success) => { | ||
| // Create body bytes which requires a flag saying whether compressed, then | ||
| // message len, then actual message. So we create a Bytes for those 5 prepend | ||
| // parts, then stream it alongside the user-provided Vec. This allows us to | ||
| // avoid copying the vec | ||
| let mut body_prepend = BytesMut::with_capacity(5); | ||
| body_prepend.put_u8(0); // 0 means no compression | ||
| body_prepend.put_u32(success.proto.len() as u32); | ||
| let stream = stream::iter(vec![ | ||
| Ok::<_, Status>(Frame::data(Bytes::from(body_prepend))), | ||
| Ok::<_, Status>(Frame::data(Bytes::from(success.proto))), | ||
| ]); | ||
| let stream_body = StreamBody::new(stream); | ||
| let full_body = BoxBody::new(stream_body).boxed(); | ||
|
|
||
| // Build response appending headers | ||
| let mut resp_builder = Response::builder() | ||
| .status(200) | ||
| .header(http::header::CONTENT_TYPE, GRPC_CONTENT_TYPE); | ||
| for (key, value) in success.headers.iter() { | ||
| resp_builder = resp_builder.header(key, value); | ||
| } | ||
| Ok(resp_builder | ||
| .body(tonic::body::Body::new(full_body)) | ||
| .map_err(|e| anyhow!(e))?) | ||
| } | ||
| Err(status) => Ok(status.into_http()), | ||
| } | ||
| }) | ||
| } | ||
| } | ||
Uh oh!
There was an error while loading. Please reload this page.