added multi client support for overlay

Author: mkatychev

crates/wasm-pkg-client/src/lib.rs

@@ -46,6 +46,7 @@ use tokio::io::AsyncSeekExt;
 use tokio::sync::RwLock;
 use tokio_util::io::SyncIoBridge;
 use wasm_pkg_common::config::RegistryConfig;
+use wasm_pkg_common::metadata::{LOCAL_PROTOCOL, OCI_PROTOCOL, WARG_PROTOCOL};
 pub use wasm_pkg_common::{
     config::{Config, CustomConfig, RegistryMapping},
     digest::ContentDigest,
@@ -219,76 +220,6 @@ impl Client {
             .ok_or_else(|| Error::NoRegistryForNamespace(package.namespace().clone()))
     }
 
-    fn resolve_backend(
-        &self,
-        registry: &Registry,
-        registry_meta: RegistryMetadata,
-    ) -> Result<InnerClient, Error> {
-        let registry_config = self
-            .config
-            .registry_config(&registry)
-            .cloned()
-            .unwrap_or_default();
-
-        let backend_type = match registry_config.default_backend() {
-            // If the local config specifies a backend type, use it
-            Some(backend_type) => Some(backend_type),
-            None => {
-                // If the registry metadata indicates a preferred protocol, use it
-                let preferred_protocol = registry_meta.preferred_protocol();
-                // ...except registry metadata cannot force a local backend
-                if preferred_protocol == Some("local") {
-                    return Err(Error::InvalidRegistryMetadata(anyhow!(
-                        "registry metadata with 'local' protocol not allowed"
-                    )));
-                }
-                preferred_protocol
-            }
-        }
-        // Otherwise use the default backend
-        .unwrap_or("oci");
-
-        todo!();
-    }
-
-    fn resolve_metadata(
-        &self,
-        package: &PackageRef,
-        registry_config: &RegistryConfig,
-        registry_override: bool,
-        registry: &Registry,
-    ) -> &RegistryMetadata {
-        if let Some(metadata) = self
-            .config
-            .package_registry_override(package)
-            .and_then(|mapping| mapping.metadata())
-        {
-            return metadata;
-        }
-
-        let _ = self
-            .config
-            .namespace_registry(package.namespace())
-            .and_then(|mapping| {
-                // If the overridden registry matches the registry we are trying to resolve, we
-                // should use the metadata, otherwise we'll need to fetch the metadata from the
-                // registry
-                match (mapping, registry_override) {
-                    (RegistryMapping::Custom(custom), true) if custom.registry == *registry => {
-                        Some(custom.metadata.clone())
-                    }
-                    (RegistryMapping::Custom(custom), false) => Some(custom.metadata.clone()),
-                    _ => None,
-                }
-            });
-        todo!();
-    }
-
-    async fn has_registry(&self, registry: &Registry) -> bool {
-        let sources = self.sources.read().await;
-        sources.contains_key(registry)
-    }
-
     async fn resolve_source(
         &self,
         package: &PackageRef,
@@ -308,7 +239,7 @@ impl Client {
             .unwrap_or_default();
 
         // Skip fetching metadata for "local" source
-        let should_fetch_meta = registry_config.default_backend() != Some("local");
+        let should_fetch_meta = registry_config.default_backend() != LOCAL_PROTOCOL.into();
         let maybe_metadata = self
             .config
             .package_registry_override(package)
@@ -353,7 +284,7 @@ impl Client {
                 // If the registry metadata indicates a preferred protocol, use it
                 let preferred_protocol = registry_meta.preferred_protocol();
                 // ...except registry metadata cannot force a local backend
-                if preferred_protocol == Some("local") {
+                if preferred_protocol == Some(LOCAL_PROTOCOL) {
                     return Err(Error::InvalidRegistryMetadata(anyhow!(
                         "registry metadata with 'local' protocol not allowed"
                     )));
@@ -362,17 +293,17 @@ impl Client {
             }
         }
         // Otherwise use the default backend
-        .unwrap_or("oci");
+        .unwrap_or(OCI_PROTOCOL);
         tracing::debug!(?backend_type, "Resolved backend type");
 
         let source: InnerClient = match backend_type {
-            "local" => Box::new(LocalBackend::new(registry_config)?),
-            "oci" => Box::new(OciBackend::new(
+            LOCAL_PROTOCOL => Box::new(LocalBackend::new(registry_config)?),
+            OCI_PROTOCOL => Box::new(OciBackend::new(
                 &registry,
                 &registry_config,
                 &registry_meta,
             )?),
-            "warg" => {
+            WARG_PROTOCOL => {
                 Box::new(WargBackend::new(&registry, &registry_config, &registry_meta).await?)
             }
             other => {

crates/wasm-pkg-client/src/overlay/mod.rs

@@ -1,4 +1,4 @@
-use std::io::Cursor;
+use std::{collections::HashMap, io::Cursor};
 
 use async_trait::async_trait;
 use tempfile::TempDir;
@@ -14,28 +14,34 @@ use crate::{
 
 pub(crate) struct OverlayBackend {
     local: LocalBackend,
-    remote: InnerClient,
+    remotes: HashMap<PackageRef, InnerClient>,
     _handle: TempDir,
 }
 
 impl OverlayBackend {
-    fn new(remote: InnerClient) -> Result<Self, Error> {
+    fn new(remotes: HashMap<PackageRef, InnerClient>) -> Result<Self, Error> {
         let handle = TempDir::new()?;
         let root = handle.path().to_owned();
         let local = LocalBackend { root };
         Ok(Self {
             local,
-            remote,
+            remotes,
             _handle: handle,
         })
     }
+
+    fn remote(&self, package: &PackageRef) -> Result<&InnerClient, Error> {
+        self.remotes
+            .get(package)
+            .ok_or_else(|| Error::InvalidPackageRef(package.to_string()))
+    }
 }
 
 #[async_trait]
 impl PackageLoader for OverlayBackend {
     async fn list_all_versions(&self, package: &PackageRef) -> Result<Vec<VersionInfo>, Error> {
         let mut versions = self.local.list_all_versions(package).await?;
-        let mut remote_versions = self.remote.list_all_versions(package).await?;
+        let mut remote_versions = self.remote(package)?.list_all_versions(package).await?;
         versions.append(&mut remote_versions);
         versions.sort();
         versions.dedup();
@@ -47,7 +53,7 @@ impl PackageLoader for OverlayBackend {
             return Ok(release);
         }
         tracing::debug!(%package, %version, method = "get_release", "OverlayBackend falling back to remote");
-        self.remote.get_release(package, version).await
+        self.remote(package)?.get_release(package, version).await
     }
 
     async fn stream_content_unvalidated(
@@ -82,13 +88,14 @@ impl PackagePublisher for OverlayBackend {
         let mut local_data = Box::pin(Cursor::new(Vec::new()));
         tokio::io::copy(&mut data, &mut local_data).await?;
 
-        self.local
-            .publish(&package, &version, local_data.clone(), dry_run)
-            .await?;
         if dry_run {
+            self.local
+                .publish(&package, &version, local_data.clone(), dry_run)
+                .await?;
             return Ok(());
         }
-        self.remote
+
+        self.remote(package)?
             .publish(&package, &version, local_data, dry_run)
             .await?;
 

crates/wasm-pkg-common/src/config.rs

@@ -204,12 +204,7 @@ impl Config {
         let namespace = package.namespace();
         // look in `self.package_registry_overrides `
         // then in `self.namespace_registries`
-        if let Some(reg) = self
-            .package_registry_overrides
-            .get(package)
-            .or_else(|| self.namespace_registries.get(namespace))
-            .map(|pkg| pkg.registry())
-        {
+        if let Some(reg) = self.resolve_mapping(package).map(|ns| ns.registry()) {
             return Some(reg);
         } else if let Some(reg) = self.default_registry.as_ref() {
             return Some(reg);
@@ -218,6 +213,15 @@ impl Config {
         self.fallback_namespace_registries.get(namespace)
     }
 
+    pub fn resolve_mapping(&self, package: &PackageRef) -> Option<&RegistryMapping> {
+        let namespace = package.namespace();
+        // look in `self.package_registry_overrides `
+        // then in `self.namespace_registries`
+        self.package_registry_overrides
+            .get(package)
+            .or_else(|| self.namespace_registries.get(namespace))
+    }
+
     /// Returns the default registry.
     pub fn default_registry(&self) -> Option<&Registry> {
         self.default_registry.as_ref()

crates/wasm-pkg-common/src/metadata.rs

@@ -41,7 +41,7 @@ pub const OCI_PROTOCOL: &str = "oci";
 /// Warg registry key
 pub const WARG_PROTOCOL: &str = "warg";
 /// Local filesystem key
-pub const LOCAL_FS: &str = "local";
+pub const LOCAL_PROTOCOL: &str = "local";
 
 impl RegistryMetadata {
     /// Returns the registry's preferred protocol.

crates/wasm-pkg-common/src/package.rs

@@ -23,7 +23,7 @@ impl PackageRef {
     }
 
     /// Returns the namespace of the package.
-    pub fn namespace(&self) -> &Label {
+    pub const fn namespace(&self) -> &Label {
         &self.namespace
     }