[PATCH datacenter-manager v2 12/20] parallel fetcher: support a custom client factory

[Lukas Wagner <l.wagner@proxmox.com>]

Add a client_factory() builder method and thread the resulting handle
through fetch_remote/fetch_node, defaulting to the client factory of
the current PdmApplication when none is set explicitly. This lets
callers, and later tests, supply a fake client factory instead of
always going through the real one.

Signed-off-by: Lukas Wagner <l.wagner@proxmox.com>
---
 server/src/parallel_fetcher.rs | 34 ++++++++++++++++++++++++++++++++--
 1 file changed, 32 insertions(+), 2 deletions(-)

diff --git a/server/src/parallel_fetcher.rs b/server/src/parallel_fetcher.rs
index 0819eb86..74ba91b3 100644
--- a/server/src/parallel_fetcher.rs
+++ b/server/src/parallel_fetcher.rs
@@ -73,7 +73,8 @@ use pve_api_types::ClusterNodeIndexResponse;
 
 use pdm_api_types::remotes::{Remote, RemoteType};
 
-use crate::connection;
+use crate::connection::ClientFactory;
+use crate::context;
 
 /// Maximum number of parallel outgoing API requests.
 pub const DEFAULT_MAX_CONNECTIONS: usize = 20;
@@ -229,6 +230,7 @@ impl<T> NodeResponse<T> {
 pub struct ParallelFetcherBuilder<C> {
     max_connections: Option<usize>,
     max_connections_per_remote: Option<usize>,
+    client_factory: Option<Arc<dyn ClientFactory + Send + Sync>>,
     context: C,
 }
 
@@ -238,6 +240,7 @@ impl<C> ParallelFetcherBuilder<C> {
             context,
             max_connections: None,
             max_connections_per_remote: None,
+            client_factory: None,
         }
     }
 
@@ -255,6 +258,12 @@ impl<C> ParallelFetcherBuilder<C> {
         self
     }
 
+    /// Set the client factory that should be used.
+    pub fn client_factory(mut self, client_factory: Arc<dyn ClientFactory + Send + Sync>) -> Self {
+        self.client_factory = Some(client_factory);
+        self
+    }
+
     /// Build the [`ParallelFetcher`] instance.
     pub fn build(self) -> ParallelFetcher<C> {
         ParallelFetcher {
@@ -262,6 +271,9 @@ impl<C> ParallelFetcherBuilder<C> {
             max_connections_per_remote: self
                 .max_connections_per_remote
                 .unwrap_or(DEFAULT_MAX_CONNECTIONS_PER_REMOTE),
+            client_factory: self
+                .client_factory
+                .unwrap_or_else(|| context::pdm_application().client_factory_shared()),
             context: self.context,
         }
     }
@@ -279,12 +291,15 @@ pub struct ParallelFetcherArgs<C> {
     /// The node. This may be 'localhost' for PBS remotes or if using
     /// [`ParallelFetcher::do_for_all_remotes`].
     pub node: String,
+    /// A handle to the client factory.
+    pub client_factory: Arc<dyn ClientFactory + Send + Sync>,
 }
 
 /// Helper for parallelizing API requests to multiple remotes/nodes.
 pub struct ParallelFetcher<C> {
     max_connections: usize,
     max_connections_per_remote: usize,
+    client_factory: Arc<dyn ClientFactory + Send + Sync>,
     context: C,
 }
 
@@ -318,13 +333,16 @@ impl<C: Clone + Send + 'static> ParallelFetcher<C> {
         for remote in remotes {
             let semaphore = Arc::clone(&total_connections_semaphore);
 
+            let client_factory = Arc::clone(&self.client_factory);
             let f = func.clone();
+
             let future = Self::fetch_remote(
                 remote,
                 self.context.clone(),
                 semaphore,
                 f,
                 self.max_connections_per_remote,
+                client_factory,
             );
 
             if let Some(log_context) = LogContext::current() {
@@ -356,6 +374,7 @@ impl<C: Clone + Send + 'static> ParallelFetcher<C> {
         semaphore: Arc<Semaphore>,
         func: F,
         max_connections_per_remote: usize,
+        client_factory: Arc<dyn ClientFactory + Send + Sync>,
     ) -> RemoteResponse<MultipleNodesResponse<T>>
     where
         F: Fn(ParallelFetcherArgs<C>) -> Ft + Clone + Send + 'static,
@@ -371,8 +390,10 @@ impl<C: Clone + Send + 'static> ParallelFetcher<C> {
             RemoteType::Pve => {
                 let remote_clone = remote.clone();
 
+                let cf = Arc::clone(&client_factory);
+
                 let nodes = match async move {
-                    let client = connection::make_pve_client(&remote_clone)?;
+                    let client = cf.make_pve_client(&remote_clone)?;
                     let nodes = client.list_nodes().await?;
 
                     Ok::<Vec<ClusterNodeIndexResponse>, Error>(nodes)
@@ -407,12 +428,14 @@ impl<C: Clone + Send + 'static> ParallelFetcher<C> {
                     let remote_clone = remote.clone();
                     let node_name = node.node.clone();
                     let context_clone = context.clone();
+                    let client_factory = Arc::clone(&client_factory);
 
                     let future = Self::fetch_node(
                         func_clone,
                         context_clone,
                         remote_clone,
                         node_name,
+                        client_factory,
                         permit,
                         Some(per_remote_connections_permit),
                     );
@@ -441,6 +464,7 @@ impl<C: Clone + Send + 'static> ParallelFetcher<C> {
                     context,
                     remote.clone(),
                     "localhost".into(),
+                    client_factory,
                     permit.unwrap(), // Always set to `Some` at this point
                     None,
                 )
@@ -464,6 +488,7 @@ impl<C: Clone + Send + 'static> ParallelFetcher<C> {
         context: C,
         remote: Remote,
         node: String,
+        client_factory: Arc<dyn ClientFactory + Send + Sync>,
         _permit: OwnedSemaphorePermit,
         _per_remote_connections_permit: Option<OwnedSemaphorePermit>,
     ) -> NodeResponse<T>
@@ -478,6 +503,7 @@ impl<C: Clone + Send + 'static> ParallelFetcher<C> {
             context,
             remote,
             node: node.clone(),
+            client_factory,
         };
 
         let result = func(parallel_fetcher_context).await;
@@ -514,6 +540,9 @@ impl<C: Clone + Send + 'static> ParallelFetcher<C> {
 
             let context = self.context.clone();
             let func = func.clone();
+
+            let client_factory = Arc::clone(&self.client_factory);
+
             let future = async move {
                 let permit = total_connections_semaphore.acquire_owned().await.unwrap();
 
@@ -525,6 +554,7 @@ impl<C: Clone + Send + 'static> ParallelFetcher<C> {
                         context,
                         remote,
                         "localhost".into(),
+                        client_factory,
                         permit,
                         None,
                     )
-- 
2.47.3