Skip to content

perf(fetch): share retained response body backing #393

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ldm0 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(fetch): share retained response body backing #393

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion moli-fetch/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ pub use request::{
pub use request_policy::{is_bad_port, should_request_be_blocked_due_to_bad_port};
pub use response::{
NegotiatedHttpVersion, NetworkRequestExtraInfo, NetworkResponseExtraInfo, RawResponse,
RedirectInfo, Response, ResponseBody, ResponseHead,
RedirectInfo, Response, ResponseBody, ResponseHead, SharedResponseBodyBytes,
};
pub use runtime::PendingStreamingRawResponse;
pub use runtime::{
Expand Down
219 changes: 180 additions & 39 deletions moli-fetch/src/response.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
use anyhow::{Result, anyhow};
use moli_cookie_jar::{StoredCookieQueryReport, StoredCookieSetReport};
use std::sync::Arc;
use url::Url;

use crate::{StreamingHtmlResponse, StreamingRawResponse};
Expand Down Expand Up @@ -50,21 +51,100 @@ pub struct ResponseHead {
pub negotiated_http_version: Option<NegotiatedHttpVersion>,
}

#[derive(Clone, Debug)]
pub struct SharedResponseBodyBytes {
backing: SharedResponseBodyBytesBacking,
}

#[derive(Clone, Debug)]
enum SharedResponseBodyBytesBacking {
Text(Arc<String>),
Bytes(Arc<Vec<u8>>),
}

impl SharedResponseBodyBytes {
pub fn from_bytes(bytes: Vec<u8>) -> Self {
Self {
backing: SharedResponseBodyBytesBacking::Bytes(Arc::new(bytes)),
}
}

fn from_text(text: Arc<String>) -> Self {
Self {
backing: SharedResponseBodyBytesBacking::Text(text),
}
}

fn from_shared_bytes(bytes: Arc<Vec<u8>>) -> Self {
Self {
backing: SharedResponseBodyBytesBacking::Bytes(bytes),
}
}

pub fn as_slice(&self) -> &[u8] {
match &self.backing {
SharedResponseBodyBytesBacking::Text(text) => text.as_bytes(),
SharedResponseBodyBytesBacking::Bytes(bytes) => bytes,
}
}

pub fn capacity(&self) -> usize {
match &self.backing {
SharedResponseBodyBytesBacking::Text(text) => text.capacity(),
SharedResponseBodyBytesBacking::Bytes(bytes) => bytes.capacity(),
}
}
}

fn unwrap_shared_string(value: Arc<String>) -> String {
Arc::try_unwrap(value).unwrap_or_else(|shared| (*shared).clone())
}

fn unwrap_shared_bytes(value: Arc<Vec<u8>>) -> Vec<u8> {
Arc::try_unwrap(value).unwrap_or_else(|shared| (*shared).clone())
}

#[derive(Debug)]
pub enum ResponseBody {
MaterializedText { text: String, bytes: Vec<u8> },
MaterializedBytes(Vec<u8>),
MaterializedText {
text: Arc<String>,
exact_bytes: Option<Arc<Vec<u8>>>,
},
MaterializedBytes(Arc<Vec<u8>>),
StreamingText(Box<StreamingHtmlResponse>),
StreamingBytes(Box<StreamingRawResponse>),
}

impl ResponseBody {
pub fn materialized_text(text: String, bytes: Vec<u8>) -> Self {
Self::MaterializedText { text, bytes }
let exact_bytes = (!bytes.is_empty()).then(|| Arc::new(bytes));
Self::MaterializedText {
text: Arc::new(text),
exact_bytes,
}
}

/// Builds a text body from exact bytes without retaining a second copy
/// when the bytes are already valid UTF-8.
pub fn lossy_text_from_bytes(bytes: Vec<u8>) -> Self {
match String::from_utf8(bytes) {
Ok(text) => Self::MaterializedText {
text: Arc::new(text),
exact_bytes: None,
},
Err(error) => {
let bytes = error.into_bytes();
let text = String::from_utf8_lossy(&bytes).into_owned();
Self::MaterializedText {
text: Arc::new(text),
exact_bytes: Some(Arc::new(bytes)),
}
}
}
}

pub fn materialized_bytes(bytes: Vec<u8>) -> Self {
Self::MaterializedBytes(bytes)
Self::MaterializedBytes(Arc::new(bytes))
}

pub fn is_streaming(&self) -> bool {
Expand All @@ -73,38 +153,33 @@ impl ResponseBody {

pub fn try_into_materialized_bytes(self) -> std::result::Result<Vec<u8>, Self> {
match self {
Self::MaterializedBytes(bytes) => Ok(bytes),
Self::MaterializedText { text, bytes } => {
if bytes.is_empty() && !text.is_empty() {
Ok(text.into_bytes())
} else {
Ok(bytes)
}
}
Self::MaterializedBytes(bytes) => Ok(unwrap_shared_bytes(bytes)),
Self::MaterializedText { text, exact_bytes } => Ok(exact_bytes.map_or_else(
|| unwrap_shared_string(text).into_bytes(),
unwrap_shared_bytes,
)),
Self::StreamingText(_) | Self::StreamingBytes(_) => Err(self),
}
}

pub fn as_materialized_bytes(&self) -> Option<&[u8]> {
match self {
Self::MaterializedBytes(bytes) => Some(bytes),
Self::MaterializedText { text, bytes } => {
if bytes.is_empty() && !text.is_empty() {
Some(text.as_bytes())
} else {
Some(bytes)
}
}
Self::MaterializedText { text, exact_bytes } => Some(
exact_bytes
.as_deref()
.map_or(text.as_bytes(), Vec::as_slice),
),
Self::StreamingText(_) | Self::StreamingBytes(_) => None,
}
}

pub fn clone_materialized(&self) -> Option<Self> {
match self {
Self::MaterializedBytes(bytes) => Some(Self::MaterializedBytes(bytes.clone())),
Self::MaterializedText { text, bytes } => Some(Self::MaterializedText {
Self::MaterializedText { text, exact_bytes } => Some(Self::MaterializedText {
text: text.clone(),
bytes: bytes.clone(),
exact_bytes: exact_bytes.clone(),
}),
Self::StreamingText(_) | Self::StreamingBytes(_) => None,
}
Expand All @@ -119,29 +194,43 @@ impl ResponseBody {

pub fn try_into_lossy_materialized_text(self) -> std::result::Result<(String, Vec<u8>), Self> {
match self {
Self::MaterializedText { text, bytes } => Ok((text, bytes)),
Self::MaterializedText { text, exact_bytes } => {
let bytes =
exact_bytes.map_or_else(|| text.as_bytes().to_vec(), unwrap_shared_bytes);
let text = unwrap_shared_string(text);
Ok((text, bytes))
}
Self::MaterializedBytes(bytes) => {
let bytes = unwrap_shared_bytes(bytes);
let text = String::from_utf8_lossy(&bytes).into_owned();
Ok((text, bytes))
}
Self::StreamingText(_) | Self::StreamingBytes(_) => Err(self),
}
}

/// Converts an already-materialized body to its text representation.
pub fn try_into_materialized_text_body(self) -> std::result::Result<Self, Self> {
match self {
Self::MaterializedText { .. } => Ok(self),
Self::MaterializedBytes(bytes) => {
Ok(Self::lossy_text_from_bytes(unwrap_shared_bytes(bytes)))
}
Self::StreamingText(_) | Self::StreamingBytes(_) => Err(self),
}
}

/// Drains any streaming source and returns a complete byte body.
///
/// This is an explicit compatibility boundary. Prefer chunked consumption
/// when the caller can avoid building a full response body in memory.
pub async fn into_materialized_bytes(self) -> Result<Vec<u8>> {
match self {
Self::MaterializedBytes(bytes) => Ok(bytes),
Self::MaterializedText { text, bytes } => {
if bytes.is_empty() && !text.is_empty() {
Ok(text.into_bytes())
} else {
Ok(bytes)
}
}
Self::MaterializedBytes(bytes) => Ok(unwrap_shared_bytes(bytes)),
Self::MaterializedText { text, exact_bytes } => Ok(exact_bytes.map_or_else(
|| unwrap_shared_string(text).into_bytes(),
unwrap_shared_bytes,
)),
Self::StreamingText(response) => {
let mut response = *response;
let mut bytes = Vec::new();
Expand All @@ -167,8 +256,14 @@ impl ResponseBody {
/// the exact bytes used to derive it.
pub async fn into_lossy_materialized_text(self) -> Result<(String, Vec<u8>)> {
match self {
Self::MaterializedText { text, bytes } => Ok((text, bytes)),
Self::MaterializedText { text, exact_bytes } => {
let bytes =
exact_bytes.map_or_else(|| text.as_bytes().to_vec(), unwrap_shared_bytes);
let text = unwrap_shared_string(text);
Ok((text, bytes))
}
Self::MaterializedBytes(bytes) => {
let bytes = unwrap_shared_bytes(bytes);
let text = String::from_utf8_lossy(&bytes).into_owned();
Ok((text, bytes))
}
Expand All @@ -191,6 +286,47 @@ impl ResponseBody {
}
}
}

/// Drains a body source and returns a materialized text body.
pub async fn into_materialized_text_body(self) -> Result<Self> {
match self {
Self::MaterializedText { .. } => Ok(self),
Self::MaterializedBytes(bytes) => {
Ok(Self::lossy_text_from_bytes(unwrap_shared_bytes(bytes)))
}
Self::StreamingText(response) => {
let mut response = *response;
let mut text = String::new();
while let Some(chunk) = response.next_chunk().await {
text.push_str(&chunk);
}
response.finish().await?;
Ok(Self::MaterializedText {
text: Arc::new(text),
exact_bytes: None,
})
}
Self::StreamingBytes(response) => {
let bytes = Self::StreamingBytes(response)
.into_materialized_bytes()
.await?;
Ok(Self::lossy_text_from_bytes(bytes))
}
}
}

pub fn shared_materialized_bytes(&self) -> Option<SharedResponseBodyBytes> {
match self {
Self::MaterializedText { text, exact_bytes } => Some(exact_bytes.as_ref().map_or_else(
|| SharedResponseBodyBytes::from_text(Arc::clone(text)),
|bytes| SharedResponseBodyBytes::from_shared_bytes(Arc::clone(bytes)),
)),
Self::MaterializedBytes(bytes) => Some(SharedResponseBodyBytes::from_shared_bytes(
Arc::clone(bytes),
)),
Self::StreamingText(_) | Self::StreamingBytes(_) => None,
}
}
}

#[derive(Debug)]
Expand Down Expand Up @@ -260,6 +396,12 @@ impl Response {
self.body_bytes().to_vec()
}

pub fn shared_body_bytes(&self) -> SharedResponseBodyBytes {
self.body
.shared_materialized_bytes()
.expect("Response body should remain materialized")
}

pub fn materialized_body(&self) -> ResponseBody {
self.body
.clone_materialized()
Expand Down Expand Up @@ -289,26 +431,25 @@ impl Response {
}

pub fn from_head_and_text_body(head: ResponseHead, body: String) -> Self {
let body_bytes = body.as_bytes().to_vec();
Self::from_head_and_body(head, body, body_bytes)
Self::from_head_and_body(head, body, Vec::new())
}

/// Builds a materialized text response from exact bytes by deriving the
/// compatibility text view with UTF-8 replacement semantics.
pub fn from_head_and_lossy_body_bytes(head: ResponseHead, body_bytes: Vec<u8>) -> Self {
let body = String::from_utf8_lossy(&body_bytes).into_owned();
Self::from_head_and_body(head, body, body_bytes)
Self::from_head_and_materialized_body(head, ResponseBody::lossy_text_from_bytes(body_bytes))
.expect("materialized text body should build a Response")
}

pub fn from_head_and_materialized_body(head: ResponseHead, body: ResponseBody) -> Result<Self> {
let (body, body_bytes) = body.try_into_lossy_materialized_text().map_err(|_| {
let body = body.try_into_materialized_text_body().map_err(|_| {
anyhow!("cannot build materialized Response from streaming response body")
})?;
Ok(Self {
final_url: head.final_url,
status: head.status,
headers: head.headers,
body: ResponseBody::materialized_text(body, body_bytes),
body,
request_cookie_report: head.request_cookie_report,
cookie_set_reports: head.cookie_set_reports,
redirected: head.redirected,
Expand All @@ -320,12 +461,12 @@ impl Response {
}

pub async fn from_head_and_body_source(head: ResponseHead, body: ResponseBody) -> Result<Self> {
let (body, body_bytes) = body.into_lossy_materialized_text().await?;
let body = body.into_materialized_text_body().await?;
Ok(Self {
final_url: head.final_url,
status: head.status,
headers: head.headers,
body: ResponseBody::materialized_text(body, body_bytes),
body,
request_cookie_report: head.request_cookie_report,
cookie_set_reports: head.cookie_set_reports,
redirected: head.redirected,
Expand Down
Loading
Loading