chore: Reduce how much code gets monomorphized (#1032)
This commit is contained in:
+92
-57
@@ -2,7 +2,7 @@ use crate::codec::compression::{CompressionEncoding, EnabledCompressionEncodings
|
|||||||
use crate::{
|
use crate::{
|
||||||
body::BoxBody,
|
body::BoxBody,
|
||||||
client::GrpcService,
|
client::GrpcService,
|
||||||
codec::{encode_client, Codec, Streaming},
|
codec::{encode_client, Codec, Decoder, Streaming},
|
||||||
request::SanitizeHeaders,
|
request::SanitizeHeaders,
|
||||||
Code, Request, Response, Status,
|
Code, Request, Response, Status,
|
||||||
};
|
};
|
||||||
@@ -30,6 +30,10 @@ use std::fmt;
|
|||||||
/// [gRPC protocol definition]: https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests
|
/// [gRPC protocol definition]: https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests
|
||||||
pub struct Grpc<T> {
|
pub struct Grpc<T> {
|
||||||
inner: T,
|
inner: T,
|
||||||
|
config: GrpcConfig,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct GrpcConfig {
|
||||||
origin: Uri,
|
origin: Uri,
|
||||||
/// Which compression encodings does the client accept?
|
/// Which compression encodings does the client accept?
|
||||||
accept_compression_encodings: EnabledCompressionEncodings,
|
accept_compression_encodings: EnabledCompressionEncodings,
|
||||||
@@ -40,12 +44,7 @@ pub struct Grpc<T> {
|
|||||||
impl<T> Grpc<T> {
|
impl<T> Grpc<T> {
|
||||||
/// Creates a new gRPC client with the provided [`GrpcService`].
|
/// Creates a new gRPC client with the provided [`GrpcService`].
|
||||||
pub fn new(inner: T) -> Self {
|
pub fn new(inner: T) -> Self {
|
||||||
Self {
|
Self::with_origin(inner, Uri::default())
|
||||||
inner,
|
|
||||||
origin: Uri::default(),
|
|
||||||
send_compression_encodings: None,
|
|
||||||
accept_compression_encodings: EnabledCompressionEncodings::default(),
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Creates a new gRPC client with the provided [`GrpcService`] and `Uri`.
|
/// Creates a new gRPC client with the provided [`GrpcService`] and `Uri`.
|
||||||
@@ -55,9 +54,11 @@ impl<T> Grpc<T> {
|
|||||||
pub fn with_origin(inner: T, origin: Uri) -> Self {
|
pub fn with_origin(inner: T, origin: Uri) -> Self {
|
||||||
Self {
|
Self {
|
||||||
inner,
|
inner,
|
||||||
origin,
|
config: GrpcConfig {
|
||||||
send_compression_encodings: None,
|
origin,
|
||||||
accept_compression_encodings: EnabledCompressionEncodings::default(),
|
send_compression_encodings: None,
|
||||||
|
accept_compression_encodings: EnabledCompressionEncodings::default(),
|
||||||
|
},
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -88,7 +89,7 @@ impl<T> Grpc<T> {
|
|||||||
/// # };
|
/// # };
|
||||||
/// ```
|
/// ```
|
||||||
pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self {
|
pub fn send_compressed(mut self, encoding: CompressionEncoding) -> Self {
|
||||||
self.send_compression_encodings = Some(encoding);
|
self.config.send_compression_encodings = Some(encoding);
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -119,7 +120,7 @@ impl<T> Grpc<T> {
|
|||||||
/// # };
|
/// # };
|
||||||
/// ```
|
/// ```
|
||||||
pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self {
|
pub fn accept_compressed(mut self, encoding: CompressionEncoding) -> Self {
|
||||||
self.accept_compression_encodings.enable(encoding);
|
self.config.accept_compression_encodings.enable(encoding);
|
||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -226,6 +227,73 @@ impl<T> Grpc<T> {
|
|||||||
M1: Send + Sync + 'static,
|
M1: Send + Sync + 'static,
|
||||||
M2: Send + Sync + 'static,
|
M2: Send + Sync + 'static,
|
||||||
{
|
{
|
||||||
|
let request = request
|
||||||
|
.map(|s| encode_client(codec.encoder(), s, self.config.send_compression_encodings))
|
||||||
|
.map(BoxBody::new);
|
||||||
|
|
||||||
|
let request = self.config.prepare_request(request, path);
|
||||||
|
|
||||||
|
let response = self
|
||||||
|
.inner
|
||||||
|
.call(request)
|
||||||
|
.await
|
||||||
|
.map_err(Status::from_error_generic)?;
|
||||||
|
|
||||||
|
let decoder = codec.decoder();
|
||||||
|
|
||||||
|
self.create_response(decoder, response)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Keeping this code in a separate function from Self::streaming lets functions that return the
|
||||||
|
// same output share the generated binary code
|
||||||
|
fn create_response<M2>(
|
||||||
|
&self,
|
||||||
|
decoder: impl Decoder<Item = M2, Error = Status> + Send + 'static,
|
||||||
|
response: http::Response<T::ResponseBody>,
|
||||||
|
) -> Result<Response<Streaming<M2>>, Status>
|
||||||
|
where
|
||||||
|
T: GrpcService<BoxBody>,
|
||||||
|
T::ResponseBody: Body + Send + 'static,
|
||||||
|
<T::ResponseBody as Body>::Error: Into<crate::Error>,
|
||||||
|
{
|
||||||
|
let encoding = CompressionEncoding::from_encoding_header(
|
||||||
|
response.headers(),
|
||||||
|
self.config.accept_compression_encodings,
|
||||||
|
)?;
|
||||||
|
|
||||||
|
let status_code = response.status();
|
||||||
|
let trailers_only_status = Status::from_header_map(response.headers());
|
||||||
|
|
||||||
|
// We do not need to check for trailers if the `grpc-status` header is present
|
||||||
|
// with a valid code.
|
||||||
|
let expect_additional_trailers = if let Some(status) = trailers_only_status {
|
||||||
|
if status.code() != Code::Ok {
|
||||||
|
return Err(status);
|
||||||
|
}
|
||||||
|
|
||||||
|
false
|
||||||
|
} else {
|
||||||
|
true
|
||||||
|
};
|
||||||
|
|
||||||
|
let response = response.map(|body| {
|
||||||
|
if expect_additional_trailers {
|
||||||
|
Streaming::new_response(decoder, body, status_code, encoding)
|
||||||
|
} else {
|
||||||
|
Streaming::new_empty(decoder, body)
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
Ok(Response::from_http(response))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl GrpcConfig {
|
||||||
|
fn prepare_request(
|
||||||
|
&self,
|
||||||
|
request: Request<http_body::combinators::UnsyncBoxBody<bytes::Bytes, Status>>,
|
||||||
|
path: PathAndQuery,
|
||||||
|
) -> http::Request<http_body::combinators::UnsyncBoxBody<bytes::Bytes, Status>> {
|
||||||
let scheme = self.origin.scheme().cloned();
|
let scheme = self.origin.scheme().cloned();
|
||||||
let authority = self.origin.authority().cloned();
|
let authority = self.origin.authority().cloned();
|
||||||
|
|
||||||
@@ -236,10 +304,6 @@ impl<T> Grpc<T> {
|
|||||||
|
|
||||||
let uri = Uri::from_parts(parts).expect("path_and_query only is valid Uri");
|
let uri = Uri::from_parts(parts).expect("path_and_query only is valid Uri");
|
||||||
|
|
||||||
let request = request
|
|
||||||
.map(|s| encode_client(codec.encoder(), s, self.send_compression_encodings))
|
|
||||||
.map(BoxBody::new);
|
|
||||||
|
|
||||||
let mut request = request.into_http(
|
let mut request = request.into_http(
|
||||||
uri,
|
uri,
|
||||||
http::Method::POST,
|
http::Method::POST,
|
||||||
@@ -274,41 +338,7 @@ impl<T> Grpc<T> {
|
|||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
let response = self
|
request
|
||||||
.inner
|
|
||||||
.call(request)
|
|
||||||
.await
|
|
||||||
.map_err(|err| Status::from_error(err.into()))?;
|
|
||||||
|
|
||||||
let encoding = CompressionEncoding::from_encoding_header(
|
|
||||||
response.headers(),
|
|
||||||
self.accept_compression_encodings,
|
|
||||||
)?;
|
|
||||||
|
|
||||||
let status_code = response.status();
|
|
||||||
let trailers_only_status = Status::from_header_map(response.headers());
|
|
||||||
|
|
||||||
// We do not need to check for trailers if the `grpc-status` header is present
|
|
||||||
// with a valid code.
|
|
||||||
let expect_additional_trailers = if let Some(status) = trailers_only_status {
|
|
||||||
if status.code() != Code::Ok {
|
|
||||||
return Err(status);
|
|
||||||
}
|
|
||||||
|
|
||||||
false
|
|
||||||
} else {
|
|
||||||
true
|
|
||||||
};
|
|
||||||
|
|
||||||
let response = response.map(|body| {
|
|
||||||
if expect_additional_trailers {
|
|
||||||
Streaming::new_response(codec.decoder(), body, status_code, encoding)
|
|
||||||
} else {
|
|
||||||
Streaming::new_empty(codec.decoder(), body)
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
Ok(Response::from_http(response))
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -316,9 +346,11 @@ impl<T: Clone> Clone for Grpc<T> {
|
|||||||
fn clone(&self) -> Self {
|
fn clone(&self) -> Self {
|
||||||
Self {
|
Self {
|
||||||
inner: self.inner.clone(),
|
inner: self.inner.clone(),
|
||||||
origin: self.origin.clone(),
|
config: GrpcConfig {
|
||||||
send_compression_encodings: self.send_compression_encodings,
|
origin: self.config.origin.clone(),
|
||||||
accept_compression_encodings: self.accept_compression_encodings,
|
send_compression_encodings: self.config.send_compression_encodings,
|
||||||
|
accept_compression_encodings: self.config.accept_compression_encodings,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -329,13 +361,16 @@ impl<T: fmt::Debug> fmt::Debug for Grpc<T> {
|
|||||||
|
|
||||||
f.field("inner", &self.inner);
|
f.field("inner", &self.inner);
|
||||||
|
|
||||||
f.field("origin", &self.origin);
|
f.field("origin", &self.config.origin);
|
||||||
|
|
||||||
f.field("compression_encoding", &self.send_compression_encodings);
|
f.field(
|
||||||
|
"compression_encoding",
|
||||||
|
&self.config.send_compression_encodings,
|
||||||
|
);
|
||||||
|
|
||||||
f.field(
|
f.field(
|
||||||
"accept_compression_encodings",
|
"accept_compression_encodings",
|
||||||
&self.accept_compression_encodings,
|
&self.config.accept_compression_encodings,
|
||||||
);
|
);
|
||||||
|
|
||||||
f.finish()
|
f.finish()
|
||||||
|
|||||||
+169
-143
@@ -21,6 +21,10 @@ const BUFFER_SIZE: usize = 8 * 1024;
|
|||||||
/// to fetch the message stream and trailing metadata
|
/// to fetch the message stream and trailing metadata
|
||||||
pub struct Streaming<T> {
|
pub struct Streaming<T> {
|
||||||
decoder: Box<dyn Decoder<Item = T, Error = Status> + Send + 'static>,
|
decoder: Box<dyn Decoder<Item = T, Error = Status> + Send + 'static>,
|
||||||
|
inner: StreamingInner,
|
||||||
|
}
|
||||||
|
|
||||||
|
struct StreamingInner {
|
||||||
body: BoxBody,
|
body: BoxBody,
|
||||||
state: State,
|
state: State,
|
||||||
direction: Direction,
|
direction: Direction,
|
||||||
@@ -96,20 +100,157 @@ impl<T> Streaming<T> {
|
|||||||
{
|
{
|
||||||
Self {
|
Self {
|
||||||
decoder: Box::new(decoder),
|
decoder: Box::new(decoder),
|
||||||
body: body
|
inner: StreamingInner {
|
||||||
.map_data(|mut buf| buf.copy_to_bytes(buf.remaining()))
|
body: body
|
||||||
.map_err(|err| Status::map_error(err.into()))
|
.map_data(|mut buf| buf.copy_to_bytes(buf.remaining()))
|
||||||
.boxed_unsync(),
|
.map_err(|err| Status::map_error(err.into()))
|
||||||
state: State::ReadHeader,
|
.boxed_unsync(),
|
||||||
direction,
|
state: State::ReadHeader,
|
||||||
buf: BytesMut::with_capacity(BUFFER_SIZE),
|
direction,
|
||||||
trailers: None,
|
buf: BytesMut::with_capacity(BUFFER_SIZE),
|
||||||
decompress_buf: BytesMut::new(),
|
trailers: None,
|
||||||
encoding,
|
decompress_buf: BytesMut::new(),
|
||||||
|
encoding,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl StreamingInner {
|
||||||
|
fn decode_chunk(&mut self) -> Result<Option<DecodeBuf<'_>>, Status> {
|
||||||
|
if let State::ReadHeader = self.state {
|
||||||
|
if self.buf.remaining() < HEADER_SIZE {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
let compression_encoding = match self.buf.get_u8() {
|
||||||
|
0 => None,
|
||||||
|
1 => {
|
||||||
|
{
|
||||||
|
if self.encoding.is_some() {
|
||||||
|
self.encoding
|
||||||
|
} else {
|
||||||
|
// https://grpc.github.io/grpc/core/md_doc_compression.html
|
||||||
|
// An ill-constructed message with its Compressed-Flag bit set but lacking a grpc-encoding
|
||||||
|
// entry different from identity in its metadata MUST fail with INTERNAL status,
|
||||||
|
// its associated description indicating the invalid Compressed-Flag condition.
|
||||||
|
return Err(Status::new(Code::Internal, "protocol error: received message with compressed-flag but no grpc-encoding was specified"));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
f => {
|
||||||
|
trace!("unexpected compression flag");
|
||||||
|
let message = if let Direction::Response(status) = self.direction {
|
||||||
|
format!(
|
||||||
|
"protocol error: received message with invalid compression flag: {} (valid flags are 0 and 1) while receiving response with status: {}",
|
||||||
|
f, status
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
format!("protocol error: received message with invalid compression flag: {} (valid flags are 0 and 1), while sending request", f)
|
||||||
|
};
|
||||||
|
return Err(Status::new(Code::Internal, message));
|
||||||
|
}
|
||||||
|
};
|
||||||
|
let len = self.buf.get_u32() as usize;
|
||||||
|
self.buf.reserve(len);
|
||||||
|
|
||||||
|
self.state = State::ReadBody {
|
||||||
|
compression: compression_encoding,
|
||||||
|
len,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
if let State::ReadBody { len, compression } = self.state {
|
||||||
|
// if we haven't read enough of the message then return and keep
|
||||||
|
// reading
|
||||||
|
if self.buf.remaining() < len || self.buf.len() < len {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
let decode_buf = if let Some(encoding) = compression {
|
||||||
|
self.decompress_buf.clear();
|
||||||
|
|
||||||
|
if let Err(err) = decompress(encoding, &mut self.buf, &mut self.decompress_buf, len)
|
||||||
|
{
|
||||||
|
let message = if let Direction::Response(status) = self.direction {
|
||||||
|
format!(
|
||||||
|
"Error decompressing: {}, while receiving response with status: {}",
|
||||||
|
err, status
|
||||||
|
)
|
||||||
|
} else {
|
||||||
|
format!("Error decompressing: {}, while sending request", err)
|
||||||
|
};
|
||||||
|
return Err(Status::new(Code::Internal, message));
|
||||||
|
}
|
||||||
|
let decompressed_len = self.decompress_buf.len();
|
||||||
|
DecodeBuf::new(&mut self.decompress_buf, decompressed_len)
|
||||||
|
} else {
|
||||||
|
DecodeBuf::new(&mut self.buf, len)
|
||||||
|
};
|
||||||
|
|
||||||
|
return Ok(Some(decode_buf));
|
||||||
|
}
|
||||||
|
|
||||||
|
Ok(None)
|
||||||
|
}
|
||||||
|
|
||||||
|
// Returns Some(()) if data was found or None if the loop in `poll_next` should break
|
||||||
|
fn poll_data(&mut self, cx: &mut Context<'_>) -> Poll<Result<Option<()>, Status>> {
|
||||||
|
let chunk = match ready!(Pin::new(&mut self.body).poll_data(cx)) {
|
||||||
|
Some(Ok(d)) => Some(d),
|
||||||
|
Some(Err(e)) => {
|
||||||
|
let _ = std::mem::replace(&mut self.state, State::Error);
|
||||||
|
let err: crate::Error = e.into();
|
||||||
|
debug!("decoder inner stream error: {:?}", err);
|
||||||
|
let status = Status::from_error(err);
|
||||||
|
return Poll::Ready(Err(status));
|
||||||
|
}
|
||||||
|
None => None,
|
||||||
|
};
|
||||||
|
|
||||||
|
Poll::Ready(if let Some(data) = chunk {
|
||||||
|
self.buf.put(data);
|
||||||
|
Ok(Some(()))
|
||||||
|
} else {
|
||||||
|
// FIXME: improve buf usage.
|
||||||
|
if self.buf.has_remaining() {
|
||||||
|
trace!("unexpected EOF decoding stream");
|
||||||
|
Err(Status::new(
|
||||||
|
Code::Internal,
|
||||||
|
"Unexpected EOF decoding stream.".to_string(),
|
||||||
|
))
|
||||||
|
} else {
|
||||||
|
Ok(None)
|
||||||
|
}
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
fn poll_response(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Status>> {
|
||||||
|
if let Direction::Response(status) = self.direction {
|
||||||
|
match ready!(Pin::new(&mut self.body).poll_trailers(cx)) {
|
||||||
|
Ok(trailer) => {
|
||||||
|
if let Err(e) = crate::status::infer_grpc_status(trailer.as_ref(), status) {
|
||||||
|
if let Some(e) = e {
|
||||||
|
return Poll::Ready(Err(e));
|
||||||
|
} else {
|
||||||
|
return Poll::Ready(Ok(()));
|
||||||
|
}
|
||||||
|
} else {
|
||||||
|
self.trailers = trailer.map(MetadataMap::from_headers);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Err(e) => {
|
||||||
|
let err: crate::Error = e.into();
|
||||||
|
debug!("decoder inner trailers error: {:?}", err);
|
||||||
|
let status = Status::from_error(err);
|
||||||
|
return Poll::Ready(Err(status));
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
Poll::Ready(Ok(()))
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl<T> Streaming<T> {
|
impl<T> Streaming<T> {
|
||||||
/// Fetch the next message from this stream.
|
/// Fetch the next message from this stream.
|
||||||
///
|
///
|
||||||
@@ -165,7 +306,7 @@ impl<T> Streaming<T> {
|
|||||||
pub async fn trailers(&mut self) -> Result<Option<MetadataMap>, Status> {
|
pub async fn trailers(&mut self) -> Result<Option<MetadataMap>, Status> {
|
||||||
// Shortcut to see if we already pulled the trailers in the stream step
|
// Shortcut to see if we already pulled the trailers in the stream step
|
||||||
// we need to do that so that the stream can error on trailing grpc-status
|
// we need to do that so that the stream can error on trailing grpc-status
|
||||||
if let Some(trailers) = self.trailers.take() {
|
if let Some(trailers) = self.inner.trailers.take() {
|
||||||
return Ok(Some(trailers));
|
return Ok(Some(trailers));
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -174,13 +315,13 @@ impl<T> Streaming<T> {
|
|||||||
|
|
||||||
// Since we call poll_trailers internally on poll_next we need to
|
// Since we call poll_trailers internally on poll_next we need to
|
||||||
// check if it got cached again.
|
// check if it got cached again.
|
||||||
if let Some(trailers) = self.trailers.take() {
|
if let Some(trailers) = self.inner.trailers.take() {
|
||||||
return Ok(Some(trailers));
|
return Ok(Some(trailers));
|
||||||
}
|
}
|
||||||
|
|
||||||
// Trailers were not caught during poll_next and thus lets poll for
|
// Trailers were not caught during poll_next and thus lets poll for
|
||||||
// them manually.
|
// them manually.
|
||||||
let map = future::poll_fn(|cx| Pin::new(&mut self.body).poll_trailers(cx))
|
let map = future::poll_fn(|cx| Pin::new(&mut self.inner.body).poll_trailers(cx))
|
||||||
.await
|
.await
|
||||||
.map_err(|e| Status::from_error(Box::new(e)));
|
.map_err(|e| Status::from_error(Box::new(e)));
|
||||||
|
|
||||||
@@ -188,90 +329,16 @@ impl<T> Streaming<T> {
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn decode_chunk(&mut self) -> Result<Option<T>, Status> {
|
fn decode_chunk(&mut self) -> Result<Option<T>, Status> {
|
||||||
if let State::ReadHeader = self.state {
|
match self.inner.decode_chunk()? {
|
||||||
if self.buf.remaining() < HEADER_SIZE {
|
Some(mut decode_buf) => match self.decoder.decode(&mut decode_buf)? {
|
||||||
return Ok(None);
|
Some(msg) => {
|
||||||
}
|
self.inner.state = State::ReadHeader;
|
||||||
|
|
||||||
let compression_encoding = match self.buf.get_u8() {
|
|
||||||
0 => None,
|
|
||||||
1 => {
|
|
||||||
{
|
|
||||||
if self.encoding.is_some() {
|
|
||||||
self.encoding
|
|
||||||
} else {
|
|
||||||
// https://grpc.github.io/grpc/core/md_doc_compression.html
|
|
||||||
// An ill-constructed message with its Compressed-Flag bit set but lacking a grpc-encoding
|
|
||||||
// entry different from identity in its metadata MUST fail with INTERNAL status,
|
|
||||||
// its associated description indicating the invalid Compressed-Flag condition.
|
|
||||||
return Err(Status::new(Code::Internal, "protocol error: received message with compressed-flag but no grpc-encoding was specified"));
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
f => {
|
|
||||||
trace!("unexpected compression flag");
|
|
||||||
let message = if let Direction::Response(status) = self.direction {
|
|
||||||
format!(
|
|
||||||
"protocol error: received message with invalid compression flag: {} (valid flags are 0 and 1) while receiving response with status: {}",
|
|
||||||
f, status
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
format!("protocol error: received message with invalid compression flag: {} (valid flags are 0 and 1), while sending request", f)
|
|
||||||
};
|
|
||||||
return Err(Status::new(Code::Internal, message));
|
|
||||||
}
|
|
||||||
};
|
|
||||||
let len = self.buf.get_u32() as usize;
|
|
||||||
self.buf.reserve(len);
|
|
||||||
|
|
||||||
self.state = State::ReadBody {
|
|
||||||
compression: compression_encoding,
|
|
||||||
len,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
if let State::ReadBody { len, compression } = self.state {
|
|
||||||
// if we haven't read enough of the message then return and keep
|
|
||||||
// reading
|
|
||||||
if self.buf.remaining() < len || self.buf.len() < len {
|
|
||||||
return Ok(None);
|
|
||||||
}
|
|
||||||
|
|
||||||
let decoding_result = if let Some(encoding) = compression {
|
|
||||||
self.decompress_buf.clear();
|
|
||||||
|
|
||||||
if let Err(err) = decompress(encoding, &mut self.buf, &mut self.decompress_buf, len)
|
|
||||||
{
|
|
||||||
let message = if let Direction::Response(status) = self.direction {
|
|
||||||
format!(
|
|
||||||
"Error decompressing: {}, while receiving response with status: {}",
|
|
||||||
err, status
|
|
||||||
)
|
|
||||||
} else {
|
|
||||||
format!("Error decompressing: {}, while sending request", err)
|
|
||||||
};
|
|
||||||
return Err(Status::new(Code::Internal, message));
|
|
||||||
}
|
|
||||||
let decompressed_len = self.decompress_buf.len();
|
|
||||||
self.decoder.decode(&mut DecodeBuf::new(
|
|
||||||
&mut self.decompress_buf,
|
|
||||||
decompressed_len,
|
|
||||||
))
|
|
||||||
} else {
|
|
||||||
self.decoder.decode(&mut DecodeBuf::new(&mut self.buf, len))
|
|
||||||
};
|
|
||||||
|
|
||||||
return match decoding_result {
|
|
||||||
Ok(Some(msg)) => {
|
|
||||||
self.state = State::ReadHeader;
|
|
||||||
Ok(Some(msg))
|
Ok(Some(msg))
|
||||||
}
|
}
|
||||||
Ok(None) => Ok(None),
|
None => Ok(None),
|
||||||
Err(e) => Err(e),
|
},
|
||||||
};
|
None => Ok(None),
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(None)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -280,7 +347,7 @@ impl<T> Stream for Streaming<T> {
|
|||||||
|
|
||||||
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
|
||||||
loop {
|
loop {
|
||||||
if let State::Error = &self.state {
|
if let State::Error = &self.inner.state {
|
||||||
return Poll::Ready(None);
|
return Poll::Ready(None);
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -291,57 +358,16 @@ impl<T> Stream for Streaming<T> {
|
|||||||
return Poll::Ready(Some(Ok(item)));
|
return Poll::Ready(Some(Ok(item)));
|
||||||
}
|
}
|
||||||
|
|
||||||
let chunk = match ready!(Pin::new(&mut self.body).poll_data(cx)) {
|
match ready!(self.inner.poll_data(cx))? {
|
||||||
Some(Ok(d)) => Some(d),
|
Some(()) => (),
|
||||||
Some(Err(e)) => {
|
None => break,
|
||||||
let _ = std::mem::replace(&mut self.state, State::Error);
|
|
||||||
let err: crate::Error = e.into();
|
|
||||||
debug!("decoder inner stream error: {:?}", err);
|
|
||||||
let status = Status::from_error(err);
|
|
||||||
return Poll::Ready(Some(Err(status)));
|
|
||||||
}
|
|
||||||
None => None,
|
|
||||||
};
|
|
||||||
|
|
||||||
if let Some(data) = chunk {
|
|
||||||
self.buf.put(data);
|
|
||||||
} else {
|
|
||||||
// FIXME: improve buf usage.
|
|
||||||
if self.buf.has_remaining() {
|
|
||||||
trace!("unexpected EOF decoding stream");
|
|
||||||
return Poll::Ready(Some(Err(Status::new(
|
|
||||||
Code::Internal,
|
|
||||||
"Unexpected EOF decoding stream.".to_string(),
|
|
||||||
))));
|
|
||||||
} else {
|
|
||||||
break;
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
if let Direction::Response(status) = self.direction {
|
Poll::Ready(match ready!(self.inner.poll_response(cx)) {
|
||||||
match ready!(Pin::new(&mut self.body).poll_trailers(cx)) {
|
Ok(()) => None,
|
||||||
Ok(trailer) => {
|
Err(err) => Some(Err(err)),
|
||||||
if let Err(e) = crate::status::infer_grpc_status(trailer.as_ref(), status) {
|
})
|
||||||
if let Some(e) = e {
|
|
||||||
return Some(Err(e)).into();
|
|
||||||
} else {
|
|
||||||
return Poll::Ready(None);
|
|
||||||
}
|
|
||||||
} else {
|
|
||||||
self.trailers = trailer.map(MetadataMap::from_headers);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
Err(e) => {
|
|
||||||
let err: crate::Error = e.into();
|
|
||||||
debug!("decoder inner trailers error: {:?}", err);
|
|
||||||
let status = Status::from_error(err);
|
|
||||||
return Some(Err(status)).into();
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
Poll::Ready(None)
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
+109
-82
@@ -58,66 +58,80 @@ where
|
|||||||
T: Encoder<Error = Status>,
|
T: Encoder<Error = Status>,
|
||||||
U: Stream<Item = Result<T::Item, Status>>,
|
U: Stream<Item = Result<T::Item, Status>>,
|
||||||
{
|
{
|
||||||
async_stream::stream! {
|
let mut buf = BytesMut::with_capacity(BUFFER_SIZE);
|
||||||
let mut buf = BytesMut::with_capacity(BUFFER_SIZE);
|
|
||||||
|
|
||||||
let compression_encoding = if compression_override == SingleMessageCompressionOverride::Disable {
|
let compression_encoding = if compression_override == SingleMessageCompressionOverride::Disable
|
||||||
None
|
{
|
||||||
} else {
|
None
|
||||||
compression_encoding
|
} else {
|
||||||
};
|
compression_encoding
|
||||||
|
};
|
||||||
|
|
||||||
let mut uncompression_buf = if compression_encoding.is_some() {
|
let mut uncompression_buf = if compression_encoding.is_some() {
|
||||||
BytesMut::with_capacity(BUFFER_SIZE)
|
BytesMut::with_capacity(BUFFER_SIZE)
|
||||||
} else {
|
} else {
|
||||||
BytesMut::new()
|
BytesMut::new()
|
||||||
};
|
};
|
||||||
|
|
||||||
futures_util::pin_mut!(source);
|
source.map(move |result| {
|
||||||
|
let item = result?;
|
||||||
|
|
||||||
loop {
|
encode_item(
|
||||||
match source.next().await {
|
&mut encoder,
|
||||||
Some(Ok(item)) => {
|
&mut buf,
|
||||||
buf.reserve(HEADER_SIZE);
|
&mut uncompression_buf,
|
||||||
unsafe {
|
compression_encoding,
|
||||||
buf.advance_mut(HEADER_SIZE);
|
item,
|
||||||
}
|
)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
if let Some(encoding) = compression_encoding {
|
fn encode_item<T>(
|
||||||
uncompression_buf.clear();
|
encoder: &mut T,
|
||||||
|
buf: &mut BytesMut,
|
||||||
encoder.encode(item, &mut EncodeBuf::new(&mut uncompression_buf))
|
uncompression_buf: &mut BytesMut,
|
||||||
.map_err(|err| Status::internal(format!("Error encoding: {}", err)))?;
|
compression_encoding: Option<CompressionEncoding>,
|
||||||
|
item: T::Item,
|
||||||
let uncompressed_len = uncompression_buf.len();
|
) -> Result<Bytes, Status>
|
||||||
|
where
|
||||||
compress(
|
T: Encoder<Error = Status>,
|
||||||
encoding,
|
{
|
||||||
&mut uncompression_buf,
|
buf.reserve(HEADER_SIZE);
|
||||||
&mut buf,
|
unsafe {
|
||||||
uncompressed_len,
|
buf.advance_mut(HEADER_SIZE);
|
||||||
).map_err(|err| Status::internal(format!("Error compressing: {}", err)))?;
|
|
||||||
} else {
|
|
||||||
encoder.encode(item, &mut EncodeBuf::new(&mut buf))
|
|
||||||
.map_err(|err| Status::internal(format!("Error encoding: {}", err)))?;
|
|
||||||
}
|
|
||||||
|
|
||||||
// now that we know length, we can write the header
|
|
||||||
let len = buf.len() - HEADER_SIZE;
|
|
||||||
assert!(len <= std::u32::MAX as usize);
|
|
||||||
{
|
|
||||||
let mut buf = &mut buf[..HEADER_SIZE];
|
|
||||||
buf.put_u8(compression_encoding.is_some() as u8);
|
|
||||||
buf.put_u32(len as u32);
|
|
||||||
}
|
|
||||||
|
|
||||||
yield Ok(buf.split_to(len + HEADER_SIZE).freeze());
|
|
||||||
},
|
|
||||||
Some(Err(status)) => yield Err(status),
|
|
||||||
None => break,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
if let Some(encoding) = compression_encoding {
|
||||||
|
uncompression_buf.clear();
|
||||||
|
|
||||||
|
encoder
|
||||||
|
.encode(item, &mut EncodeBuf::new(uncompression_buf))
|
||||||
|
.map_err(|err| Status::internal(format!("Error encoding: {}", err)))?;
|
||||||
|
|
||||||
|
let uncompressed_len = uncompression_buf.len();
|
||||||
|
|
||||||
|
compress(encoding, uncompression_buf, buf, uncompressed_len)
|
||||||
|
.map_err(|err| Status::internal(format!("Error compressing: {}", err)))?;
|
||||||
|
} else {
|
||||||
|
encoder
|
||||||
|
.encode(item, &mut EncodeBuf::new(buf))
|
||||||
|
.map_err(|err| Status::internal(format!("Error encoding: {}", err)))?;
|
||||||
|
}
|
||||||
|
|
||||||
|
// now that we know length, we can write the header
|
||||||
|
Ok(finish_encoding(compression_encoding, buf))
|
||||||
|
}
|
||||||
|
|
||||||
|
fn finish_encoding(compression_encoding: Option<CompressionEncoding>, buf: &mut BytesMut) -> Bytes {
|
||||||
|
let len = buf.len() - HEADER_SIZE;
|
||||||
|
assert!(len <= std::u32::MAX as usize);
|
||||||
|
{
|
||||||
|
let mut buf = &mut buf[..HEADER_SIZE];
|
||||||
|
buf.put_u8(compression_encoding.is_some() as u8);
|
||||||
|
buf.put_u32(len as u32);
|
||||||
|
}
|
||||||
|
|
||||||
|
buf.split_to(len + HEADER_SIZE).freeze()
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
@@ -131,6 +145,11 @@ enum Role {
|
|||||||
pub(crate) struct EncodeBody<S> {
|
pub(crate) struct EncodeBody<S> {
|
||||||
#[pin]
|
#[pin]
|
||||||
inner: S,
|
inner: S,
|
||||||
|
state: EncodeState,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
struct EncodeState {
|
||||||
error: Option<Status>,
|
error: Option<Status>,
|
||||||
role: Role,
|
role: Role,
|
||||||
is_end_stream: bool,
|
is_end_stream: bool,
|
||||||
@@ -143,18 +162,44 @@ where
|
|||||||
pub(crate) fn new_client(inner: S) -> Self {
|
pub(crate) fn new_client(inner: S) -> Self {
|
||||||
Self {
|
Self {
|
||||||
inner,
|
inner,
|
||||||
error: None,
|
state: EncodeState {
|
||||||
role: Role::Client,
|
error: None,
|
||||||
is_end_stream: false,
|
role: Role::Client,
|
||||||
|
is_end_stream: false,
|
||||||
|
},
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn new_server(inner: S) -> Self {
|
pub(crate) fn new_server(inner: S) -> Self {
|
||||||
Self {
|
Self {
|
||||||
inner,
|
inner,
|
||||||
error: None,
|
state: EncodeState {
|
||||||
role: Role::Server,
|
error: None,
|
||||||
is_end_stream: false,
|
role: Role::Server,
|
||||||
|
is_end_stream: false,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl EncodeState {
|
||||||
|
fn trailers(&mut self) -> Result<Option<HeaderMap>, Status> {
|
||||||
|
match self.role {
|
||||||
|
Role::Client => Ok(None),
|
||||||
|
Role::Server => {
|
||||||
|
if self.is_end_stream {
|
||||||
|
return Ok(None);
|
||||||
|
}
|
||||||
|
|
||||||
|
let status = if let Some(status) = self.error.take() {
|
||||||
|
self.is_end_stream = true;
|
||||||
|
status
|
||||||
|
} else {
|
||||||
|
Status::new(Code::Ok, "")
|
||||||
|
};
|
||||||
|
|
||||||
|
Ok(Some(status.to_header_map()?))
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -167,7 +212,7 @@ where
|
|||||||
type Error = Status;
|
type Error = Status;
|
||||||
|
|
||||||
fn is_end_stream(&self) -> bool {
|
fn is_end_stream(&self) -> bool {
|
||||||
self.is_end_stream
|
self.state.is_end_stream
|
||||||
}
|
}
|
||||||
|
|
||||||
fn poll_data(
|
fn poll_data(
|
||||||
@@ -177,10 +222,10 @@ where
|
|||||||
let mut self_proj = self.project();
|
let mut self_proj = self.project();
|
||||||
match ready!(self_proj.inner.try_poll_next_unpin(cx)) {
|
match ready!(self_proj.inner.try_poll_next_unpin(cx)) {
|
||||||
Some(Ok(d)) => Some(Ok(d)).into(),
|
Some(Ok(d)) => Some(Ok(d)).into(),
|
||||||
Some(Err(status)) => match self_proj.role {
|
Some(Err(status)) => match self_proj.state.role {
|
||||||
Role::Client => Some(Err(status)).into(),
|
Role::Client => Some(Err(status)).into(),
|
||||||
Role::Server => {
|
Role::Server => {
|
||||||
*self_proj.error = Some(status);
|
self_proj.state.error = Some(status);
|
||||||
None.into()
|
None.into()
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
@@ -192,24 +237,6 @@ where
|
|||||||
self: Pin<&mut Self>,
|
self: Pin<&mut Self>,
|
||||||
_cx: &mut Context<'_>,
|
_cx: &mut Context<'_>,
|
||||||
) -> Poll<Result<Option<HeaderMap>, Status>> {
|
) -> Poll<Result<Option<HeaderMap>, Status>> {
|
||||||
match self.role {
|
Poll::Ready(self.project().state.trailers())
|
||||||
Role::Client => Poll::Ready(Ok(None)),
|
|
||||||
Role::Server => {
|
|
||||||
let self_proj = self.project();
|
|
||||||
|
|
||||||
if *self_proj.is_end_stream {
|
|
||||||
return Poll::Ready(Ok(None));
|
|
||||||
}
|
|
||||||
|
|
||||||
let status = if let Some(status) = self_proj.error.take() {
|
|
||||||
*self_proj.is_end_stream = true;
|
|
||||||
status
|
|
||||||
} else {
|
|
||||||
Status::new(Code::Ok, "")
|
|
||||||
};
|
|
||||||
|
|
||||||
Poll::Ready(Ok(Some(status.to_header_map()?)))
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -303,6 +303,13 @@ impl Status {
|
|||||||
Status::new(Code::Unauthenticated, message)
|
Status::new(Code::Unauthenticated, message)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg_attr(not(feature = "transport"), allow(dead_code))]
|
||||||
|
pub(crate) fn from_error_generic(
|
||||||
|
err: impl Into<Box<dyn Error + Send + Sync + 'static>>,
|
||||||
|
) -> Status {
|
||||||
|
Self::from_error(err.into())
|
||||||
|
}
|
||||||
|
|
||||||
#[cfg_attr(not(feature = "transport"), allow(dead_code))]
|
#[cfg_attr(not(feature = "transport"), allow(dead_code))]
|
||||||
pub(crate) fn from_error(err: Box<dyn Error + Send + Sync + 'static>) -> Status {
|
pub(crate) fn from_error(err: Box<dyn Error + Send + Sync + 'static>) -> Status {
|
||||||
Status::try_from_error(err).unwrap_or_else(|err| {
|
Status::try_from_error(err).unwrap_or_else(|err| {
|
||||||
|
|||||||
@@ -23,13 +23,7 @@ where
|
|||||||
IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
|
IO: AsyncRead + AsyncWrite + Connected + Unpin + Send + 'static,
|
||||||
IE: Into<crate::Error>,
|
IE: Into<crate::Error>,
|
||||||
{
|
{
|
||||||
async_stream::try_stream! {
|
incoming.err_into().map_ok(ServerIo::new_io)
|
||||||
futures_util::pin_mut!(incoming);
|
|
||||||
|
|
||||||
while let Some(stream) = incoming.try_next().await? {
|
|
||||||
yield ServerIo::new_io(stream);
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(feature = "tls")]
|
#[cfg(feature = "tls")]
|
||||||
|
|||||||
Reference in New Issue
Block a user