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