chore: Reduce how much code gets monomorphized (#1032)

This commit is contained in:
Markus Westerlind
2022-08-23 14:29:15 -04:00
committed by GitHub
parent 0a2a2f3ac5
commit 523a55044f
5 changed files with 378 additions and 289 deletions
+92 -57
View File
@@ -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
View File
@@ -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
View File
@@ -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()?)))
}
}
} }
} }
+7
View File
@@ -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| {
+1 -7
View File
@@ -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")]