Yield after channel send and move cpu tasks to thread

Author: konstin

crates/puffin-client/src/cached_client.rs

@@ -104,7 +104,7 @@ impl CachedClient {
     /// client.
     #[instrument(skip_all)]
     pub async fn get_cached_with_callback<
-        Payload: Serialize + DeserializeOwned + Send,
+        Payload: Serialize + DeserializeOwned + Send + 'static,
         CallBackError,
         Callback,
         CallbackReturn,
@@ -172,7 +172,7 @@ impl CachedClient {
         }
     }
 
-    async fn read_cache<Payload: Serialize + DeserializeOwned + Send>(
+    async fn read_cache<Payload: Serialize + DeserializeOwned + Send + 'static>(
         cache_entry: &CacheEntry,
     ) -> Option<DataWithCachePolicy<Payload>> {
         let read_span = info_span!("read_cache", file = %cache_entry.path().display());
@@ -185,8 +185,12 @@ impl CachedClient {
                 "parse_cache",
                 path = %cache_entry.path().display()
             );
-            let parse_result = parse_span
-                .in_scope(|| rmp_serde::from_slice::<DataWithCachePolicy<Payload>>(&cached));
+            let parse_result = tokio::task::spawn_blocking(move || {
+                parse_span
+                    .in_scope(|| rmp_serde::from_slice::<DataWithCachePolicy<Payload>>(&cached))
+            })
+            .await
+            .expect("Tokio executor failed, was there a panic?");
             match parse_result {
                 Ok(data) => Some(data),
                 Err(err) => {

crates/puffin-resolver/src/resolver/mod.rs

@@ -5,7 +5,9 @@ use std::sync::Arc;
 
 use anyhow::Result;
 use dashmap::{DashMap, DashSet};
-use futures::channel::mpsc::UnboundedReceiver;
+use futures::channel::mpsc::{
+    UnboundedReceiver as MpscUnboundedReceiver, UnboundedSender as MpscUnboundedSender,
+};
 use futures::{FutureExt, StreamExt};
 use itertools::Itertools;
 use pubgrub::error::PubGrubError;
@@ -239,7 +241,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
     #[instrument(skip_all)]
     async fn solve(
         &self,
-        request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
+        request_sink: &MpscUnboundedSender<Request>,
     ) -> Result<ResolutionGraph, ResolveError> {
         let root = PubGrubPackage::Root(self.project.clone());
 
@@ -384,7 +386,7 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
         &self,
         package: &PubGrubPackage,
         priorities: &mut PubGrubPriorities,
-        request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
+        request_sink: &MpscUnboundedSender<Request>,
     ) -> Result<(), ResolveError> {
         match package {
             PubGrubPackage::Root(_) => {}
@@ -411,14 +413,16 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                 }
             }
         }
+        // Yield after sending on a channel to allow the subscribers to continue
+        tokio::task::yield_now().await;
         Ok(())
     }
 
     /// Visit the set of [`PubGrubPackage`] candidates prior to selection. This allows us to fetch
     /// metadata for all of the packages in parallel.
     fn pre_visit<'data>(
         packages: impl Iterator<Item = (&'data PubGrubPackage, &'data Range<Version>)>,
-        request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
+        request_sink: &MpscUnboundedSender<Request>,
     ) -> Result<(), ResolveError> {
         // Iterate over the potential packages, and fetch file metadata for any of them. These
         // represent our current best guesses for the versions that we _might_ select.
@@ -439,9 +443,9 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
         package: &PubGrubPackage,
         range: &Range<Version>,
         pins: &mut FilePins,
-        request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
+        request_sink: &MpscUnboundedSender<Request>,
     ) -> Result<Option<Version>, ResolveError> {
-        return match package {
+        let result = match package {
             PubGrubPackage::Root(_) => Ok(Some(MIN_VERSION.clone())),
 
             PubGrubPackage::Python(PubGrubPython::Installed) => {
@@ -581,15 +585,19 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
                 Ok(Some(version))
             }
         };
+        // Yield after sending on a channel to allow the subscribers to continue
+        tokio::task::yield_now().await;
+        result
     }
 
     /// Given a candidate package and version, return its dependencies.
+    #[instrument(skip_all, fields(%package, %version))]
     async fn get_dependencies(
         &self,
         package: &PubGrubPackage,
         version: &Version,
         priorities: &mut PubGrubPriorities,
-        request_sink: &futures::channel::mpsc::UnboundedSender<Request>,
+        request_sink: &MpscUnboundedSender<Request>,
     ) -> Result<Dependencies, ResolveError> {
         match package {
             PubGrubPackage::Root(_) => {
@@ -714,7 +722,10 @@ impl<'a, Provider: ResolverProvider> Resolver<'a, Provider> {
     }
 
     /// Fetch the metadata for a stream of packages and versions.
-    async fn fetch(&self, request_stream: UnboundedReceiver<Request>) -> Result<(), ResolveError> {
+    async fn fetch(
+        &self,
+        request_stream: MpscUnboundedReceiver<Request>,
+    ) -> Result<(), ResolveError> {
         let mut response_stream = request_stream
             .map(|request| self.process_request(request))
             .buffer_unordered(50);
@@ -904,10 +915,10 @@ impl Display for Request {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
             Request::Package(package_name) => {
-                write!(f, "Package {package_name}")
+                write!(f, "Versions {package_name}")
             }
             Request::Dist(dist) => {
-                write!(f, "Dist {dist}")
+                write!(f, "Metadata {dist}")
             }
             Request::Prefetch(package_name, range) => {
                 write!(f, "Prefetch {package_name} {range}")

crates/puffin-resolver/src/resolver/provider.rs

@@ -1,13 +1,15 @@
 use std::future::Future;
+use std::ops::Deref;
+use std::sync::Arc;
 
 use anyhow::Result;
 use chrono::{DateTime, Utc};
 use futures::FutureExt;
 use url::Url;
 
-use distribution_types::Dist;
+use distribution_types::{Dist, IndexUrl};
 use platform_tags::Tags;
-use puffin_client::{FlatIndex, RegistryClient};
+use puffin_client::{FlatIndex, RegistryClient, SimpleMetadata};
 use puffin_distribution::DistributionDatabase;
 use puffin_normalize::PackageName;
 use puffin_traits::{BuildContext, NoBinary};
@@ -45,17 +47,29 @@ pub trait ResolverProvider: Send + Sync {
 /// The main IO backend for the resolver, which does cached requests network requests using the
 /// [`RegistryClient`] and [`DistributionDatabase`].
 pub struct DefaultResolverProvider<'a, Context: BuildContext + Send + Sync> {
-    /// The [`RegistryClient`] used to query the index.
-    client: &'a RegistryClient,
     /// The [`DistributionDatabase`] used to build source distributions.
     fetcher: DistributionDatabase<'a, Context>,
+    inner: Arc<DefaultResolverProviderInner>,
+}
+
+pub struct DefaultResolverProviderInner {
+    /// The [`RegistryClient`] used to query the index.
+    client: RegistryClient,
     /// These are the entries from `--find-links` that act as overrides for index responses.
-    flat_index: &'a FlatIndex,
-    tags: &'a Tags,
+    flat_index: FlatIndex,
+    tags: Tags,
     python_requirement: PythonRequirement,
     exclude_newer: Option<DateTime<Utc>>,
     allowed_yanks: AllowedYanks,
-    no_binary: &'a NoBinary,
+    no_binary: NoBinary,
+}
+
+impl<'a, Context: BuildContext + Send + Sync> Deref for DefaultResolverProvider<'a, Context> {
+    type Target = DefaultResolverProviderInner;
+
+    fn deref(&self) -> &Self::Target {
+        self.inner.as_ref()
+    }
 }
 
 impl<'a, Context: BuildContext + Send + Sync> DefaultResolverProvider<'a, Context> {
@@ -72,50 +86,69 @@ impl<'a, Context: BuildContext + Send + Sync> DefaultResolverProvider<'a, Contex
         no_binary: &'a NoBinary,
     ) -> Self {
         Self {
-            client,
             fetcher,
-            flat_index,
-            tags,
-            python_requirement,
-            exclude_newer,
-            allowed_yanks,
-            no_binary,
+            inner: Arc::new(DefaultResolverProviderInner {
+                client: client.clone(),
+                flat_index: flat_index.clone(),
+                tags: tags.clone(),
+                python_requirement,
+                exclude_newer,
+                allowed_yanks,
+                no_binary: no_binary.clone(),
+            }),
         }
     }
 }
 
+fn simple_to_version_map(
+    self_arc: Arc<DefaultResolverProviderInner>,
+    result: Result<(IndexUrl, SimpleMetadata), puffin_client::Error>,
+    package_name: &PackageName,
+) -> VersionMapResponse {
+    match result {
+        Ok((index, metadata)) => Ok(VersionMap::from_metadata(
+            metadata,
+            &package_name,
+            &index,
+            &self_arc.tags,
+            &self_arc.python_requirement,
+            &self_arc.allowed_yanks,
+            self_arc.exclude_newer.as_ref(),
+            self_arc.flat_index.get(&package_name).cloned(),
+            &self_arc.no_binary,
+        )),
+        Err(err) => match err.into_kind() {
+            kind @ (puffin_client::ErrorKind::PackageNotFound(_)
+            | puffin_client::ErrorKind::NoIndex(_)) => {
+                if let Some(flat_index) = self_arc.flat_index.get(&package_name).cloned() {
+                    Ok(VersionMap::from(flat_index))
+                } else {
+                    Err(kind.into())
+                }
+            }
+            kind => Err(kind.into()),
+        },
+    }
+}
+
 impl<'a, Context: BuildContext + Send + Sync> ResolverProvider
     for DefaultResolverProvider<'a, Context>
 {
     fn get_version_map<'io>(
         &'io self,
         package_name: &'io PackageName,
     ) -> impl Future<Output = VersionMapResponse> + Send + 'io {
-        self.client
-            .simple(package_name)
-            .map(move |result| match result {
-                Ok((index, metadata)) => Ok(VersionMap::from_metadata(
-                    metadata,
-                    package_name,
-                    &index,
-                    self.tags,
-                    &self.python_requirement,
-                    &self.allowed_yanks,
-                    self.exclude_newer.as_ref(),
-                    self.flat_index.get(package_name).cloned(),
-                    self.no_binary,
-                )),
-                Err(err) => match err.into_kind() {
-                    kind @ (puffin_client::ErrorKind::PackageNotFound(_)
-                    | puffin_client::ErrorKind::NoIndex(_)) => {
-                        if let Some(flat_index) = self.flat_index.get(package_name).cloned() {
-                            Ok(VersionMap::from(flat_index))
-                        } else {
-                            Err(kind.into())
-                        }
-                    }
-                    kind => Err(kind.into()),
-                },
+        let package_name_owned = package_name.clone();
+        let self_arc = self.inner.clone();
+        self.inner
+            .client
+            .simple(&package_name)
+            .then(|result| async move {
+                tokio::task::spawn_blocking(move || {
+                    simple_to_version_map(self_arc, result, &package_name_owned)
+                })
+                .await
+                .unwrap()
             })
     }
 

crates/puffin-traits/src/lib.rs

@@ -160,7 +160,7 @@ impl Display for BuildKind {
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub enum NoBinary {
     /// Allow installation of any wheel.
     None,