diff --git a/Cargo.lock b/Cargo.lock index b0db1676..ec65e738 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1585,6 +1585,7 @@ dependencies = [ "tempfile", "text-stub-library", "tokio", + "tokio-util", "url", "version-compare", "zip", diff --git a/Cargo.toml b/Cargo.toml index 82f02982..fbdf1a7d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -37,6 +37,7 @@ tar = "0.4.40" tempfile = "3.10.0" text-stub-library = "0.9.0" tokio = "1.43.1" +tokio-util = "0.7.13" url = "2.5.0" version-compare = "0.1.1" zip = "0.6.6" diff --git a/src/github.rs b/src/github.rs index 0d3bc971..bf94e7c6 100644 --- a/src/github.rs +++ b/src/github.rs @@ -17,10 +17,9 @@ use { }, rayon::prelude::*, reqwest::{Client, StatusCode}, - reqwest_middleware::{self, ClientWithMiddleware}, reqwest_retry::{ - default_on_request_failure, policies::ExponentialBackoff, RetryTransientMiddleware, - Retryable, RetryableStrategy, + default_on_request_failure, policies::ExponentialBackoff, RetryPolicy, Retryable, + RetryableStrategy, }, sha2::{Digest, Sha256}, std::{ @@ -28,6 +27,7 @@ use { io::Read, path::PathBuf, str::FromStr, + time::{Duration, SystemTime}, }, url::Url, zip::ZipArchive, @@ -65,12 +65,19 @@ async fn fetch_artifact( Ok(res) } +enum UploadSource { + Filename(PathBuf), + Data(Bytes), +} + async fn upload_release_artifact( - client: &ClientWithMiddleware, + client: &Client, + retry_policy: &impl RetryPolicy, + retryable_strategy: &impl RetryableStrategy, auth_token: String, release: &Release, filename: String, - data: Bytes, + body: UploadSource, dry_run: bool, ) -> Result<()> { if release.assets.iter().any(|asset| asset.name == filename) { @@ -93,17 +100,51 @@ async fn upload_release_artifact( return Ok(()); } - // Octocrab doesn't yet support release artifact upload. And the low-level HTTP API - // forces the use of strings on us. So we have to make our own HTTP client. - - let response = client - .put(url) - .header("Authorization", format!("Bearer {auth_token}")) - .header("Content-Length", data.len()) - .header("Content-Type", "application/x-tar") - .body(data) - .send() - .await?; + // Octocrab's high-level API for uploading release artifacts doesn't yet support streaming + // bodies, and their low-level API isn't more helpful than using our own HTTP client. + // + // Because we are streaming the body, we can't use the standard retry middleware for reqwest + // (see e.g. https://github.com/seanmonstar/reqwest/issues/2416), so we have to recreate the + // request on each retry and handle the retry logic ourself. This logic is inspired by + // uv/crates/uv-publish/src/lib.rs (which has the same problem), which in turn is inspired by + // reqwest-middleware/reqwest-retry/src/middleware.rs. + // + // (While Octocrab's API would work fine for the non-streaming case, we just use this function + // for both cases so that we can make a homogeneous Vec later in the file.) + + let mut n_past_retries = 0; + let start_time = SystemTime::now(); + let response = loop { + let request = client + .put(url.clone()) + .header("Authorization", format!("Bearer {auth_token}")) + .header("Content-Type", "application/octet-stream"); + let request = match body { + UploadSource::Filename(ref path) => { + let file = tokio::fs::File::open(&path).await?; + let len = file.metadata().await?.len(); + request.header("Content-Length", len).body(file) + } + UploadSource::Data(ref bytes) => request + .header("Content-Length", bytes.len()) + .body(bytes.clone()), + }; + 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); + if let reqwest_retry::RetryDecision::Retry { execute_after } = retry_decision { + println!("retrying {url}: {result:?}"); + let duration = execute_after + .duration_since(SystemTime::now()) + .unwrap_or_else(|_| Duration::default()); + tokio::time::sleep(duration).await; + n_past_retries += 1; + continue; + } + } + break result?; + }; if !response.status().is_success() { return Err(anyhow!("HTTP {}", response.status())); @@ -112,6 +153,20 @@ async fn upload_release_artifact( Ok(()) } +fn new_github_client(args: &ArgMatches) -> Result<(Octocrab, String)> { + let token = args + .get_one::("token") + .expect("token should be specified") + .to_string(); + let github_uri = args.get_one::("github-uri"); + + let mut builder = OctocrabBuilder::new().personal_token(token.clone()); + if let Some(github_uri) = github_uri { + builder = builder.base_uri(github_uri.clone())?; + } + Ok((builder.build()?, token)) +} + pub async fn command_fetch_release_distributions(args: &ArgMatches) -> Result<()> { let dest_dir = args .get_one::("dest") @@ -121,13 +176,7 @@ pub async fn command_fetch_release_distributions(args: &ArgMatches) -> Result<() .expect("organization should be set"); let repo = args.get_one::("repo").expect("repo should be set"); - let client = OctocrabBuilder::new() - .personal_token( - args.get_one::("token") - .expect("token should be required argument") - .to_string(), - ) - .build()?; + let (client, _) = new_github_client(args)?; let release_version_range = pep440_rs::VersionSpecifier::from_str(">=3.9")?; @@ -207,10 +256,8 @@ pub async fn command_fetch_release_distributions(args: &ArgMatches) -> Result<() .await?; for artifact in artifacts { - if matches!( - artifact.name.as_str(), - "pythonbuild" | "toolchain" - ) || artifact.name.contains("install-only") + if matches!(artifact.name.as_str(), "pythonbuild" | "toolchain") + || artifact.name.contains("install-only") { continue; } @@ -358,10 +405,6 @@ pub async fn command_upload_release_distributions(args: &ArgMatches) -> Result<( .get_one::("tag") .expect("tag should be specified"); let ignore_missing = args.get_flag("ignore_missing"); - let token = args - .get_one::("token") - .expect("token should be specified") - .to_string(); let organization = args .get_one::("organization") .expect("organization should be specified"); @@ -451,9 +494,7 @@ pub async fn command_upload_release_distributions(args: &ArgMatches) -> Result<( return Err(anyhow!("missing {} release artifacts", missing.len())); } - let client = OctocrabBuilder::new() - .personal_token(token.clone()) - .build()?; + let (client, token) = new_github_client(args)?; let repo_handler = client.repos(organization, repo); let releases = repo_handler.releases(); @@ -473,12 +514,7 @@ pub async fn command_upload_release_distributions(args: &ArgMatches) -> Result<( let mut digests = BTreeMap::new(); let retry_policy = ExponentialBackoff::builder().build_with_max_retries(5); - let raw_client = reqwest_middleware::ClientBuilder::new(Client::new()) - .with(RetryTransientMiddleware::new_with_policy_and_strategy( - retry_policy, - GitHubUploadRetryStrategy, - )) - .build(); + let raw_client = Client::new(); { let mut fs = vec![]; @@ -488,29 +524,39 @@ pub async fn command_upload_release_distributions(args: &ArgMatches) -> Result<( continue; } - let file_data = Bytes::copy_from_slice(&std::fs::read(dist_dir.join(&source))?); - - let mut digest = Sha256::new(); - digest.update(&file_data); - - let digest = hex::encode(digest.finalize()); - - digests.insert(dest.clone(), digest.clone()); - + let local_filename = dist_dir.join(&source); fs.push(upload_release_artifact( &raw_client, + &retry_policy, + &GitHubUploadRetryStrategy, token.clone(), &release, dest.clone(), - file_data, + UploadSource::Filename(local_filename.clone()), dry_run, )); + + // 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 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()); fs.push(upload_release_artifact( &raw_client, + &retry_policy, + &GitHubUploadRetryStrategy, token.clone(), &release, - 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())), dry_run, )); } @@ -532,10 +578,12 @@ pub async fn command_upload_release_distributions(args: &ArgMatches) -> Result<( upload_release_artifact( &raw_client, + &retry_policy, + &GitHubUploadRetryStrategy, token.clone(), &release, "SHA256SUMS".to_string(), - Bytes::copy_from_slice(shasums.as_bytes()), + UploadSource::Data(Bytes::copy_from_slice(shasums.as_bytes())), dry_run, ) .await?; diff --git a/src/main.rs b/src/main.rs index df441e13..c8d422e8 100644 --- a/src/main.rs +++ b/src/main.rs @@ -70,6 +70,12 @@ fn main_impl() -> Result<()> { .action(ArgAction::Set) .default_value("python-build-standalone") .help("GitHub repository name"), + ) + .arg( + Arg::new("github-uri") + .long("github-uri") + .action(ArgAction::Set) + .help("Alternative GitHub URI"), ), ); @@ -154,6 +160,12 @@ fn main_impl() -> Result<()> { .action(ArgAction::Set) .default_value("python-build-standalone") .help("GitHub repository name"), + ) + .arg( + Arg::new("github-uri") + .long("github-uri") + .action(ArgAction::Set) + .help("Alternative GitHub URI"), ), );