Pre-allocate the buffer based on the expected Content-Length with the Rust HTTP client

changelog.d/19498.misc

@@ -0,0 +1 @@
+Pre-allocate the buffer based on the expected `Content-Length` with the Rust HTTP client.

rust/src/http_client.rs

@@ -16,6 +16,7 @@ use std::{collections::HashMap, future::Future, sync::OnceLock};
 
 use anyhow::Context;
 use futures::TryStreamExt;
+use headers::HeaderMapExt;
 use once_cell::sync::OnceCell;
 use pyo3::{create_exception, exceptions::PyException, prelude::*};
 use reqwest::RequestBuilder;
@@ -235,8 +236,47 @@ impl HttpClient {
 
             let status = response.status();
 
+            // Find the expected `Content-Length` so we can pre-allocate the buffer
+            // necessary to read the response. It's expected that not every request will
+            // have a `Content-Length` header.
+            //
+            // `response.content_length()` does exist but the "value does not directly
+            // represents the value of the `Content-Length` header, but rather the size
+            // of the response’s body"
+            // (https://docs.rs/reqwest/latest/reqwest/struct.Response.html#method.content_length)
+            // and we want to avoid reading the entire body at this point because we
+            // purposely stream it below until the `response_limit`.
+            let content_length = {
+                let content_length = response
+                    .headers()
+                    .typed_get::<headers::ContentLength>()
+                    // We need a `usize` for the `Vec::with_capacity(...)` usage below
+                    .and_then(|content_length| content_length.0.try_into().ok());
+
+                // Sanity check that the request isn't too large from the information
+                // they told us (may be inaccurate so we also check below as we actually
+                // read the bytes)
+                if let Some(content_length_bytes) = content_length {
+                    if content_length_bytes > response_limit {
+                        Err(anyhow::anyhow!(
+                            "Response size (defined by `Content-Length`) too large"
+                        ))?;
+                    }
+                }
+
+                content_length
+            };
+
+            // Stream the response to avoid allocating a giant object on the server
+            // above our expected `response_limit`.
             let mut stream = response.bytes_stream();
-            let mut buffer = Vec::new();
+            // Pre-allocate the buffer based on the expected `Content-Length`
+            let mut buffer = Vec::with_capacity(
+                content_length
+                    // Default to pre-allocating nothing when the request doesn't have a
+                    // `Content-Length` header
+                    .unwrap_or(0),
+            );
             while let Some(chunk) = stream.try_next().await.context("reading body")? {
                 if buffer.len() + chunk.len() > response_limit {
                     Err(anyhow::anyhow!("Response size too large"))?;