Skip to content

Commit 38924fe

Browse files
committed
perf(client): stream proxy response body instead of buffering
1 parent 77e23d1 commit 38924fe

3 files changed

Lines changed: 22 additions & 18 deletions

File tree

src/client.rs

Lines changed: 12 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,9 @@
11
use arc_swap::ArcSwap;
22
use cached::proc_macro::cached;
33
use futures_lite::future::block_on;
4-
use futures_lite::{future::Boxed, FutureExt};
5-
use http_body_util::BodyExt;
4+
use futures_lite::{future::Boxed, FutureExt, StreamExt};
5+
use http_body_util::{BodyExt, StreamBody};
6+
use hyper::body::Frame;
67
use hyper::{header, Method, Request, Response};
78
use log::{error, info, trace, warn};
89
use percent_encoding::{percent_encode, CONTROLS};
@@ -18,7 +19,7 @@ use std::result::Result;
1819

1920
use crate::dbg_msg;
2021
use crate::oauth::{force_refresh_token, token_daemon, Oauth, OauthBackendImpl};
21-
use crate::server::RequestExt;
22+
use crate::server::{Body, RequestExt};
2223
use crate::utils::{format_url, Post};
2324

2425
const REDDIT_URL_BASE: &str = "https://oauth.reddit.com";
@@ -69,8 +70,6 @@ pub fn build_client() -> WreqClient {
6970
.expect("Should always be able to build a client")
7071
}
7172

72-
use crate::server::{full, Body};
73-
7473
/// Convert a wreq Response into a hyper Response<Body>.
7574
/// wreq and hyper now both use http v1.x, so header types are directly compatible.
7675
trait IntoHyperResponse {
@@ -93,15 +92,18 @@ impl IntoHyperResponse for wreq::Response {
9392
}
9493
}
9594

96-
let bytes = self.bytes().await.map_err(|e| e.to_string())?;
97-
// wreq already decompresses; remove the content-encoding header so
98-
// downstream code doesn't try to decompress again.
95+
// wreq already decompresses; drop content-length since the final
96+
// size is unknown without buffering, relying on chunked encoding.
9997
if let Some(h) = builder.headers_mut() {
10098
h.remove(header::CONTENT_ENCODING);
101-
h.insert(header::CONTENT_LENGTH, bytes.len().into());
99+
h.remove(header::CONTENT_LENGTH);
102100
}
103101

104-
builder.body(full(bytes)).map_err(|e| e.to_string())
102+
let stream = self
103+
.bytes_stream()
104+
.map(|r| r.map(Frame::data).map_err(|e| e.to_string()));
105+
106+
builder.body(BodyExt::boxed(StreamBody::new(stream))).map_err(|e| e.to_string())
105107
}
106108
}
107109

src/server.rs

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -29,17 +29,19 @@ use tokio::net::TcpListener;
2929

3030
use crate::dbg_msg;
3131

32-
/// The unified body type used for all responses.
33-
pub type Body = BoxBody<Bytes, Infallible>;
32+
/// Unified response body type. Error is `String` so proxy streams can
33+
/// propagate upstream failures; `full()`/`empty()` are backed by
34+
/// `Full`/`Empty<_, Infallible>` and can never produce one.
35+
pub type Body = BoxBody<Bytes, String>;
3436

3537
/// Create a response body from a string or bytes.
3638
pub fn full<T: Into<Bytes>>(chunk: T) -> Body {
37-
Full::new(chunk.into()).boxed()
39+
Full::new(chunk.into()).map_err(|e: Infallible| match e {}).boxed()
3840
}
3941

4042
/// Create an empty response body.
4143
pub fn empty() -> Body {
42-
Empty::<Bytes>::new().boxed()
44+
Empty::<Bytes>::new().map_err(|e: Infallible| match e {}).boxed()
4345
}
4446

4547
const BANNED_USER_AGENTS: &[&str] = &[
@@ -604,8 +606,8 @@ async fn compress_response(req_headers: &HeaderMap<header::HeaderValue>, res: &m
604606
let body_bytes: Vec<u8> = {
605607
// Swap body with empty, collect old body
606608
let old_body = std::mem::replace(res.body_mut(), empty());
607-
// Body error type is Infallible, so collect cannot fail
608-
old_body.collect().await.expect("Body<Infallible> collect cannot fail").to_bytes().to_vec()
609+
// full/empty bodies cannot produce a String error; only proxy stream bodies can.
610+
old_body.collect().await.expect("static response body collection cannot fail").to_bytes().to_vec()
609611
};
610612

611613
// Don't bother compressing tiny responses

src/settings.rs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ pub async fn set(req: Request<Body>) -> Result<Response<Body>, String> {
7171
.filter_map(|header| Cookie::parse(header.to_str().unwrap_or_default()).ok())
7272
.collect();
7373

74-
// Collect the body bytes (Body error is Infallible, so unwrap is safe)
74+
// Request body is pre-buffered as Full<Bytes>; the String error cannot fire.
7575
let body_bytes = body.collect().await.unwrap().to_bytes();
7676

7777
let form = url::form_urlencoded::parse(&body_bytes).collect::<HashMap<_, _>>();
@@ -263,7 +263,7 @@ pub async fn update(req: Request<Body>) -> Result<Response<Body>, String> {
263263
}
264264

265265
pub async fn encoded_restore(req: Request<Body>) -> Result<Response<Body>, String> {
266-
// Body error is Infallible, so unwrap is safe
266+
// Request body is pre-buffered as Full<Bytes>; the String error cannot fire.
267267
let body = req.into_body().collect().await.unwrap().to_bytes();
268268

269269
if body.len() > 1024 * 1024 {

0 commit comments

Comments
 (0)