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 {
type Response = tonic::Response<()>;
type Response = tonic::Response<tonic::body::BoxAsyncBody>;
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>> {
Ok(()).into()
@@ -81,60 +81,35 @@ pub fn server(attr: TokenStream, item: TokenStream) -> TokenStream {
match request.uri().path() {
"/helloworld.Greeter/SayHello" => {
// let kind = greeter::methods::SayHello(self.inner.clone());
// greeter::ResponseFuture { kind: greeter::Kind::SayHello(kind) }
self.inner.stream(request).await?;
let inner = self.inner.clone();
let fut = async move {
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!("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));
+10 -6
View File
@@ -32,13 +32,17 @@ impl MyGreeter {
Ok(Response::new(()))
}
pub async fn server_stream(&self, request: Request<()>) -> Result<impl Stream, Status> {
unimplemented!()
}
// pub async fn streaming(&self, request: Request<impl Stream>) -> Result<Response<impl Stream>, Status> {
// unimplemented!()
// }
pub async fn client_stream(&self, request: Request<impl Stream>) -> Result<(), Status> {
unimplemented!()
}
// pub async fn server_stream(&self, request: Request<()>) -> Result<impl Stream, Status> {
// unimplemented!()
// }
// pub async fn client_stream(&self, request: Request<impl Stream>) -> Result<(), Status> {
// unimplemented!()
// }
}
#[tokio::test]
+81
View File
@@ -4,10 +4,91 @@ use futures_core::Stream;
use futures_util::{ready, TryStreamExt};
use http::HeaderMap;
use http_body::Body;
use std::pin::Pin;
use std::task::{Context, Poll};
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> {
inner: S,
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 {
type Encode;
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
pub mod body;
pub mod codec;
#[doc(hidden)]
pub mod error;
pub mod metadata;
pub mod server;
mod body;
mod codec;
mod request;
mod response;
mod server;
mod status;
pub use request::Request;
@@ -40,6 +41,8 @@ pub mod _codegen {
pub use tower_service::Service;
pub type ResponseFuture<T> =
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 use http::*;
+148 -182
View File
@@ -1,206 +1,172 @@
#![allow(dead_code)]
use crate::{Code, Request, Response, Status};
use async_stream::stream;
use bytes::{Buf, BufMut, Bytes, BytesMut, IntoBuf};
use futures_core::{Stream, TryStream};
use futures_util::{future, stream, StreamExt, TryStreamExt};
use crate::{
codec::{self, Codec},
Request, Status,
};
use futures_core::{Future, TryStream};
use futures_util::{future, stream, TryStreamExt};
use http_body::Body;
use std::future::Future;
use tokio_codec::{Decoder, Encoder};
use tower_service::Service;
use tracing::{debug, trace};
use std::pin::Pin;
pub trait Codec {
type Encode;
type Decode;
pub struct Grpc<T> {
codec: T,
}
pub struct Encode<T, U> {
encoder: T,
source: U,
// type UnaryFuture<B> = Once<Ready<Result<B, Status>>>;
// type ResponseBody = impl Stream<Item = Result<crate::body::BytesBuf, Status>>;
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
T: Encoder,
U: TryStream<Ok = T::Item, Error = Status> + Unpin,
T: Codec,
T::Decode: Unpin + 'static,
T::Encode: Unpin + 'static,
{
pub fn new(encoder: T, source: U) -> Self {
Encode { encoder, source }
pub fn new(codec: T) -> Self {
Self { codec }
}
pub fn encode<'a>(
&'a mut self,
buf: &'a mut BytesMut,
) -> impl Stream<Item = Result<crate::body::BytesBuf, Status>> + 'a {
stream! {
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
pub async fn unary<S, B>(
&mut self,
mut service: S,
req: http::Request<B>,
) -> http::Response<impl TryStream<Ok = crate::body::BytesBuf, Error = Status>>
where
S: UnaryService<T::Decode, Response = T::Encode>,
B: Body,
B::Error: Into<crate::Error>,
{
stream! {
loop {
// TODO: use try_stream! and ?
if let Some(item) = self.decode_chunk().unwrap() {
yield Ok(item);
}
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 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 {
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;
}
}
}
}
http::Response::new(body)
}
fn decode_chunk(&mut self) -> Result<Option<T::Item>, Status> {
let buf = (&self.buf).into_buf();
pub async fn server_streaming<S, B>(
&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 {
if buf.remaining() < 5 {
return Ok(None);
}
http::Response::new(body)
}
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 = (&self.buf[..]).into_buf().get_u32_be() as usize;
pub async fn client_streaming<S, B>(
&mut self,
mut service: S,
req: http::Request<B>,
) -> http::Response<impl TryStream<Ok = crate::body::BytesBuf, Error = Status>>
where
S: ClientStreamingService<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 message = response.into_inner();
let source = stream::once(future::ok(message));
let body = codec::encode(self.codec.encoder(), source).await;
self.state = State::ReadBody {
compression: is_compressed,
len,
}
}
http::Response::new(body)
}
if let State::ReadBody { len, .. } = self.state {
if buf.remaining() < len {
return Ok(None);
}
pub async fn streaming<S, B>(
&mut self,
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) {
Ok(Some(msg)) => {
self.state = State::ReadHeader;
return Ok(Some(msg));
}
Err(e) => {
return Err(e);
}
}
}
Ok(None)
http::Response::new(body)
}
}
#[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 {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(