Merge pull request #1003 from finitecomputer/claude/blob-transfers-at-once

sync: blobs move four at a time, up and down

Author: futurepaul

Cargo.lock

@@ -24,6 +24,9 @@ serde_json = "1"
 # TLS trusts the public roots and the machine's own: a company's CA (its
 # intranet's, or its TLS-inspecting proxy's), installed in the OS's store
 reqwest = { version = "0.12", default-features = false, features = ["blocking", "rustls-tls", "rustls-tls-native-roots"] }
+# a request's body kept once for its retries, an answer's taken in place
+# (api.rs `attempt_body`, `owned`): reqwest's own buffer type
+bytes = "1"
 sha2 = "0.10"
 base64 = "0.22"
 hex = "0.4"

cli/Cargo.toml

@@ -3,6 +3,7 @@
 // the platform's `ErrorBody` and a success into its fragment_proto type.
 use crate::auth::Identity;
 use anyhow::{Context, Result};
+use bytes::Bytes;
 use fragment_proto::{ErrorBody, ErrorCode};
 use serde::de::DeserializeOwned;
 use serde_json::Value;
@@ -403,7 +404,7 @@ impl Client {
     }
 
     fn request(&self, method: &str, path: &str, body: Option<Vec<u8>>) -> Result<Resp> {
-        let body = body.unwrap_or_default();
+        let body = Bytes::from(body.unwrap_or_default());
         let timeout = timeout_for(body.len() as u64);
         self.send(method, path, body, Replay::of(method), Signed::Body, timeout, None)
     }
@@ -412,7 +413,7 @@ impl Client {
     /// allows (`send_retrying`); `content_type` names the body's type, when
     /// it has one.
     #[allow(clippy::too_many_arguments)]
-    fn send(&self, method: &str, path: &str, body: Vec<u8>, replay: Replay, signed: Signed, timeout: Duration, content_type: Option<&str>) -> Result<Resp> {
+    fn send(&self, method: &str, path: &str, body: Bytes, replay: Replay, signed: Signed, timeout: Duration, content_type: Option<&str>) -> Result<Resp> {
         let path = &self.resolved(path)?;
         let url = self.url(method, path);
         let verb: reqwest::Method = method.parse()?;
@@ -423,7 +424,7 @@ impl Client {
             if let Some(t) = content_type {
                 req = req.header("content-type", t);
             }
-            if body.is_empty() { req } else { req.body(body.clone()) }
+            if body.is_empty() { req } else { req.body(attempt_body(&body)) }
         };
         let (host, coded) = (&self.host, |code, msg| Err(anyhow::Error::new(CodedError { code, msg })));
         match send_retrying(build, replay, self.verbose.then_some(what.as_str())) {
@@ -450,14 +451,14 @@ impl Client {
     /// A bounded slow operation, such as icon generation. It retains the
     /// normal non-idempotent POST retry policy.
     pub fn post_json_with_timeout(&self, path: &str, v: &impl serde::Serialize, timeout: Duration) -> Result<Resp> {
-        self.send("POST", path, serde_json::to_vec(v)?, Replay::ConnectOnly, Signed::Body, timeout, None)
+        self.send("POST", path, Bytes::from(serde_json::to_vec(v)?), Replay::ConnectOnly, Signed::Body, timeout, None)
     }
     /// A POST whose body carries its idempotency id (an operation call):
     /// retried like a read (`Replay::ById`).
     pub fn post_json_by_id(&self, path: &str, v: &impl serde::Serialize) -> Result<Resp> {
         let body = serde_json::to_vec(v)?;
         let timeout = timeout_for(body.len() as u64);
-        self.send("POST", path, body, Replay::ById, Signed::Body, timeout, None)
+        self.send("POST", path, Bytes::from(body), Replay::ById, Signed::Body, timeout, None)
     }
     pub fn put_json(&self, path: &str, v: &impl serde::Serialize) -> Result<Resp> {
         self.request("PUT", path, Some(serde_json::to_vec(v)?))
@@ -470,12 +471,12 @@ impl Client {
     /// `content_type` is what a page reading it gets (`__blob`).
     pub fn put_blob(&self, path: &str, bytes: Vec<u8>, content_type: Option<&str>) -> Result<Resp> {
         let timeout = timeout_for(bytes.len() as u64);
-        self.send("PUT", path, bytes, Replay::Safe, Signed::Url, timeout, content_type)
+        self.send("PUT", path, Bytes::from(bytes), Replay::Safe, Signed::Url, timeout, content_type)
     }
     /// A GET whose answer is known to be about `bytes` long (a blob), given
     /// the time that takes.
     pub fn get_sized(&self, path: &str, bytes: u64) -> Result<Resp> {
-        self.send("GET", path, Vec::new(), Replay::Safe, Signed::Body, timeout_for(bytes), None)
+        self.send("GET", path, Bytes::new(), Replay::Safe, Signed::Body, timeout_for(bytes), None)
     }
     /// A signed GET whose answer goes to `to` as it arrives (an export's
     /// archive, which may be gigabytes), never held whole: its bytes. A
@@ -489,7 +490,7 @@ impl Client {
         let mut resp = self.http.get(&url).header(header, value).timeout(Duration::from_secs(6 * 3600)).send().with_context(|| format!("GET {path}"))?;
         let status = resp.status().as_u16();
         if !(200..300).contains(&status) {
-            let body = resp.bytes().map(|b| b.to_vec()).unwrap_or_default();
+            let body = resp.bytes().map(owned).unwrap_or_default();
             return Err(anyhow::Error::new(Resp { status, body }.refusal()));
         }
         let mut file = std::fs::File::create(to).with_context(|| format!("writing {}", to.display()))?;
@@ -550,6 +551,21 @@ pub fn resolve_label(label: &str, listed: &[fragment_proto::ListedFragment]) ->
     }
 }
 
+/// The body one attempt of a request sends: the buffer kept for its
+/// retries itself, shared (a reference, not a copy). A blob's PUT (up to
+/// `BLOB_MAX_BYTES`) once held two whole copies, the kept one and the
+/// attempt's.
+pub fn attempt_body(kept: &Bytes) -> Bytes {
+    kept.clone()
+}
+
+/// An answer's bytes as the `Vec` a `Resp` holds: the buffer reqwest read
+/// them into, taken in place when the answer holds it alone (a body read
+/// whole does), not copied out beside it as a blob's GET once was.
+pub fn owned(answer: Bytes) -> Vec<u8> {
+    Vec::from(answer)
+}
+
 pub fn send_retrying(build: impl Fn() -> reqwest::blocking::RequestBuilder, replay: Replay, log: Option<&str>) -> std::result::Result<(u16, Vec<u8>), Failed> {
     let (mut last, mut reached) = (None, false);
     for attempt in 0..REQUEST_ATTEMPTS {
@@ -562,7 +578,7 @@ pub fn send_retrying(build: impl Fn() -> reqwest::blocking::RequestBuilder, repl
         let answer = build().send().and_then(|resp| {
             crate::notice::heard(resp.url().host_str(), resp.headers().get("x-fragment-cli").and_then(|v| v.to_str().ok()));
             let status = resp.status().as_u16();
-            resp.bytes().map(|b| (status, b.to_vec()))
+            resp.bytes().map(|b| (status, owned(b)))
         });
         let ms = t0.elapsed().as_millis();
         match answer {
@@ -602,6 +618,38 @@ pub fn encode_q(s: &str) -> String {
 mod tests {
     use super::*;
 
+    /// Goal: a transfer's bytes are one buffer, not two (Codex on
+    /// 69037874: four 256 MiB blobs moving at once held 2 GiB). Method:
+    /// the buffers themselves, by address. An attempt sends the body kept
+    /// for its retries, shared; an answer read whole by this reqwest (a
+    /// 4 MiB one from a local socket, in many chunks) becomes the `Vec` a
+    /// `Resp` holds in place. Each fails with a copy (`clone` of a `Vec`,
+    /// `to_vec`: checked).
+    #[test]
+    fn a_transfer_holds_its_bytes_once() {
+        let kept = Bytes::from(vec![7u8; 1 << 20]);
+        assert_eq!(attempt_body(&kept).as_ptr(), kept.as_ptr(), "an attempt's body is the kept buffer");
+
+        use std::io::{Read, Write};
+        const LEN: usize = 4 << 20;
+        let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
+        let url = format!("http://{}/blob", listener.local_addr().unwrap());
+        let served = std::thread::spawn(move || {
+            let (mut conn, _) = listener.accept().unwrap();
+            let mut asked = [0u8; 1024];
+            let _ = conn.read(&mut asked).unwrap();
+            write!(conn, "HTTP/1.1 200 OK\r\ncontent-length: {LEN}\r\nconnection: close\r\n\r\n").unwrap();
+            for _ in 0..LEN / (64 << 10) {
+                conn.write_all(&[9u8; 64 << 10]).unwrap();
+            }
+        });
+        let answer = reqwest::blocking::get(&url).unwrap().bytes().unwrap();
+        let at = answer.as_ptr();
+        let body = owned(answer);
+        assert_eq!((body.as_ptr(), body.len()), (at, LEN), "an answer's Vec is the buffer reqwest read it into");
+        served.join().unwrap();
+    }
+
     use fragment_fakes::http::{Response as FakeResponse, Server};
     use std::sync::atomic::{AtomicU32, Ordering};
     use std::sync::Arc;

cli/src/api.rs

@@ -4,7 +4,7 @@
 //! downloads a pointer's bytes when it pulls one, so a folder always holds
 //! real files, at the threshold the host names (`blobMinBytes` in status).
 
-use std::sync::OnceLock;
+use std::sync::Mutex;
 
 use anyhow::{anyhow, bail, Result};
 use fragment_core::blob::{self, sha256_hex};
@@ -15,23 +15,26 @@ use crate::api::Client;
 pub struct Blobs<'a> {
     client: &'a Client,
     name: &'a str,
-    /// shared by a pass's reads at once (`sync::READS_AT_ONCE`)
-    min: OnceLock<Option<u64>>,
+    /// the host's threshold once asked: a pass's transfers at once
+    /// (`sync::BLOB_TRANSFERS_AT_ONCE`) share it, and ask it once
+    min: Mutex<Option<Option<u64>>>,
 }
 
 impl<'a> Blobs<'a> {
     pub fn new(client: &'a Client, name: &'a str) -> Blobs<'a> {
-        Blobs { client, name, min: OnceLock::new() }
+        Blobs { client, name, min: Mutex::new(None) }
     }
 
     /// The host's blob threshold, asked once per pass (`None`: it keeps
     /// everything in git).
     fn min_bytes(&self) -> Result<Option<u64>> {
-        if let Some(m) = self.min.get() {
-            return Ok(*m);
+        let mut min = self.min.lock().unwrap_or_else(|e| e.into_inner());
+        if let Some(m) = *min {
+            return Ok(m);
         }
         let status: FragmentStatus = self.client.call_as(self.client.get(&format!("/api/f/{}/status", self.name))?)?;
-        Ok(*self.min.get_or_init(|| status.blob_min_bytes))
+        *min = Some(status.blob_min_bytes);
+        Ok(status.blob_min_bytes)
     }
 
     fn path(&self, sha: &str) -> String {

cli/src/blobs.rs

@@ -207,15 +207,16 @@ impl CodeStorage {
     /// semantics belong to the callers' bounded retry loops.
     fn req(&self, method: &str, path: &str, body: Option<Vec<u8>>, content_type: Option<&str>) -> Result<(u16, Vec<u8>), CsError> {
         let url = format!("{}/api/repos/{}{}", self.server, self.repo_seg(), path);
-        let body = body.unwrap_or_default();
+        // kept once for its retries, each attempt sending it shared
+        let body = bytes::Bytes::from(body.unwrap_or_default());
         let timeout = timeout_for(body.len() as u64);
         let verb: reqwest::Method = method.parse().map_err(|_| CsError::Malformed(format!("bad method {method}")))?;
         let build = || {
             let mut r = self.http.request(verb.clone(), &url).bearer_auth(&self.token).timeout(timeout);
             if let Some(ct) = content_type {
                 r = r.header("content-type", ct);
             }
-            if body.is_empty() { r } else { r.body(body.clone()) }
+            if body.is_empty() { r } else { r.body(crate::api::attempt_body(&body)) }
         };
         send_retrying(build, Replay::of(method), None).map_err(|failed| match failed {
             Failed::Unbuilt(e) => CsError::Malformed(format!("building {method} {path}: {e}")),

cli/src/codestorage.rs

@@ -2201,7 +2201,7 @@ fn run(cli: Cli) -> Result<()> {
                 req = req.header(api::AGENT_HEADER, &mode.agent);
             }
             let resp = req.send()?;
-            let resp = api::Resp { status: resp.status().as_u16(), body: resp.bytes()?.to_vec() };
+            let resp = api::Resp { status: resp.status().as_u16(), body: api::owned(resp.bytes()?) };
             let v = c.call(resp)?;
             json_exit(j, &v);
             println!("{}", serde_json::to_string(&v)?);