Fetch `--find-links` indexes in parallel (#934)

## Summary

Removes a TODO.

## Test Plan

Tested manually with:

```shell
cargo run -p puffin-cli -- \
    pip compile requirements.in -n \
    --find-links 'https://download.pytorch.org/whl/torch_stable.html' \
    --find-links 'https://storage.googleapis.com/jax-releases/jax_cuda_releases.html' \
    --verbose
```

And inspecting the logs to ensure that the two requests were kicked off
concrurently.
This commit is contained in:
Charlie Marsh 2024-01-16 05:37:35 -05:00 committed by GitHub
parent 2f8f126f2f
commit b50e5fcbc5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 29 additions and 23 deletions

View File

@ -2,6 +2,7 @@ use std::collections::btree_map::Entry;
use std::collections::BTreeMap; use std::collections::BTreeMap;
use std::path::PathBuf; use std::path::PathBuf;
use futures::StreamExt;
use reqwest::Response; use reqwest::Response;
use rustc_hash::FxHashMap; use rustc_hash::FxHashMap;
use tracing::{debug, info_span, instrument, warn, Instrument}; use tracing::{debug, info_span, instrument, warn, Instrument};
@ -53,30 +54,35 @@ impl<'a> FlatIndexClient<'a> {
&self, &self,
indexes: impl Iterator<Item = &FlatIndexLocation>, indexes: impl Iterator<Item = &FlatIndexLocation>,
) -> Result<Vec<FlatIndexEntry>, FlatIndexError> { ) -> Result<Vec<FlatIndexEntry>, FlatIndexError> {
let mut dists = Vec::new(); let mut fetches = futures::stream::iter(indexes)
// TODO(konstin): Parallelize reads over flat indexes. .map(|index| async move {
for flat_index in indexes { let entries = match index {
let index_dists = match flat_index { FlatIndexLocation::Path(path) => Self::read_from_directory(path)
FlatIndexLocation::Path(path) => Self::read_from_directory(path) .map_err(|err| FlatIndexError::FindLinksDirectory(path.clone(), err))?,
.map_err(|err| FlatIndexError::FindLinksDirectory(path.clone(), err))?, FlatIndexLocation::Url(url) => self
FlatIndexLocation::Url(url) => self .read_from_url(url)
.read_from_url(url) .await
.await .map_err(|err| FlatIndexError::FindLinksUrl(url.clone(), err))?,
.map_err(|err| FlatIndexError::FindLinksUrl(url.clone(), err))?, };
}; if entries.is_empty() {
if index_dists.is_empty() { warn!("No packages found in `--find-links` entry: {}", index);
warn!("No packages found in `--find-links` entry: {}", flat_index); } else {
} else { debug!(
debug!( "Found {} package{} in `--find-links` entry: {}",
"Found {} package{} in `--find-links` entry: {}", entries.len(),
index_dists.len(), if entries.len() == 1 { "" } else { "s" },
if index_dists.len() == 1 { "" } else { "s" }, index
flat_index );
); }
} Ok::<Vec<FlatIndexEntry>, FlatIndexError>(entries)
dists.extend(index_dists); })
.buffered(16);
let mut results = Vec::new();
while let Some(entries) = fetches.next().await.transpose()? {
results.extend(entries);
} }
Ok(dists) Ok(results)
} }
/// Read a flat remote index from a `--find-links` URL. /// Read a flat remote index from a `--find-links` URL.