Skip to content

Commit 1b3fa89

Browse files
authored
Make compression encoding configuration more malleable (#1757)
1 parent d312dcc commit 1b3fa89

File tree

1 file changed

+43
-43
lines changed

1 file changed

+43
-43
lines changed

tonic/src/codec/compression.rs

Lines changed: 43 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use crate::{metadata::MetadataValue, Status};
2-
use bytes::{Buf, BytesMut};
2+
use bytes::{Buf, BufMut, BytesMut};
33
#[cfg(feature = "gzip")]
44
use flate2::read::{GzDecoder, GzEncoder};
55
use std::fmt;
@@ -10,62 +10,64 @@ pub(crate) const ENCODING_HEADER: &str = "grpc-encoding";
1010
pub(crate) const ACCEPT_ENCODING_HEADER: &str = "grpc-accept-encoding";
1111

1212
/// Struct used to configure which encodings are enabled on a server or channel.
13+
///
14+
/// Represents an ordered list of compression encodings that are enabled.
1315
#[derive(Debug, Default, Clone, Copy)]
1416
pub struct EnabledCompressionEncodings {
15-
#[cfg(feature = "gzip")]
16-
pub(crate) gzip: bool,
17-
#[cfg(feature = "zstd")]
18-
pub(crate) zstd: bool,
17+
inner: [Option<CompressionEncoding>; 2],
1918
}
2019

2120
impl EnabledCompressionEncodings {
22-
/// Check if a [`CompressionEncoding`] is enabled.
23-
pub fn is_enabled(&self, encoding: CompressionEncoding) -> bool {
24-
match encoding {
25-
#[cfg(feature = "gzip")]
26-
CompressionEncoding::Gzip => self.gzip,
27-
#[cfg(feature = "zstd")]
28-
CompressionEncoding::Zstd => self.zstd,
29-
}
30-
}
31-
3221
/// Enable a [`CompressionEncoding`].
22+
///
23+
/// Adds the new encoding to the end of the encoding list.
3324
pub fn enable(&mut self, encoding: CompressionEncoding) {
34-
match encoding {
35-
#[cfg(feature = "gzip")]
36-
CompressionEncoding::Gzip => self.gzip = true,
37-
#[cfg(feature = "zstd")]
38-
CompressionEncoding::Zstd => self.zstd = true,
25+
for e in self.inner.iter_mut() {
26+
match e {
27+
Some(e) if *e == encoding => return,
28+
None => {
29+
*e = Some(encoding);
30+
return;
31+
}
32+
_ => continue,
33+
}
3934
}
4035
}
4136

37+
/// Remove the last [`CompressionEncoding`].
38+
pub fn pop(&mut self) -> Option<CompressionEncoding> {
39+
self.inner
40+
.iter_mut()
41+
.rev()
42+
.find(|entry| entry.is_some())?
43+
.take()
44+
}
45+
4246
pub(crate) fn into_accept_encoding_header_value(self) -> Option<http::HeaderValue> {
43-
match (self.is_gzip_enabled(), self.is_zstd_enabled()) {
44-
(true, false) => Some(http::HeaderValue::from_static("gzip,identity")),
45-
(false, true) => Some(http::HeaderValue::from_static("zstd,identity")),
46-
(true, true) => Some(http::HeaderValue::from_static("gzip,zstd,identity")),
47-
(false, false) => None,
47+
let mut value = BytesMut::new();
48+
for encoding in self.inner.into_iter().flatten() {
49+
if !value.is_empty() {
50+
value.put_slice(b",");
51+
}
52+
value.put_slice(encoding.as_str().as_bytes());
4853
}
49-
}
5054

51-
#[cfg(feature = "gzip")]
52-
const fn is_gzip_enabled(&self) -> bool {
53-
self.gzip
54-
}
55+
if value.is_empty() {
56+
return None;
57+
}
5558

56-
#[cfg(not(feature = "gzip"))]
57-
const fn is_gzip_enabled(&self) -> bool {
58-
false
59+
value.put_slice(b",identity");
60+
Some(http::HeaderValue::from_maybe_shared(value).unwrap())
5961
}
6062

61-
#[cfg(feature = "zstd")]
62-
const fn is_zstd_enabled(&self) -> bool {
63-
self.zstd
63+
/// Check if a [`CompressionEncoding`] is enabled.
64+
pub fn is_enabled(&self, encoding: CompressionEncoding) -> bool {
65+
self.inner.iter().any(|e| *e == Some(encoding))
6466
}
6567

66-
#[cfg(not(feature = "zstd"))]
67-
const fn is_zstd_enabled(&self) -> bool {
68-
false
68+
/// Check if any [`CompressionEncoding`]s are enabled.
69+
pub fn is_empty(&self) -> bool {
70+
self.inner.iter().all(|e| e.is_none())
6971
}
7072
}
7173

@@ -95,7 +97,7 @@ impl CompressionEncoding {
9597
map: &http::HeaderMap,
9698
enabled_encodings: EnabledCompressionEncodings,
9799
) -> Option<Self> {
98-
if !enabled_encodings.is_gzip_enabled() && !enabled_encodings.is_zstd_enabled() {
100+
if enabled_encodings.is_empty() {
99101
return None;
100102
}
101103

@@ -157,9 +159,7 @@ impl CompressionEncoding {
157159
}
158160
}
159161

160-
#[allow(missing_docs)]
161-
#[cfg(any(feature = "gzip", feature = "zstd"))]
162-
pub(crate) fn as_str(&self) -> &'static str {
162+
pub(crate) fn as_str(self) -> &'static str {
163163
match self {
164164
#[cfg(feature = "gzip")]
165165
CompressionEncoding::Gzip => "gzip",

0 commit comments

Comments
 (0)