feat(codec): compression support (#692)

* Initial compression support

* Support configuring compression on `Server`

* Minor clean up

* Test that compression is actually happening

* Clean up some todos

* channels compressing requests

* Move compression to be on the codecs

* Test sending compressed request to server that doesn't support it

* Clean up a bit

* Compress server streams

* Compress client streams

* Bidirectional streaming compression

* Handle receiving unsupported encoding

* Clean up

* Add note to future self

* Support disabling compression for individual responses

* Add docs

* Add compression examples

* Disable compression behind feature flag

* Add some docs

* Make flate2 optional dependency

* Fix docs wording

* Format

* Reply with which encodings are supported

* Convert tests to use mocked io

* Fix lints

* Use separate counters

* Don't make a long stream

* Address review feedback
This commit is contained in:
David Pedersen
2021-07-02 11:25:03 -04:00
committed by GitHub
parent 7677ad6476
commit 0583cff80f
30 changed files with 2191 additions and 98 deletions
+4
View File
@@ -40,6 +40,7 @@ tls-roots-common = ["tls"]
tls-roots = ["tls-roots-common", "rustls-native-certs"]
tls-webpki-roots = ["tls-roots-common", "webpki-roots"]
prost = ["prost1", "prost-derive"]
compression = ["flate2"]
# [[bench]]
# name = "bench_main"
@@ -82,6 +83,9 @@ tokio-rustls = { version = "0.22", optional = true }
rustls-native-certs = { version = "0.5", optional = true }
webpki-roots = { version = "0.21.1", optional = true }
# compression
flate2 = { version = "1.0", optional = true }
[dev-dependencies]
tokio = { version = "1.0", features = ["rt", "macros"] }
static_assertions = "1.0"
+1 -1
View File
@@ -22,7 +22,7 @@ macro_rules! bench {
b.iter(|| {
rt.block_on(async {
let decoder = MockDecoder::new($message_size);
let mut stream = Streaming::new_request(decoder, body.clone());
let mut stream = Streaming::new_request(decoder, body.clone(), None);
let mut count = 0;
while let Some(msg) = stream.message().await.unwrap() {
+152 -4
View File
@@ -1,3 +1,5 @@
#[cfg(feature = "compression")]
use crate::codec::compression::{CompressionEncoding, EnabledCompressionEncodings};
use crate::{
body::BoxBody,
client::GrpcService,
@@ -28,12 +30,102 @@ use std::fmt;
/// [gRPC protocol definition]: https://github.com/grpc/grpc/blob/master/doc/PROTOCOL-HTTP2.md#requests
pub struct Grpc<T> {
inner: T,
#[cfg(feature = "compression")]
/// Which compression encodings does the client accept?
accept_compression_encodings: EnabledCompressionEncodings,
#[cfg(feature = "compression")]
/// The compression encoding that will be applied to requests.
send_compression_encodings: Option<CompressionEncoding>,
}
impl<T> Grpc<T> {
/// Creates a new gRPC client with the provided [`GrpcService`].
pub fn new(inner: T) -> Self {
Self { inner }
Self {
inner,
#[cfg(feature = "compression")]
send_compression_encodings: None,
#[cfg(feature = "compression")]
accept_compression_encodings: EnabledCompressionEncodings::default(),
}
}
/// Compress requests with `gzip`.
///
/// Requires the server to accept `gzip` otherwise it might return an error.
///
/// # Example
///
/// The most common way of using this is through a client generated by tonic-build:
///
/// ```rust
/// use tonic::transport::Channel;
/// # struct TestClient<T>(T);
/// # impl<T> TestClient<T> {
/// # fn new(channel: T) -> Self { Self(channel) }
/// # fn send_gzip(self) -> Self { self }
/// # }
///
/// # async {
/// let channel = Channel::builder("127.0.0.1:3000".parse().unwrap())
/// .connect()
/// .await
/// .unwrap();
///
/// let client = TestClient::new(channel).send_gzip();
/// # };
/// ```
#[cfg(feature = "compression")]
#[cfg_attr(docsrs, doc(cfg(feature = "compression")))]
pub fn send_gzip(mut self) -> Self {
self.send_compression_encodings = Some(CompressionEncoding::Gzip);
self
}
#[doc(hidden)]
#[cfg(not(feature = "compression"))]
pub fn send_gzip(self) -> Self {
panic!(
"`send_gzip` called on a client but the `compression` feature is not enabled on tonic"
);
}
/// Enable accepting `gzip` compressed responses.
///
/// Requires the server to also support sending compressed responses.
///
/// # Example
///
/// The most common way of using this is through a client generated by tonic-build:
///
/// ```rust
/// use tonic::transport::Channel;
/// # struct TestClient<T>(T);
/// # impl<T> TestClient<T> {
/// # fn new(channel: T) -> Self { Self(channel) }
/// # fn accept_gzip(self) -> Self { self }
/// # }
///
/// # async {
/// let channel = Channel::builder("127.0.0.1:3000".parse().unwrap())
/// .connect()
/// .await
/// .unwrap();
///
/// let client = TestClient::new(channel).accept_gzip();
/// # };
/// ```
#[cfg(feature = "compression")]
#[cfg_attr(docsrs, doc(cfg(feature = "compression")))]
pub fn accept_gzip(mut self) -> Self {
self.accept_compression_encodings.enable_gzip();
self
}
#[doc(hidden)]
#[cfg(not(feature = "compression"))]
pub fn accept_gzip(self) -> Self {
panic!("`accept_gzip` called on a client but the `compression` feature is not enabled on tonic");
}
/// Check if the inner [`GrpcService`] is able to accept a new request.
@@ -145,7 +237,14 @@ impl<T> Grpc<T> {
let uri = Uri::from_parts(parts).expect("path_and_query only is valid Uri");
let request = request
.map(|s| encode_client(codec.encoder(), s))
.map(|s| {
encode_client(
codec.encoder(),
s,
#[cfg(feature = "compression")]
self.send_compression_encodings,
)
})
.map(BoxBody::new);
let mut request = request.into_http(uri);
@@ -160,12 +259,38 @@ impl<T> Grpc<T> {
.headers_mut()
.insert(CONTENT_TYPE, HeaderValue::from_static("application/grpc"));
#[cfg(feature = "compression")]
{
if let Some(encoding) = self.send_compression_encodings {
request.headers_mut().insert(
crate::codec::compression::ENCODING_HEADER,
encoding.into_header_value(),
);
}
if let Some(header_value) = self
.accept_compression_encodings
.into_accept_encoding_header_value()
{
request.headers_mut().insert(
crate::codec::compression::ACCEPT_ENCODING_HEADER,
header_value,
);
}
}
let response = self
.inner
.call(request)
.await
.map_err(|err| Status::from_error(err.into()))?;
#[cfg(feature = "compression")]
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());
@@ -183,7 +308,13 @@ impl<T> Grpc<T> {
let response = response.map(|body| {
if expect_additional_trailers {
Streaming::new_response(codec.decoder(), body, status_code)
Streaming::new_response(
codec.decoder(),
body,
status_code,
#[cfg(feature = "compression")]
encoding,
)
} else {
Streaming::new_empty(codec.decoder(), body)
}
@@ -197,12 +328,29 @@ impl<T: Clone> Clone for Grpc<T> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
#[cfg(feature = "compression")]
send_compression_encodings: self.send_compression_encodings,
#[cfg(feature = "compression")]
accept_compression_encodings: self.accept_compression_encodings,
}
}
}
impl<T: fmt::Debug> fmt::Debug for Grpc<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Grpc").field("inner", &self.inner).finish()
let mut f = f.debug_struct("Grpc");
f.field("inner", &self.inner);
#[cfg(feature = "compression")]
f.field("compression_encoding", &self.send_compression_encodings);
#[cfg(feature = "compression")]
f.field(
"accept_compression_encodings",
&self.accept_compression_encodings,
);
f.finish()
}
}
+189
View File
@@ -0,0 +1,189 @@
use super::encode::BUFFER_SIZE;
use crate::{metadata::MetadataValue, Status};
use bytes::{Buf, BufMut, BytesMut};
use flate2::read::{GzDecoder, GzEncoder};
use std::fmt;
pub(crate) const ENCODING_HEADER: &str = "grpc-encoding";
pub(crate) const ACCEPT_ENCODING_HEADER: &str = "grpc-accept-encoding";
/// Struct used to configure which encodings are enabled on a server or channel.
#[derive(Debug, Default, Clone, Copy)]
pub struct EnabledCompressionEncodings {
pub(crate) gzip: bool,
}
impl EnabledCompressionEncodings {
/// Check if `gzip` compression is enabled.
pub fn gzip(self) -> bool {
self.gzip
}
/// Enable `gzip` compression.
pub fn enable_gzip(&mut self) {
self.gzip = true;
}
pub(crate) fn into_accept_encoding_header_value(self) -> Option<http::HeaderValue> {
let Self { gzip } = self;
if gzip {
Some(http::HeaderValue::from_static("gzip,identity"))
} else {
None
}
}
}
/// The compression encodings Tonic supports.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum CompressionEncoding {
#[allow(missing_docs)]
Gzip,
}
impl CompressionEncoding {
/// Based on the `grpc-accept-encoding` header, pick an encoding to use.
pub(crate) fn from_accept_encoding_header(
map: &http::HeaderMap,
enabled_encodings: EnabledCompressionEncodings,
) -> Option<Self> {
let header_value = map.get(ACCEPT_ENCODING_HEADER)?;
let header_value_str = header_value.to_str().ok()?;
let EnabledCompressionEncodings { gzip } = enabled_encodings;
split_by_comma(header_value_str).find_map(|value| match value {
"gzip" if gzip => Some(CompressionEncoding::Gzip),
_ => None,
})
}
/// Get the value of `grpc-encoding` header. Returns an error if the encoding isn't supported.
pub(crate) fn from_encoding_header(
map: &http::HeaderMap,
enabled_encodings: EnabledCompressionEncodings,
) -> Result<Option<Self>, Status> {
let header_value = if let Some(value) = map.get(ENCODING_HEADER) {
value
} else {
return Ok(None);
};
let header_value_str = if let Ok(value) = header_value.to_str() {
value
} else {
return Ok(None);
};
let EnabledCompressionEncodings { gzip } = enabled_encodings;
match header_value_str {
"gzip" if gzip => Ok(Some(CompressionEncoding::Gzip)),
other => {
let mut status = Status::unimplemented(format!(
"Content is compressed with `{}` which isn't supported",
other
));
let header_value = enabled_encodings
.into_accept_encoding_header_value()
.map(MetadataValue::unchecked_from_header_value)
.unwrap_or_else(|| MetadataValue::from_static("identity"));
status
.metadata_mut()
.insert(ACCEPT_ENCODING_HEADER, header_value);
Err(status)
}
}
}
pub(crate) fn into_header_value(self) -> http::HeaderValue {
match self {
CompressionEncoding::Gzip => http::HeaderValue::from_static("gzip"),
}
}
}
impl fmt::Display for CompressionEncoding {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
CompressionEncoding::Gzip => write!(f, "gzip"),
}
}
}
fn split_by_comma(s: &str) -> impl Iterator<Item = &str> {
s.trim().split(',').map(|s| s.trim())
}
/// Compress `len` bytes from `decompressed_buf` into `out_buf`.
pub(crate) fn compress(
encoding: CompressionEncoding,
decompressed_buf: &mut BytesMut,
out_buf: &mut BytesMut,
len: usize,
) -> Result<(), std::io::Error> {
let capacity = ((len / BUFFER_SIZE) + 1) * BUFFER_SIZE;
out_buf.reserve(capacity);
match encoding {
CompressionEncoding::Gzip => {
let mut gzip_encoder = GzEncoder::new(
&decompressed_buf[0..len],
// FIXME: support customizing the compression level
flate2::Compression::new(6),
);
let mut out_writer = out_buf.writer();
std::io::copy(&mut gzip_encoder, &mut out_writer)?;
}
}
decompressed_buf.advance(len);
Ok(())
}
/// Decompress `len` bytes from `compressed_buf` into `out_buf`.
pub(crate) fn decompress(
encoding: CompressionEncoding,
compressed_buf: &mut BytesMut,
out_buf: &mut BytesMut,
len: usize,
) -> Result<(), std::io::Error> {
let estimate_decompressed_len = len * 2;
let capacity = ((estimate_decompressed_len / BUFFER_SIZE) + 1) * BUFFER_SIZE;
out_buf.reserve(capacity);
match encoding {
CompressionEncoding::Gzip => {
let mut gzip_decoder = GzDecoder::new(&compressed_buf[0..len]);
let mut out_writer = out_buf.writer();
std::io::copy(&mut gzip_decoder, &mut out_writer)?;
}
}
compressed_buf.advance(len);
Ok(())
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum SingleMessageCompressionOverride {
/// Inherit whatever compression is already configured. If the stream is compressed this
/// message will also be configured.
///
/// This is the default.
Inherit,
/// Don't compress this message, even if compression is enabled on the stream.
Disable,
}
impl Default for SingleMessageCompressionOverride {
fn default() -> Self {
Self::Inherit
}
}
+97 -18
View File
@@ -1,4 +1,6 @@
use super::{DecodeBuf, Decoder};
#[cfg(feature = "compression")]
use super::compression::{decompress, CompressionEncoding};
use super::{DecodeBuf, Decoder, HEADER_SIZE};
use crate::{body::BoxBody, metadata::MetadataMap, Code, Status};
use bytes::{Buf, BufMut, BytesMut};
use futures_core::Stream;
@@ -25,6 +27,10 @@ pub struct Streaming<T> {
direction: Direction,
buf: BytesMut,
trailers: Option<MetadataMap>,
#[cfg(feature = "compression")]
decompress_buf: BytesMut,
#[cfg(feature = "compression")]
encoding: Option<CompressionEncoding>,
}
impl<T> Unpin for Streaming<T> {}
@@ -43,13 +49,24 @@ enum Direction {
}
impl<T> Streaming<T> {
pub(crate) fn new_response<B, D>(decoder: D, body: B, status_code: StatusCode) -> Self
pub(crate) fn new_response<B, D>(
decoder: D,
body: B,
status_code: StatusCode,
#[cfg(feature = "compression")] encoding: Option<CompressionEncoding>,
) -> Self
where
B: Body + Send + Sync + 'static,
B::Error: Into<crate::Error>,
D: Decoder<Item = T, Error = Status> + Send + Sync + 'static,
{
Self::new(decoder, body, Direction::Response(status_code))
Self::new(
decoder,
body,
Direction::Response(status_code),
#[cfg(feature = "compression")]
encoding,
)
}
pub(crate) fn new_empty<B, D>(decoder: D, body: B) -> Self
@@ -58,20 +75,41 @@ impl<T> Streaming<T> {
B::Error: Into<crate::Error>,
D: Decoder<Item = T, Error = Status> + Send + Sync + 'static,
{
Self::new(decoder, body, Direction::EmptyResponse)
Self::new(
decoder,
body,
Direction::EmptyResponse,
#[cfg(feature = "compression")]
None,
)
}
#[doc(hidden)]
pub fn new_request<B, D>(decoder: D, body: B) -> Self
pub fn new_request<B, D>(
decoder: D,
body: B,
#[cfg(feature = "compression")] encoding: Option<CompressionEncoding>,
) -> Self
where
B: Body + Send + Sync + 'static,
B::Error: Into<crate::Error>,
D: Decoder<Item = T, Error = Status> + Send + Sync + 'static,
{
Self::new(decoder, body, Direction::Request)
Self::new(
decoder,
body,
Direction::Request,
#[cfg(feature = "compression")]
encoding,
)
}
fn new<B, D>(decoder: D, body: B, direction: Direction) -> Self
fn new<B, D>(
decoder: D,
body: B,
direction: Direction,
#[cfg(feature = "compression")] encoding: Option<CompressionEncoding>,
) -> Self
where
B: Body + Send + Sync + 'static,
B::Error: Into<crate::Error>,
@@ -87,6 +125,10 @@ impl<T> Streaming<T> {
direction,
buf: BytesMut::with_capacity(BUFFER_SIZE),
trailers: None,
#[cfg(feature = "compression")]
decompress_buf: BytesMut::new(),
#[cfg(feature = "compression")]
encoding,
}
}
}
@@ -156,18 +198,21 @@ impl<T> Streaming<T> {
fn decode_chunk(&mut self) -> Result<Option<T>, Status> {
if let State::ReadHeader = self.state {
if self.buf.remaining() < 5 {
if self.buf.remaining() < HEADER_SIZE {
return Ok(None);
}
let is_compressed = match self.buf.get_u8() {
0 => false,
1 => {
trace!("message compressed, compression not supported yet");
return Err(Status::new(
Code::Unimplemented,
"Message compressed, compression not supported yet.".to_string(),
));
if cfg!(feature = "compression") {
true
} else {
return Err(Status::new(
Code::Unimplemented,
"Message compressed, compression support not enabled.".to_string(),
));
}
}
f => {
trace!("unexpected compression flag");
@@ -191,17 +236,51 @@ impl<T> Streaming<T> {
}
}
if let State::ReadBody { len, .. } = &self.state {
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);
}
return match self
.decoder
.decode(&mut DecodeBuf::new(&mut self.buf, *len))
{
let decoding_result = if *compression {
#[cfg(feature = "compression")]
{
self.decompress_buf.clear();
if let Err(err) = decompress(
self.encoding.unwrap_or_else(|| {
unreachable!("message was compressed but `Streaming.encoding` was `None`. This is a bug in Tonic. Please file an issue")
}),
&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,
))
}
#[cfg(not(feature = "compression"))]
unreachable!("should not take this branch if compression is disabled")
} 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))
+76 -12
View File
@@ -1,4 +1,6 @@
use super::{EncodeBuf, Encoder};
#[cfg(feature = "compression")]
use super::compression::{compress, CompressionEncoding, SingleMessageCompressionOverride};
use super::{EncodeBuf, Encoder, HEADER_SIZE};
use crate::{Code, Status};
use bytes::{BufMut, Bytes, BytesMut};
use futures_core::{Stream, TryStream};
@@ -11,62 +13,124 @@ use std::{
task::{Context, Poll},
};
const BUFFER_SIZE: usize = 8 * 1024;
pub(super) const BUFFER_SIZE: usize = 8 * 1024;
pub(crate) fn encode_server<T, U>(
encoder: T,
source: U,
#[cfg(feature = "compression")] compression_encoding: Option<CompressionEncoding>,
#[cfg(feature = "compression")] compression_override: SingleMessageCompressionOverride,
) -> EncodeBody<impl Stream<Item = Result<Bytes, Status>>>
where
T: Encoder<Error = Status> + Send + Sync + 'static,
T::Item: Send + Sync,
U: Stream<Item = Result<T::Item, Status>> + Send + Sync + 'static,
{
let stream = encode(encoder, source).into_stream();
let stream = encode(
encoder,
source,
#[cfg(feature = "compression")]
compression_encoding,
#[cfg(feature = "compression")]
compression_override,
)
.into_stream();
EncodeBody::new_server(stream)
}
pub(crate) fn encode_client<T, U>(
encoder: T,
source: U,
#[cfg(feature = "compression")] compression_encoding: Option<CompressionEncoding>,
) -> EncodeBody<impl Stream<Item = Result<Bytes, Status>>>
where
T: Encoder<Error = Status> + Send + Sync + 'static,
T::Item: Send + Sync,
U: Stream<Item = T::Item> + Send + Sync + 'static,
{
let stream = encode(encoder, source.map(Ok)).into_stream();
let stream = encode(
encoder,
source.map(Ok),
#[cfg(feature = "compression")]
compression_encoding,
#[cfg(feature = "compression")]
SingleMessageCompressionOverride::default(),
)
.into_stream();
EncodeBody::new_client(stream)
}
fn encode<T, U>(mut encoder: T, source: U) -> impl TryStream<Ok = Bytes, Error = Status>
fn encode<T, U>(
mut encoder: T,
source: U,
#[cfg(feature = "compression")] compression_encoding: Option<CompressionEncoding>,
#[cfg(feature = "compression")] compression_override: SingleMessageCompressionOverride,
) -> impl TryStream<Ok = Bytes, Error = Status>
where
T: Encoder<Error = Status>,
U: Stream<Item = Result<T::Item, Status>>,
{
async_stream::stream! {
let mut buf = BytesMut::with_capacity(BUFFER_SIZE);
#[cfg(feature = "compression")]
let (compression_enabled_for_stream, mut uncompression_buf) = match compression_encoding {
Some(CompressionEncoding::Gzip) => (true, BytesMut::with_capacity(BUFFER_SIZE)),
None => (false, BytesMut::new()),
};
#[cfg(feature = "compression")]
let compress_item = compression_enabled_for_stream && compression_override == SingleMessageCompressionOverride::Inherit;
#[cfg(not(feature = "compression"))]
let compress_item = false;
futures_util::pin_mut!(source);
loop {
match source.next().await {
Some(Ok(item)) => {
buf.reserve(5);
buf.reserve(HEADER_SIZE);
unsafe {
buf.advance_mut(5);
buf.advance_mut(HEADER_SIZE);
}
if compress_item {
#[cfg(feature = "compression")]
{
uncompression_buf.clear();
encoder.encode(item, &mut EncodeBuf::new(&mut uncompression_buf))
.map_err(|err| Status::internal(format!("Error encoding: {}", err)))?;
let uncompressed_len = uncompression_buf.len();
compress(
compression_encoding.unwrap(),
&mut uncompression_buf,
&mut buf,
uncompressed_len,
).map_err(|err| Status::internal(format!("Error compressing: {}", err)))?;
}
#[cfg(not(feature = "compression"))]
unreachable!("compression disabled, should not take this branch");
} else {
encoder.encode(item, &mut EncodeBuf::new(&mut buf))
.map_err(|err| Status::internal(format!("Error encoding: {}", err)))?;
}
encoder.encode(item, &mut EncodeBuf::new(&mut buf)).map_err(drop).unwrap();
// now that we know length, we can write the header
let len = buf.len() - 5;
let len = buf.len() - HEADER_SIZE;
assert!(len <= std::u32::MAX as usize);
{
let mut buf = &mut buf[..5];
buf.put_u8(0); // byte must be 0, reserve doesn't auto-zero
let mut buf = &mut buf[..HEADER_SIZE];
buf.put_u8(compress_item as u8);
buf.put_u32(len as u32);
}
yield Ok(buf.split_to(len + 5).freeze());
yield Ok(buf.split_to(len + HEADER_SIZE).freeze());
},
Some(Err(status)) => yield Err(status),
None => break,
+16 -3
View File
@@ -4,20 +4,33 @@
//! and a protobuf codec based on prost.
mod buffer;
#[cfg(feature = "compression")]
pub(crate) mod compression;
mod decode;
mod encode;
#[cfg(feature = "prost")]
mod prost;
use crate::Status;
use std::io;
pub use self::decode::Streaming;
pub(crate) use self::encode::{encode_client, encode_server};
pub use self::buffer::{DecodeBuf, EncodeBuf};
#[cfg(feature = "compression")]
#[cfg_attr(docsrs, doc(cfg(feature = "compression")))]
pub use self::compression::{CompressionEncoding, EnabledCompressionEncodings};
pub use self::decode::Streaming;
#[cfg(feature = "prost")]
#[cfg_attr(docsrs, doc(cfg(feature = "prost")))]
pub use self::prost::ProstCodec;
use crate::Status;
pub use buffer::{DecodeBuf, EncodeBuf};
// 5 bytes
const HEADER_SIZE: usize =
// compression flag
std::mem::size_of::<u8>() +
// data length
std::mem::size_of::<u32>();
/// Trait that knows how to encode and decode gRPC messages.
pub trait Codec: Default {
+13 -4
View File
@@ -77,7 +77,10 @@ fn from_decode_error(error: prost1::DecodeError) -> crate::Status {
#[cfg(test)]
mod tests {
use crate::codec::{encode_server, DecodeBuf, Decoder, EncodeBuf, Encoder, Streaming};
use crate::codec::compression::SingleMessageCompressionOverride;
use crate::codec::{
encode_server, DecodeBuf, Decoder, EncodeBuf, Encoder, Streaming, HEADER_SIZE,
};
use crate::Status;
use bytes::{Buf, BufMut, BytesMut};
use http_body::Body;
@@ -92,7 +95,7 @@ mod tests {
let mut buf = BytesMut::new();
buf.reserve(msg.len() + 5);
buf.reserve(msg.len() + HEADER_SIZE);
buf.put_u8(0);
buf.put_u32(msg.len() as u32);
@@ -100,7 +103,7 @@ mod tests {
let body = body::MockBody::new(&buf[..], 10005, 0);
let mut stream = Streaming::new_request(decoder, body);
let mut stream = Streaming::new_request(decoder, body, None);
let mut i = 0usize;
while let Some(output_msg) = stream.message().await.unwrap() {
@@ -119,7 +122,12 @@ mod tests {
let messages = std::iter::repeat_with(move || Ok::<_, Status>(msg.clone())).take(10000);
let source = futures_util::stream::iter(messages);
let body = encode_server(encoder, source);
let body = encode_server(
encoder,
source,
None,
SingleMessageCompressionOverride::default(),
);
futures_util::pin_mut!(body);
@@ -216,6 +224,7 @@ mod tests {
}
}
#[allow(clippy::drop_ref)]
fn poll_trailers(
self: Pin<&mut Self>,
cx: &mut Context<'_>,
+2
View File
@@ -10,6 +10,8 @@ pub use std::sync::Arc;
pub use std::task::{Context, Poll};
pub use tower_service::Service;
pub type StdError = Box<dyn std::error::Error + Send + Sync + 'static>;
#[cfg(feature = "compression")]
pub use crate::codec::{CompressionEncoding, EnabledCompressionEncodings};
pub use crate::service::interceptor::InterceptedService;
pub use http_body::Body;
+4
View File
@@ -28,6 +28,9 @@
//! - `tls-webpki-roots`: Add the standard trust roots from the `webpki-roots` crate to
//! `rustls`-based gRPC clients. Not enabled by default.
//! - `prost`: Enables the [`prost`] based gRPC [`Codec`] implementation.
//! - `compression`: Enables compressing requests, responses, and streams. Note
//! that you must enable the `compression` feature on both `tonic` and
//! `tonic-build` to use it. Depends on [flate2]. Not enabled by default.
//!
//! # Structure
//!
@@ -62,6 +65,7 @@
//! [`rustls`]: https://docs.rs/rustls
//! [`client`]: client/index.html
//! [`transport`]: transport/index.html
//! [flate2]: https://crates.io/crates/flate2
#![recursion_limit = "256"]
#![allow(clippy::inconsistent_struct_constructor)]
+1 -2
View File
@@ -200,12 +200,11 @@ pub(crate) const GRPC_TIMEOUT_HEADER: &str = "grpc-timeout";
impl MetadataMap {
// Headers reserved by the gRPC protocol.
pub(crate) const GRPC_RESERVED_HEADERS: [&'static str; 7] = [
pub(crate) const GRPC_RESERVED_HEADERS: [&'static str; 6] = [
"te",
"user-agent",
"content-type",
"grpc-message",
"grpc-encoding",
"grpc-message-type",
"grpc-status",
];
+15
View File
@@ -107,6 +107,21 @@ impl<T> Response<T> {
pub fn extensions_mut(&mut self) -> &mut Extensions {
&mut self.extensions
}
/// Disable compression of the response body.
///
/// This disables compression of the body of this response, even if compression is enabled on
/// the server.
///
/// **Note**: This only has effect on responses to unary requests and responses to client to
/// server streams. Response streams (server to client stream and bidirectional streams) will
/// still be compressed according to the configuration of the server.
#[cfg(feature = "compression")]
#[cfg_attr(docsrs, doc(cfg(feature = "compression")))]
pub fn disable_compression(&mut self) {
self.extensions_mut()
.insert(crate::codec::compression::SingleMessageCompressionOverride::Disable);
}
}
#[cfg(test)]
+308 -26
View File
@@ -1,3 +1,7 @@
#[cfg(feature = "compression")]
use crate::codec::compression::{
CompressionEncoding, EnabledCompressionEncodings, SingleMessageCompressionOverride,
};
use crate::{
body::BoxBody,
codec::{encode_server, Codec, Streaming},
@@ -9,6 +13,15 @@ use futures_util::{future, stream, TryStreamExt};
use http_body::Body;
use std::fmt;
macro_rules! t {
($result:expr) => {
match $result {
Ok(value) => value,
Err(status) => return status.to_http(),
}
};
}
/// A gRPC Server handler.
///
/// This will wrap some inner [`Codec`] and provide utilities to handle
@@ -20,6 +33,12 @@ use std::fmt;
/// implements some [`Body`].
pub struct Grpc<T> {
codec: T,
/// Which compression encodings does the server accept for requests?
#[cfg(feature = "compression")]
accept_compression_encodings: EnabledCompressionEncodings,
/// Which compression encodings might the server use for responses.
#[cfg(feature = "compression")]
send_compression_encodings: EnabledCompressionEncodings,
}
impl<T> Grpc<T>
@@ -29,7 +48,121 @@ where
{
/// Creates a new gRPC server with the provided [`Codec`].
pub fn new(codec: T) -> Self {
Self { codec }
Self {
codec,
#[cfg(feature = "compression")]
accept_compression_encodings: EnabledCompressionEncodings::default(),
#[cfg(feature = "compression")]
send_compression_encodings: EnabledCompressionEncodings::default(),
}
}
/// Enable accepting `gzip` compressed requests.
///
/// If a request with an unsupported encoding is received the server will respond with
/// [`Code::UnUnimplemented`](crate::Code).
///
/// # Example
///
/// The most common way of using this is through a server generated by tonic-build:
///
/// ```rust
/// # struct Svc;
/// # struct ExampleServer<T>(T);
/// # impl<T> ExampleServer<T> {
/// # fn new(svc: T) -> Self { Self(svc) }
/// # fn accept_gzip(self) -> Self { self }
/// # }
/// # #[tonic::async_trait]
/// # trait Example {}
///
/// #[tonic::async_trait]
/// impl Example for Svc {
/// // ...
/// }
///
/// let service = ExampleServer::new(Svc).accept_gzip();
/// ```
#[cfg(feature = "compression")]
#[cfg_attr(docsrs, doc(cfg(feature = "compression")))]
pub fn accept_gzip(mut self) -> Self {
self.accept_compression_encodings.enable_gzip();
self
}
#[doc(hidden)]
#[cfg(not(feature = "compression"))]
pub fn accept_gzip(self) -> Self {
panic!("`accept_gzip` called on a server but the `compression` feature is not enabled on tonic");
}
/// Enable sending `gzip` compressed responses.
///
/// Requires the client to also support receiving compressed responses.
///
/// # Example
///
/// The most common way of using this is through a server generated by tonic-build:
///
/// ```rust
/// # struct Svc;
/// # struct ExampleServer<T>(T);
/// # impl<T> ExampleServer<T> {
/// # fn new(svc: T) -> Self { Self(svc) }
/// # fn send_gzip(self) -> Self { self }
/// # }
/// # #[tonic::async_trait]
/// # trait Example {}
///
/// #[tonic::async_trait]
/// impl Example for Svc {
/// // ...
/// }
///
/// let service = ExampleServer::new(Svc).send_gzip();
/// ```
#[cfg(feature = "compression")]
#[cfg_attr(docsrs, doc(cfg(feature = "compression")))]
pub fn send_gzip(mut self) -> Self {
self.send_compression_encodings.enable_gzip();
self
}
#[doc(hidden)]
#[cfg(not(feature = "compression"))]
pub fn send_gzip(self) -> Self {
panic!(
"`send_gzip` called on a server but the `compression` feature is not enabled on tonic"
);
}
#[cfg(feature = "compression")]
#[doc(hidden)]
pub fn apply_compression_config(
self,
accept_encodings: EnabledCompressionEncodings,
send_encodings: EnabledCompressionEncodings,
) -> Self {
let mut this = self;
let EnabledCompressionEncodings { gzip: accept_gzip } = accept_encodings;
if accept_gzip {
this = this.accept_gzip();
}
let EnabledCompressionEncodings { gzip: send_gzip } = send_encodings;
if send_gzip {
this = this.send_gzip();
}
this
}
#[cfg(not(feature = "compression"))]
#[doc(hidden)]
#[allow(unused_variables)]
pub fn apply_compression_config(self, accept_encodings: (), send_encodings: ()) -> Self {
self
}
/// Handle a single unary gRPC request.
@@ -43,13 +176,23 @@ where
B: Body + Send + Sync + 'static,
B::Error: Into<crate::Error> + Send,
{
#[cfg(feature = "compression")]
let accept_encoding = CompressionEncoding::from_accept_encoding_header(
req.headers(),
self.send_compression_encodings,
);
let request = match self.map_request_unary(req).await {
Ok(r) => r,
Err(status) => {
return self
.map_response::<stream::Once<future::Ready<Result<T::Encode, Status>>>>(Err(
status,
));
.map_response::<stream::Once<future::Ready<Result<T::Encode, Status>>>>(
Err(status),
#[cfg(feature = "compression")]
accept_encoding,
#[cfg(feature = "compression")]
SingleMessageCompressionOverride::default(),
);
}
};
@@ -58,7 +201,16 @@ where
.await
.map(|r| r.map(|m| stream::once(future::ok(m))));
self.map_response(response)
#[cfg(feature = "compression")]
let compression_override = compression_override_from_response(&response);
self.map_response(
response,
#[cfg(feature = "compression")]
accept_encoding,
#[cfg(feature = "compression")]
compression_override,
)
}
/// Handle a server side streaming request.
@@ -73,16 +225,36 @@ where
B: Body + Send + Sync + 'static,
B::Error: Into<crate::Error> + Send,
{
#[cfg(feature = "compression")]
let accept_encoding = CompressionEncoding::from_accept_encoding_header(
req.headers(),
self.send_compression_encodings,
);
let request = match self.map_request_unary(req).await {
Ok(r) => r,
Err(status) => {
return self.map_response::<S::ResponseStream>(Err(status));
return self.map_response::<S::ResponseStream>(
Err(status),
#[cfg(feature = "compression")]
accept_encoding,
#[cfg(feature = "compression")]
SingleMessageCompressionOverride::default(),
);
}
};
let response = service.call(request).await;
self.map_response(response)
self.map_response(
response,
#[cfg(feature = "compression")]
accept_encoding,
// disabling compression of individual stream items must be done on
// the items themselves
#[cfg(feature = "compression")]
SingleMessageCompressionOverride::default(),
)
}
/// Handle a client side streaming gRPC request.
@@ -96,12 +268,29 @@ where
B: Body + Send + Sync + 'static,
B::Error: Into<crate::Error> + Send + 'static,
{
let request = self.map_request_streaming(req);
#[cfg(feature = "compression")]
let accept_encoding = CompressionEncoding::from_accept_encoding_header(
req.headers(),
self.send_compression_encodings,
);
let request = t!(self.map_request_streaming(req));
let response = service
.call(request)
.await
.map(|r| r.map(|m| stream::once(future::ok(m))));
self.map_response(response)
#[cfg(feature = "compression")]
let compression_override = compression_override_from_response(&response);
self.map_response(
response,
#[cfg(feature = "compression")]
accept_encoding,
#[cfg(feature = "compression")]
compression_override,
)
}
/// Handle a bi-directional streaming gRPC request.
@@ -116,9 +305,23 @@ where
B: Body + Send + Sync + 'static,
B::Error: Into<crate::Error> + Send,
{
let request = self.map_request_streaming(req);
#[cfg(feature = "compression")]
let accept_encoding = CompressionEncoding::from_accept_encoding_header(
req.headers(),
self.send_compression_encodings,
);
let request = t!(self.map_request_streaming(req));
let response = service.call(request).await;
self.map_response(response)
self.map_response(
response,
#[cfg(feature = "compression")]
accept_encoding,
#[cfg(feature = "compression")]
SingleMessageCompressionOverride::default(),
)
}
async fn map_request_unary<B>(
@@ -129,7 +332,16 @@ where
B: Body + Send + Sync + 'static,
B::Error: Into<crate::Error> + Send,
{
#[cfg(feature = "compression")]
let request_compression_encoding = self.request_encoding_if_supported(&request)?;
let (parts, body) = request.into_parts();
#[cfg(feature = "compression")]
let stream =
Streaming::new_request(self.codec.decoder(), body, request_compression_encoding);
#[cfg(not(feature = "compression"))]
let stream = Streaming::new_request(self.codec.decoder(), body);
futures_util::pin_mut!(stream);
@@ -151,42 +363,112 @@ where
fn map_request_streaming<B>(
&mut self,
request: http::Request<B>,
) -> Request<Streaming<T::Decode>>
) -> Result<Request<Streaming<T::Decode>>, Status>
where
B: Body + Send + Sync + 'static,
B::Error: Into<crate::Error> + Send,
{
Request::from_http(request.map(|body| Streaming::new_request(self.codec.decoder(), body)))
#[cfg(feature = "compression")]
let encoding = self.request_encoding_if_supported(&request)?;
#[cfg(feature = "compression")]
let request =
request.map(|body| Streaming::new_request(self.codec.decoder(), body, encoding));
#[cfg(not(feature = "compression"))]
let request = request.map(|body| Streaming::new_request(self.codec.decoder(), body));
Ok(Request::from_http(request))
}
fn map_response<B>(
&mut self,
response: Result<crate::Response<B>, Status>,
#[cfg(feature = "compression")] accept_encoding: Option<CompressionEncoding>,
#[cfg(feature = "compression")] compression_override: SingleMessageCompressionOverride,
) -> http::Response<BoxBody>
where
B: TryStream<Ok = T::Encode, Error = Status> + Send + Sync + 'static,
{
match response {
Ok(r) => {
let (mut parts, body) = r.into_http().into_parts();
let response = match response {
Ok(r) => r,
Err(status) => return status.to_http(),
};
// Set the content type
parts.headers.insert(
http::header::CONTENT_TYPE,
http::header::HeaderValue::from_static("application/grpc"),
);
let (mut parts, body) = response.into_http().into_parts();
let body = encode_server(self.codec.encoder(), body.into_stream());
// Set the content type
parts.headers.insert(
http::header::CONTENT_TYPE,
http::header::HeaderValue::from_static("application/grpc"),
);
http::Response::from_parts(parts, BoxBody::new(body))
}
Err(status) => status.to_http(),
#[cfg(feature = "compression")]
if let Some(encoding) = accept_encoding {
// Set the content encoding
parts.headers.insert(
crate::codec::compression::ENCODING_HEADER,
encoding.into_header_value(),
);
}
let body = encode_server(
self.codec.encoder(),
body.into_stream(),
#[cfg(feature = "compression")]
accept_encoding,
#[cfg(feature = "compression")]
compression_override,
);
http::Response::from_parts(parts, BoxBody::new(body))
}
#[cfg(feature = "compression")]
fn request_encoding_if_supported<B>(
&self,
request: &http::Request<B>,
) -> Result<Option<CompressionEncoding>, Status> {
CompressionEncoding::from_encoding_header(
request.headers(),
self.accept_compression_encodings,
)
}
}
impl<T: fmt::Debug> fmt::Debug for Grpc<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Grpc").field("codec", &self.codec).finish()
let mut f = f.debug_struct("Grpc");
f.field("codec", &self.codec);
#[cfg(feature = "compression")]
f.field(
"accept_compression_encodings",
&self.accept_compression_encodings,
);
#[cfg(feature = "compression")]
f.field(
"send_compression_encodings",
&self.send_compression_encodings,
);
f.finish()
}
}
#[cfg(feature = "compression")]
fn compression_override_from_response<B, E>(
res: &Result<crate::Response<B>, E>,
) -> SingleMessageCompressionOverride {
res.as_ref()
.ok()
.and_then(|response| {
response
.extensions()
.get::<SingleMessageCompressionOverride>()
.copied()
})
.unwrap_or_default()
}