Do not load the entire artifact in memory when uploading (#618) #677

geofft · 2025-06-29T21:45:12Z

No description provided.

geofft · 2025-06-29T21:46:00Z

This compiles, but is wholly untested. I might try to actually run it tomorrow on a fork or something.

It would also be nice to be able to inject errors, somehow....

zanieb · 2025-06-30T20:48:02Z

@konstin has all the context on this pattern in uv, I'll delegate review :)

konstin · 2025-07-01T07:49:52Z

src/github.rs

+        let result = request.send().await.map_err(|e| e.into());
+
+        if retryable_strategy.handle(&result) == Some(Retryable::Transient) {
+            let retry_decision = retry_policy.should_retry(start_time, n_past_retries);


fwiw we have a more extensive retry policy in uv, but the reqwest-retry may be sufficient here.

konstin · 2025-07-01T07:51:32Z

src/github.rs

+            }
+        }
+        break result?;
+    };

    if !response.status().is_success() {


This means we don't retry on status errors, such as a 500 (only the 403 we handle in the github upload retry strategy)

This (and the above) is intended to match the existing logic using reqwest-retry... we ought to retry on a 500, though, if we don't let me fix that while we're here.

konstin · 2025-07-01T07:56:08Z

src/github.rs

+            // reqwest wants to take ownership of the body, so it's hard for us to do anything
+            // clever with reading the file once and calculating the sha256sum while we read.
+            // So we open and read the file again.
+            let mut file = tokio::fs::File::open(local_filename).await?;


nit: Wrapping the file in a BufReader gives better performance.

Does it? My understanding was that BufReader is helpful if you're doing lots of small reads, but here sha256 doesn't really have an opinion about the size of the buffer we pass in for each chunk, and I'm intentionally doing very large reads (1 MB). I think BufReader defaults to tokion::io::util::DEFAULT_BUF_SIZE = 8192 bytes, and it looks like the implementation bypasses the internal buffer if you ask for a larger read.

You're right, with the manual large buffer this is actually not necessary.

konstin · 2025-07-01T07:56:36Z

src/github.rs

                dry_run,
            ));
+
+            // reqwest wants to take ownership of the body, so it's hard for us to do anything


I think that's totally fine to read the file twice.

konstin · 2025-07-01T07:59:07Z

src/github.rs

+                if len == 0 {
+                    break;
+                };
+                hasher.update(&buf);


Do we need to slice the buf to actually read size? I don't know about the properties of SHA-512, but we're currently passing a variable number of trailing zero bytes.

Oof, yes. Not even zero bytes, whatever was at the end of the previous megabyte.

I feel like this is not the first time I have made this mistake in Rust, sigh. (This is also basically tedu's "heartbleed in Rust" from pre-1.0.)

I was surprised the api doesn't have a way to pass a whole Read or AsyncRead, that's way more idomatic and prevents such mistakes, I find myself avoiding the &mut buf APIs usually.

konstin · 2025-07-01T08:00:20Z

src/github.rs

-                format!("{}.sha256", dest),
-                Bytes::copy_from_slice(format!("{}\n", digest).as_bytes()),
+                format!("{dest}.sha256"),
+                UploadSource::Data(Bytes::copy_from_slice(format!("{digest}\n").as_bytes())),


Suggested change

UploadSource::Data(Bytes::copy_from_slice(format!("{digest}\n").as_bytes())),

UploadSource::Data(Bytes::from(format!("{digest}\n"))),

It was like that when I found it :) thanks

geofft · 2025-07-01T20:53:28Z

Marking as draft since I'm working on a mock GitHub HTTP server to do some fault injection to test things properly and I want to do those tests, but I'm mildly confident about this code at the moment.

This lets me test the release scripts against a custom, fault-injected Python server, but I suppose it might also be useful for downstream users who have GHES, maybe. Patches welcome if anyone is using this and it doesn't quite work right!

geofft · 2025-07-23T14:38:04Z

Sorry about the rebase, the relevant diff since the last push is

diff --git a/src/github.rs b/src/github.rs
index cc5c4e4..676104f 100644
--- a/src/github.rs
+++ b/src/github.rs
@@ -117,6 +117,7 @@ async fn upload_release_artifact(
     let response = loop {
         let request = client
             .put(url.clone())
+            .timeout(Duration::from_secs(60))
             .header("Authorization", format!("Bearer {auth_token}"))
             .header("Content-Type", "application/octet-stream");
         let request = match body {
@@ -134,7 +135,7 @@ async fn upload_release_artifact(
         if retryable_strategy.handle(&result) == Some(Retryable::Transient) {
             let retry_decision = retry_policy.should_retry(start_time, n_past_retries);
             if let reqwest_retry::RetryDecision::Retry { execute_after } = retry_decision {
-                println!("retrying {url}: {result:?}");
+                println!("retrying upload to {url} after {result:?}");
                 let duration = execute_after
                     .duration_since(SystemTime::now())
                     .unwrap_or_else(|_| Duration::default());

plus the addition of src/github_api_tester.py, which is a little Flask app to try to do some fault injection. The tests there pass but it doesn't seem like it's helpful to run the tests on every push, especially since they're kind of slow (since the uploader is Rust and the tests are Python, I can't take advantage of either language's async framework's autojump clock).

Adding the 60-second timeout was necessary to get the tests to pass, because for some reason the retry on 401 is hanging - the client believes it's sent the request and the server doesn't seem to be doing anything. While this is probably a bug in my test server (or its dependencies), it does bring up the point that we have no timeout on uploads and maybe one would be good. I'm not sure if 60 seconds is too small.

Oh, and if you missed it in the prior push, the thing above with mishandled buffer sizes was solved with by adding tokio-util and using its ReaderStream:

@@ -540,19 +539,15 @@ pub async fn command_upload_release_distributions(args: &ArgMatches) -> Result<(
             // reqwest wants to take ownership of the body, so it's hard for us to do anything
             // clever with reading the file once and calculating the sha256sum while we read.
             // So we open and read the file again.
-            let mut file = tokio::fs::File::open(local_filename).await?;
-            let mut hasher = Sha256::new();
-            let mut buf = vec![0; 1048576];
-            loop {
-                let len = file.read(&mut buf).await?;
-                if len == 0 {
-                    break;
-                };
-                hasher.update(&buf);
-            }
-            drop(file);
-
-            let digest = hex::encode(hasher.finalize());
+            let digest = {
+                let file = tokio::fs::File::open(local_filename).await?;
+                let mut stream = tokio_util::io::ReaderStream::with_capacity(file, 1048576);
+                let mut hasher = Sha256::new();
+                while let Some(chunk) = stream.next().await {
+                    hasher.update(&chunk?);
+                }
+                hex::encode(hasher.finalize())
+            };
             digests.insert(dest.clone(), digest.clone());
         }

Now that we're not loading every artifact in memory, hopefully the normal runner will work.

konstin

Didn't check the test server in detail, but the rust part looks good

geofft requested a review from zanieb June 29, 2025 21:45

geofft added the ci:skip label Jun 29, 2025

zanieb requested a review from konstin June 30, 2025 20:47

konstin reviewed Jul 1, 2025

View reviewed changes

geofft marked this pull request as draft July 1, 2025 20:50

geofft force-pushed the stream-upload branch 2 times, most recently from 19e271d to ed37a6a Compare July 23, 2025 14:32

Do not load the entire artifact in memory when uploading (astral-sh#618)

1e23d66

geofft force-pushed the stream-upload branch from ed37a6a to 1e23d66 Compare July 23, 2025 14:36

geofft marked this pull request as ready for review July 23, 2025 14:38

Go back to the normal runner for the releease process

14b9a59

Now that we're not loading every artifact in memory, hopefully the normal runner will work.

konstin approved these changes Jul 24, 2025

View reviewed changes

geofft merged commit a376f32 into astral-sh:main Jul 25, 2025
10 checks passed

	UploadSource::Data(Bytes::copy_from_slice(format!("{digest}\n").as_bytes())),
	UploadSource::Data(Bytes::from(format!("{digest}\n"))),

Uh oh!

Do not load the entire artifact in memory when uploading (#618) #677

Do not load the entire artifact in memory when uploading (#618) #677

Conversation

geofft commented Jun 29, 2025

Uh oh!

geofft commented Jun 29, 2025

Uh oh!

zanieb commented Jun 30, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

geofft commented Jul 1, 2025

Uh oh!

geofft commented Jul 23, 2025

Uh oh!

konstin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!