first pass at a generic grpc impl

This commit is contained in:
Lucio Franco
2019-08-13 17:33:44 -04:00
parent c11de766ca
commit 9fbfec4c62
7 changed files with 474 additions and 243 deletions
+25 -50
View File
@@ -68,9 +68,9 @@ pub fn server(attr: TokenStream, item: TokenStream) -> TokenStream {
} }
impl _codegen::Service<_codegen::http::Request<()>> for GrpcServer { impl _codegen::Service<_codegen::http::Request<()>> for GrpcServer {
type Response = tonic::Response<()>; type Response = tonic::Response<tonic::body::BoxAsyncBody>;
type Error = tonic::error::Never; type Error = tonic::error::Never;
type Future = greeter::ResponseFuture; type Future = _codegen::ResponseFuture2<Self::Response, Self::Error>;
fn poll_ready(&mut self, _cx: &mut _codegen::Context<'_>) -> _codegen::Poll<Result<(), Self::Error>> { fn poll_ready(&mut self, _cx: &mut _codegen::Context<'_>) -> _codegen::Poll<Result<(), Self::Error>> {
Ok(()).into() Ok(()).into()
@@ -81,60 +81,35 @@ pub fn server(attr: TokenStream, item: TokenStream) -> TokenStream {
match request.uri().path() { match request.uri().path() {
"/helloworld.Greeter/SayHello" => { "/helloworld.Greeter/SayHello" => {
// let kind = greeter::methods::SayHello(self.inner.clone()); let inner = self.inner.clone();
// greeter::ResponseFuture { kind: greeter::Kind::SayHello(kind) } let fut = async move {
self.inner.stream(request).await?; let codec = tonic::codec::UnitCodec::default();
let mut grpc = tonic::server::Grpc::new(codec);
let response = match grpc.unary_request(request).await {
Ok(request) => inner.#m_ident(request).await,
Err(status) => Err(status),
};
let response = grpc.unary_response(response.unwrap())
.await
.map(|b| tonic::body::BoxAsyncBody::new(b));
// .map(|b| tonic::body::BoxBody::map(b))
Ok(response)
};
Box::pin(fut)
},
"helloworld.Greeter/SayHelloStream" => {
unimplemented!() unimplemented!()
}, },
_ => unimplemented!("use grpc unimplemented") _ => unimplemented!("use grpc unimplemented")
} }
} }
} }
// TODO: get actual service name
pub mod greeter {
use tonic::_codegen::*;
pub struct ResponseFuture {
pub kind: Kind,
}
pub enum Kind {
SayHello(methods::SayHello),
}
impl Future for ResponseFuture {
type Output = Result<tonic::Response<()>, tonic::error::Never>;
fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
unimplemented!()
}
}
pub mod methods {
use tonic::_codegen::*;
pub struct SayHello(pub std::sync::Arc<super::super::#s>);
impl Service<tonic::Request<()>> for SayHello {
type Response = tonic::Response<()>;
type Error = tonic::Status;
type Future = ResponseFuture<Self::Response>;
fn poll_ready(&mut self, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Ok(()).into()
}
fn call(&mut self, request: tonic::Request<()>) -> Self::Future {
let inner = self.0.clone();
Box::pin(async move {
inner.#m_ident(request).await
})
}
}
}
}
}; };
original.extend(TokenStream::from(ts)); original.extend(TokenStream::from(ts));
+10 -6
View File
@@ -32,13 +32,17 @@ impl MyGreeter {
Ok(Response::new(())) Ok(Response::new(()))
} }
pub async fn server_stream(&self, request: Request<()>) -> Result<impl Stream, Status> { // pub async fn streaming(&self, request: Request<impl Stream>) -> Result<Response<impl Stream>, Status> {
unimplemented!() // unimplemented!()
} // }
pub async fn client_stream(&self, request: Request<impl Stream>) -> Result<(), Status> { // pub async fn server_stream(&self, request: Request<()>) -> Result<impl Stream, Status> {
unimplemented!() // unimplemented!()
} // }
// pub async fn client_stream(&self, request: Request<impl Stream>) -> Result<(), Status> {
// unimplemented!()
// }
} }
#[tokio::test] #[tokio::test]
+81
View File
@@ -4,10 +4,91 @@ use futures_core::Stream;
use futures_util::{ready, TryStreamExt}; use futures_util::{ready, TryStreamExt};
use http::HeaderMap; use http::HeaderMap;
use http_body::Body; use http_body::Body;
use std::pin::Pin;
use std::task::{Context, Poll}; use std::task::{Context, Poll};
pub type BytesBuf = <Bytes as IntoBuf>::Buf; pub type BytesBuf = <Bytes as IntoBuf>::Buf;
pub struct BoxBody {
inner: Box<dyn Body<Data = BytesBuf, Error = Status> + Send>,
}
impl BoxBody {
/// Create a new `BoxBody` mapping item and error to the default types.
pub fn map_from<B>(inner: B) -> Self
where
B: Body<Data = BytesBuf, Error = Status> + Send + 'static,
{
BoxBody {
inner: Box::new(inner),
}
}
}
impl Body for BoxBody {
type Data = BytesBuf;
type Error = Status;
fn is_end_stream(&self) -> bool {
self.inner.is_end_stream()
}
fn poll_data(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<Self::Data, Self::Error>>> {
self.inner.poll_data(cx)
}
fn poll_trailers(
&mut self,
cx: &mut Context<'_>,
) -> Poll<Result<Option<http::HeaderMap>, Self::Error>> {
self.inner.poll_trailers(cx)
}
}
pub struct BoxAsyncBody {
inner: Pin<Box<dyn Stream<Item = Result<BytesBuf, Status>> + Send>>,
error: Option<Status>,
}
impl BoxAsyncBody {
pub fn new<S>(inner: S) -> Self
where
S: Stream<Item = Result<crate::body::BytesBuf, Status>> + Send + 'static,
{
Self {
inner: Box::pin(inner),
error: None,
}
}
}
impl Body for BoxAsyncBody {
type Data = BytesBuf;
type Error = Status;
fn poll_data(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<Self::Data, Self::Error>>> {
match ready!(self.inner.try_poll_next_unpin(cx)) {
Some(Ok(d)) => Some(Ok(d)).into(),
Some(Err(status)) => {
self.error = Some(status);
None.into()
}
None => None.into(),
}
}
fn poll_trailers(&mut self, _cx: &mut Context<'_>) -> Poll<Result<Option<HeaderMap>, Status>> {
let status = if let Some(status) = self.error.take() {
status
} else {
Status::new(Code::Ok, "")
};
Poll::Ready(Ok(Some(status.to_header_map()?)))
}
}
#[derive(Debug)]
pub struct AsyncBody<S> { pub struct AsyncBody<S> {
inner: S, inner: S,
error: Option<Status>, error: Option<Status>,
+197 -1
View File
@@ -1,6 +1,202 @@
#![allow(dead_code)]
use crate::{body::BytesBuf, Status};
use async_stream::stream;
use bytes::{Buf, BufMut, BytesMut, IntoBuf};
use futures_core::TryStream;
use futures_util::{future, TryStreamExt};
use http_body::Body;
use tokio_codec::{Decoder, Encoder};
use tracing::{debug, trace};
pub trait Codec { pub trait Codec {
type Encode; type Encode;
type Decode; type Decode;
type Encoder; type Encoder: Encoder<Item = Self::Encode, Error = Status>;
type Decoder: Decoder<Item = Self::Decode, Error = Status>;
fn encoder(&mut self) -> Self::Encoder;
fn decoder(&mut self) -> Self::Decoder;
}
pub async fn encode<T, U>(
mut encoder: T,
mut source: U,
) -> impl TryStream<Ok = BytesBuf, Error = Status>
where
T: Encoder<Error = Status>,
U: TryStream<Ok = T::Item, Error = Status> + Unpin,
{
stream! {
let mut buf = BytesMut::with_capacity(1024);
loop {
match source.try_next().await {
Ok(Some(item)) => {
encoder.encode(item, &mut buf).map_err(drop).unwrap();
let len = buf.len();
yield Ok(buf.split_to(len).freeze().into_buf());
},
Ok(None) => break,
Err(status) => yield Err(status),
}
}
}
}
pub fn decode<T, B>(mut decoder: T, mut source: B) -> impl TryStream<Ok = T::Item, Error = Status>
where
T: Decoder<Error = Status>,
T::Item: Unpin + 'static,
B: Body,
B::Error: Into<crate::Error>,
{
stream! {
let mut buf = BytesMut::with_capacity(1024);
let mut state = State::ReadHeader;
loop {
// TODO: use try_stream! and ?
if let Some(item) = decode_chunk(&mut decoder, &mut buf, &mut state).unwrap() {
yield Ok(item);
}
let chunk = match future::poll_fn(|cx| source.poll_data(cx)).await {
Some(Ok(d)) => Some(d),
Some(Err(e)) => {
let err = e.into();
debug!("decoder inner stream error: {:?}", err);
let status = Status::from_error(&*err);
yield Err(status);
break;
},
None => None,
};
if let Some(data)= chunk {
buf.put(data);
} else {
if buf.has_remaining_mut() {
trace!("unexpected EOF decoding stream");
// yield Err(Status::new(
// Code::Internal,
// "Unexpected EOF decoding stream.".to_string(),
// ));
} else {
break;
}
}
// TODO: poll_trailers for Response status code
}
}
}
fn decode_chunk<T>(
decoder: &mut T,
buf1: &mut BytesMut,
state: &mut State,
) -> Result<Option<T::Item>, Status>
where
T: Decoder<Error = Status>,
{
let mut buf = (&buf1[..]).into_buf();
if let State::ReadHeader = state {
if buf.remaining() < 5 {
return Ok(None);
}
let is_compressed = match buf.get_u8() {
0 => false,
1 => {
trace!("message compressed, compression not supported yet");
return Err(crate::Status::new(
crate::Code::Unimplemented,
"Message compressed, compression not supported yet.".to_string(),
));
}
f => {
trace!("unexpected compression flag");
return Err(crate::Status::new(
crate::Code::Internal,
format!("Unexpected compression flag: {}", f),
));
}
};
let len = buf.get_u32_be() as usize;
*state = State::ReadBody {
compression: is_compressed,
len,
}
}
if let State::ReadBody { len, .. } = state {
if buf.remaining() < *len {
return Ok(None);
}
match decoder.decode(buf1) {
Ok(Some(msg)) => {
*state = State::ReadHeader;
return Ok(Some(msg));
}
Ok(None) => return Ok(None),
Err(e) => {
return Err(e);
}
}
}
Ok(None)
}
#[derive(Default)]
pub struct UnitCodec;
impl Codec for UnitCodec {
type Encode = ();
type Decode = ();
type Encoder = UnitEncoder;
type Decoder = UnitDecoder;
fn encoder(&mut self) -> Self::Encoder {
UnitEncoder
}
fn decoder(&mut self) -> Self::Decoder {
UnitDecoder
}
}
pub struct UnitEncoder;
impl Encoder for UnitEncoder {
type Item = ();
type Error = crate::Status;
fn encode(&mut self, _item: Self::Item, _buf: &mut BytesMut) -> Result<(), Self::Error> {
unimplemented!()
}
}
pub struct UnitDecoder;
impl Decoder for UnitDecoder {
type Item = ();
type Error = Status;
fn decode(&mut self, _buf: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
Ok(Some(()))
}
}
#[derive(Debug)]
enum State {
ReadHeader,
ReadBody { compression: bool, len: usize },
Done,
} }
+7 -4
View File
@@ -1,16 +1,17 @@
#![feature(async_await, type_alias_impl_trait)] #![feature(async_await)]
#![recursion_limit = "256"]
//! gRPC implementation //! gRPC implementation
pub mod body;
pub mod codec;
#[doc(hidden)] #[doc(hidden)]
pub mod error; pub mod error;
pub mod metadata; pub mod metadata;
pub mod server;
mod body;
mod codec;
mod request; mod request;
mod response; mod response;
mod server;
mod status; mod status;
pub use request::Request; pub use request::Request;
@@ -40,6 +41,8 @@ pub mod _codegen {
pub use tower_service::Service; pub use tower_service::Service;
pub type ResponseFuture<T> = pub type ResponseFuture<T> =
self::Pin<Box<dyn self::Future<Output = Result<T, crate::Status>> + Send + 'static>>; self::Pin<Box<dyn self::Future<Output = Result<T, crate::Status>> + Send + 'static>>;
pub type ResponseFuture2<T, E> =
self::Pin<Box<dyn self::Future<Output = Result<T, E>> + Send + 'static>>;
pub mod http { pub mod http {
pub use http::*; pub use http::*;
+148 -182
View File
@@ -1,206 +1,172 @@
#![allow(dead_code)] use crate::{
codec::{self, Codec},
use crate::{Code, Request, Response, Status}; Request, Status,
use async_stream::stream; };
use bytes::{Buf, BufMut, Bytes, BytesMut, IntoBuf}; use futures_core::{Future, TryStream};
use futures_core::{Stream, TryStream}; use futures_util::{future, stream, TryStreamExt};
use futures_util::{future, stream, StreamExt, TryStreamExt};
use http_body::Body; use http_body::Body;
use std::future::Future; use std::pin::Pin;
use tokio_codec::{Decoder, Encoder};
use tower_service::Service;
use tracing::{debug, trace};
pub trait Codec { pub struct Grpc<T> {
type Encode; codec: T,
type Decode;
} }
pub struct Encode<T, U> { // type UnaryFuture<B> = Once<Ready<Result<B, Status>>>;
encoder: T, // type ResponseBody = impl Stream<Item = Result<crate::body::BytesBuf, Status>>;
source: U,
pub trait UnaryService<R> {
/// Protobuf response message type
type Response;
/// Response future
type Future: Future<Output = Result<crate::Response<Self::Response>, Status>>;
/// Call the service
fn call(&mut self, request: Request<R>) -> Self::Future;
} }
impl<T, U> Encode<T, U> pub trait ServerStreamingService<R> {
/// Protobuf response message type
type Response;
/// Stream of outbound response messages
type ResponseStream: TryStream<Ok = Self::Response, Error = crate::Status> + Unpin;
/// Response future
type Future: Future<Output = Result<crate::Response<Self::ResponseStream>, crate::Status>>;
/// Call the service
fn call(&mut self, request: Request<R>) -> Self::Future;
}
pub trait ClientStreamingService<RequestStream> {
/// Protobuf response message type
type Response;
/// Response future
type Future: Future<Output = Result<crate::Response<Self::Response>, Status>>;
/// Call the service
fn call(&mut self, request: Request<RequestStream>) -> Self::Future;
}
pub trait StreamingService<RequestStream> {
/// Protobuf response message type
type Response;
/// Stream of outbound response messages
type ResponseStream: TryStream<Ok = Self::Response, Error = crate::Status> + Unpin;
/// Response future
type Future: Future<Output = Result<crate::Response<Self::ResponseStream>, crate::Status>>;
/// Call the service
fn call(&mut self, request: Request<RequestStream>) -> Self::Future;
}
type BoxStream<T> = Pin<Box<dyn TryStream<Ok = T, Error = Status> + Send + 'static>>;
impl<T> Grpc<T>
where where
T: Encoder, T: Codec,
U: TryStream<Ok = T::Item, Error = Status> + Unpin, T::Decode: Unpin + 'static,
T::Encode: Unpin + 'static,
{ {
pub fn new(encoder: T, source: U) -> Self { pub fn new(codec: T) -> Self {
Encode { encoder, source } Self { codec }
} }
pub fn encode<'a>( pub async fn unary<S, B>(
&'a mut self, &mut self,
buf: &'a mut BytesMut, mut service: S,
) -> impl Stream<Item = Result<crate::body::BytesBuf, Status>> + 'a { req: http::Request<B>,
stream! { ) -> http::Response<impl TryStream<Ok = crate::body::BytesBuf, Error = Status>>
loop {
match self.source.try_next().await {
Ok(Some(item)) => {
self.encoder.encode(item, buf).map_err(drop).unwrap();
let len = buf.len();
yield Ok(buf.split_to(len).freeze().into_buf());
},
Ok(None) => break,
Err(status) => yield Err(status),
}
}
}
}
}
pub struct Streaming<T> {
decoder: T,
buf: BytesMut,
state: State,
}
#[derive(Debug)]
enum State {
ReadHeader,
ReadBody { compression: bool, len: usize },
Done,
}
impl<T> Streaming<T>
where
T: Decoder,
T::Item: Unpin + 'static,
{
pub fn decode<'a, B>(
&'a mut self,
source: &'a mut B,
) -> impl Stream<Item = Result<T::Item, Status>> + 'a
where where
S: UnaryService<T::Decode, Response = T::Encode>,
B: Body, B: Body,
B::Error: Into<crate::Error>, B::Error: Into<crate::Error>,
{ {
stream! { let (_parts, body) = req.into_parts();
loop { let stream = codec::decode(self.codec.decoder(), body).into_stream();
// TODO: use try_stream! and ? futures_util::pin_mut!(stream);
if let Some(item) = self.decode_chunk().unwrap() { let message = stream.try_next().await.unwrap().unwrap();
yield Ok(item); let request = Request::new(message);
} let response = service.call(request).await.unwrap();
let message = response.into_inner();
let source = stream::once(future::ok(message));
let body = codec::encode(self.codec.encoder(), source).await;
let chunk = match future::poll_fn(|cx| source.poll_data(cx)).await { http::Response::new(body)
Some(Ok(d)) => Some(d),
Some(Err(e)) => {
let err = e.into();
debug!("decoder inner stream error: {:?}", err);
let status = Status::from_error(&*err);
yield Err(status);
break;
},
None => None,
};
if let Some(data)= chunk {
self.buf.put(data);
} else {
if self.buf.has_remaining_mut() {
trace!("unexpected EOF decoding stream");
yield Err(Status::new(
Code::Internal,
"Unexpected EOF decoding stream.".to_string(),
));
} else {
break;
}
}
}
}
} }
fn decode_chunk(&mut self) -> Result<Option<T::Item>, Status> { pub async fn server_streaming<S, B>(
let buf = (&self.buf).into_buf(); &mut self,
mut service: S,
req: http::Request<B>,
) -> http::Response<impl TryStream<Ok = crate::body::BytesBuf, Error = Status>>
where
S: ServerStreamingService<T::Decode, Response = T::Encode>,
B: Body,
B::Error: Into<crate::Error>,
{
let (_parts, body) = req.into_parts();
let stream = codec::decode(self.codec.decoder(), body).into_stream();
futures_util::pin_mut!(stream);
let message = stream.try_next().await.unwrap().unwrap();
let request = Request::new(message);
let response = service.call(request).await.unwrap();
let source = response.into_inner();
let body = codec::encode(self.codec.encoder(), source).await;
if let State::ReadHeader = self.state { http::Response::new(body)
if buf.remaining() < 5 { }
return Ok(None);
}
let is_compressed = match buf.get_u8() { pub async fn client_streaming<S, B>(
0 => false, &mut self,
1 => { mut service: S,
trace!("message compressed, compression not supported yet"); req: http::Request<B>,
return Err(crate::Status::new( ) -> http::Response<impl TryStream<Ok = crate::body::BytesBuf, Error = Status>>
crate::Code::Unimplemented, where
"Message compressed, compression not supported yet.".to_string(), S: ClientStreamingService<BoxStream<T::Decode>, Response = T::Encode>,
)); T::Decode: Send,
} T::Decoder: Send + 'static,
f => { B: Body + Send + 'static,
trace!("unexpected compression flag"); B::Data: Send,
return Err(crate::Status::new( B::Error: Into<crate::Error> + Send,
crate::Code::Internal, {
format!("Unexpected compression flag: {}", f), let (_parts, body) = req.into_parts();
)); let stream = codec::decode(self.codec.decoder(), body);
} let stream = Box::pin(stream) as BoxStream<T::Decode>;
}; let request = Request::new(stream);
let len = (&self.buf[..]).into_buf().get_u32_be() as usize; let response = service.call(request).await.unwrap();
let message = response.into_inner();
let source = stream::once(future::ok(message));
let body = codec::encode(self.codec.encoder(), source).await;
self.state = State::ReadBody { http::Response::new(body)
compression: is_compressed, }
len,
}
}
if let State::ReadBody { len, .. } = self.state { pub async fn streaming<S, B>(
if buf.remaining() < len { &mut self,
return Ok(None); mut service: S,
} req: http::Request<B>,
) -> http::Response<impl TryStream<Ok = crate::body::BytesBuf, Error = Status>>
where
S: StreamingService<BoxStream<T::Decode>, Response = T::Encode>,
T::Decode: Send,
T::Decoder: Send + 'static,
B: Body + Send + 'static,
B::Data: Send,
B::Error: Into<crate::Error> + Send,
{
let (_parts, body) = req.into_parts();
let stream = codec::decode(self.codec.decoder(), body);
let stream = Box::pin(stream) as BoxStream<T::Decode>;
let request = Request::new(stream);
let response = service.call(request).await.unwrap();
let source = response.into_inner();
let body = codec::encode(self.codec.encoder(), source).await;
match self.decoder.decode(&mut self.buf) { http::Response::new(body)
Ok(Some(msg)) => {
self.state = State::ReadHeader;
return Ok(Some(msg));
}
Err(e) => {
return Err(e);
}
}
}
Ok(None)
} }
} }
#[cfg(test)]
mod tests {
use crate::body::AsyncBody;
use crate::server::Encode;
use bytes::{Bytes, BytesMut};
use tokio_codec::BytesCodec;
#[test]
fn body() {
let stream = futures_util::stream::iter(vec![Ok(Bytes::new())]);
let mut encode = Encode::new(BytesCodec::new(), stream);
let mut buf = BytesMut::with_capacity(1024);
AsyncBody::new(Box::pin(encode.encode(&mut buf)));
}
}
// impl<T, U> http_body::Body for Encode<T, U>
// pub struct Grpc<T> {
// opdec: T,
// }
// impl<T: Codec> Grpc<T> {
// pub async fn unary<B>(&mut self, message: B) -> Result<Response<B>> {
// self.server_streaming(stream::once(message)).await
// }
// pub async fn server_streaming<B>(
// &mut self,
// stream: impl Stream,
// ) -> Result<Response<impl http_body::Body>> {
// unimplemetned!()
// }
// fn map_request<B>(&mut self, request: http::Request<B>) -> Request<B> {
// Request::from_http(request)
// }
// }
+6
View File
@@ -289,6 +289,12 @@ impl From<Status> for h2::Error {
} }
} }
impl From<std::io::Error> for Status {
fn from(_io: std::io::Error) -> Self {
unimplemented!()
}
}
impl fmt::Display for Status { impl fmt::Display for Status {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!( write!(