zombienet_provider/docker/
provider.rs

1use std::{
2    collections::HashMap,
3    path::{Path, PathBuf},
4    sync::{Arc, Weak},
5};
6
7use async_trait::async_trait;
8use support::fs::FileSystem;
9use tokio::sync::RwLock;
10
11use super::{client::DockerClient, namespace::DockerNamespace};
12use crate::{
13    types::ProviderCapabilities, DynNamespace, Provider, ProviderError, ProviderNamespace,
14};
15
16const PROVIDER_NAME: &str = "docker";
17
18pub struct DockerProvider<FS>
19where
20    FS: FileSystem + Send + Sync + Clone,
21{
22    weak: Weak<DockerProvider<FS>>,
23    capabilities: ProviderCapabilities,
24    tmp_dir: PathBuf,
25    docker_client: DockerClient,
26    filesystem: FS,
27    pub(super) namespaces: RwLock<HashMap<String, Arc<DockerNamespace<FS>>>>,
28}
29
30impl<FS> DockerProvider<FS>
31where
32    FS: FileSystem + Send + Sync + Clone + 'static,
33{
34    pub async fn new(filesystem: FS) -> Arc<Self> {
35        let docker_client = DockerClient::new().await.unwrap();
36
37        let provider = Arc::new_cyclic(|weak| DockerProvider {
38            weak: weak.clone(),
39            capabilities: ProviderCapabilities {
40                requires_image: true,
41                has_resources: false,
42                prefix_with_full_path: false,
43                use_default_ports_in_cmd: true,
44            },
45            tmp_dir: std::env::temp_dir(),
46            docker_client,
47            filesystem,
48            namespaces: RwLock::new(HashMap::new()),
49        });
50
51        let cloned_provider = provider.clone();
52        tokio::spawn(async move {
53            tokio::signal::ctrl_c().await.unwrap();
54            for (_, ns) in cloned_provider.namespaces().await {
55                if ns.is_detached().await {
56                    // best effort
57                    let _ = ns.destroy().await;
58                }
59            }
60
61            // exit the process (130, SIGINT)
62            std::process::exit(130)
63        });
64
65        provider
66    }
67
68    pub fn tmp_dir(mut self, tmp_dir: impl Into<PathBuf>) -> Self {
69        self.tmp_dir = tmp_dir.into();
70        self
71    }
72}
73
74#[async_trait]
75impl<FS> Provider for DockerProvider<FS>
76where
77    FS: FileSystem + Send + Sync + Clone + 'static,
78{
79    fn name(&self) -> &str {
80        PROVIDER_NAME
81    }
82
83    fn capabilities(&self) -> &ProviderCapabilities {
84        &self.capabilities
85    }
86
87    async fn namespaces(&self) -> HashMap<String, DynNamespace> {
88        self.namespaces
89            .read()
90            .await
91            .iter()
92            .map(|(name, namespace)| (name.clone(), namespace.clone() as DynNamespace))
93            .collect()
94    }
95
96    async fn create_namespace(&self) -> Result<DynNamespace, ProviderError> {
97        let namespace = DockerNamespace::new(
98            &self.weak,
99            &self.tmp_dir,
100            &self.capabilities,
101            &self.docker_client,
102            &self.filesystem,
103            None,
104        )
105        .await?;
106
107        self.namespaces
108            .write()
109            .await
110            .insert(namespace.name().to_string(), namespace.clone());
111
112        Ok(namespace)
113    }
114
115    async fn create_namespace_with_base_dir(
116        &self,
117        base_dir: &Path,
118    ) -> Result<DynNamespace, ProviderError> {
119        let namespace = DockerNamespace::new(
120            &self.weak,
121            &self.tmp_dir,
122            &self.capabilities,
123            &self.docker_client,
124            &self.filesystem,
125            Some(base_dir),
126        )
127        .await?;
128
129        self.namespaces
130            .write()
131            .await
132            .insert(namespace.name().to_string(), namespace.clone());
133
134        Ok(namespace)
135    }
136}