use std::{
collections::HashMap,
path::{Path, PathBuf},
sync::{Arc, Weak},
thread,
};
use anyhow::anyhow;
use async_trait::async_trait;
use support::{constants::THIS_IS_A_BUG, fs::FileSystem};
use tokio::sync::{Mutex, RwLock};
use tracing::{debug, trace, warn};
use uuid::Uuid;
use super::{
client::{ContainerRunOptions, DockerClient},
node::DockerNode,
DockerProvider,
};
use crate::{
constants::NAMESPACE_PREFIX,
docker::node::DockerNodeOptions,
types::{
GenerateFileCommand, GenerateFilesOptions, ProviderCapabilities, RunCommandOptions,
SpawnNodeOptions,
},
DynNode, ProviderError, ProviderNamespace, ProviderNode,
};
pub struct DockerNamespace<FS>
where
FS: FileSystem + Send + Sync + Clone,
{
weak: Weak<DockerNamespace<FS>>,
#[allow(dead_code)]
provider: Weak<DockerProvider<FS>>,
name: String,
base_dir: PathBuf,
capabilities: ProviderCapabilities,
docker_client: DockerClient,
filesystem: FS,
delete_on_drop: Arc<Mutex<bool>>,
pub(super) nodes: RwLock<HashMap<String, Arc<DockerNode<FS>>>>,
}
impl<FS> DockerNamespace<FS>
where
FS: FileSystem + Send + Sync + Clone + 'static,
{
pub(super) async fn new(
provider: &Weak<DockerProvider<FS>>,
tmp_dir: &PathBuf,
capabilities: &ProviderCapabilities,
docker_client: &DockerClient,
filesystem: &FS,
custom_base_dir: Option<&Path>,
) -> Result<Arc<Self>, ProviderError> {
let name = format!("{}{}", NAMESPACE_PREFIX, Uuid::new_v4());
let base_dir = if let Some(custom_base_dir) = custom_base_dir {
if !filesystem.exists(custom_base_dir).await {
filesystem.create_dir(custom_base_dir).await?;
} else {
warn!(
"⚠️ Using and existing directory {} as base dir",
custom_base_dir.to_string_lossy()
);
}
PathBuf::from(custom_base_dir)
} else {
let base_dir = PathBuf::from_iter([tmp_dir, &PathBuf::from(&name)]);
filesystem.create_dir(&base_dir).await?;
base_dir
};
let namespace = Arc::new_cyclic(|weak| DockerNamespace {
weak: weak.clone(),
provider: provider.clone(),
name,
base_dir,
capabilities: capabilities.clone(),
filesystem: filesystem.clone(),
docker_client: docker_client.clone(),
nodes: RwLock::new(HashMap::new()),
delete_on_drop: Arc::new(Mutex::new(true)),
});
namespace.initialize().await?;
Ok(namespace)
}
async fn initialize(&self) -> Result<(), ProviderError> {
self.initialize_zombie_scripts_volume().await?;
self.initialize_helper_binaries_volume().await?;
Ok(())
}
async fn initialize_zombie_scripts_volume(&self) -> Result<(), ProviderError> {
let local_zombie_wrapper_path =
PathBuf::from_iter([&self.base_dir, &PathBuf::from("zombie-wrapper.sh")]);
self.filesystem
.write(
&local_zombie_wrapper_path,
include_str!("../shared/scripts/zombie-wrapper.sh"),
)
.await?;
let local_helper_binaries_downloader_path = PathBuf::from_iter([
&self.base_dir,
&PathBuf::from("helper-binaries-downloader.sh"),
]);
self.filesystem
.write(
&local_helper_binaries_downloader_path,
include_str!("../shared/scripts/helper-binaries-downloader.sh"),
)
.await?;
let zombie_wrapper_volume_name = format!("{}-zombie-wrapper", self.name);
let zombie_wrapper_container_name = format!("{}-scripts", self.name);
self.docker_client
.create_volume(&zombie_wrapper_volume_name)
.await
.map_err(|err| ProviderError::CreateNamespaceFailed(self.name.clone(), err.into()))?;
self.docker_client
.container_create(
ContainerRunOptions::new("alpine:latest", vec!["tail", "-f", "/dev/null"])
.volume_mounts(HashMap::from([(
zombie_wrapper_volume_name.as_str(),
"/scripts",
)]))
.name(&zombie_wrapper_container_name)
.rm(),
)
.await
.map_err(|err| ProviderError::CreateNamespaceFailed(self.name.clone(), err.into()))?;
self.docker_client
.container_cp(
&zombie_wrapper_container_name,
&local_zombie_wrapper_path,
&PathBuf::from("/scripts/zombie-wrapper.sh"),
)
.await
.map_err(|err| ProviderError::CreateNamespaceFailed(self.name.clone(), err.into()))?;
self.docker_client
.container_cp(
&zombie_wrapper_container_name,
&local_helper_binaries_downloader_path,
&PathBuf::from("/scripts/helper-binaries-downloader.sh"),
)
.await
.map_err(|err| ProviderError::CreateNamespaceFailed(self.name.clone(), err.into()))?;
self.docker_client
.container_run(
ContainerRunOptions::new("alpine:latest", vec!["chmod", "-R", "777", "/scripts"])
.volume_mounts(HashMap::from([(
zombie_wrapper_volume_name.as_ref(),
"/scripts",
)]))
.rm(),
)
.await
.map_err(|err| ProviderError::CreateNamespaceFailed(self.name.clone(), err.into()))?;
Ok(())
}
async fn initialize_helper_binaries_volume(&self) -> Result<(), ProviderError> {
let helper_binaries_volume_name = format!("{}-helper-binaries", self.name);
let zombie_wrapper_volume_name = format!("{}-zombie-wrapper", self.name);
self.docker_client
.create_volume(&helper_binaries_volume_name)
.await
.map_err(|err| ProviderError::CreateNamespaceFailed(self.name.clone(), err.into()))?;
self.docker_client
.container_run(
ContainerRunOptions::new(
"alpine:latest",
vec!["ash", "/scripts/helper-binaries-downloader.sh"],
)
.volume_mounts(HashMap::from([
(
helper_binaries_volume_name.as_str(),
"/helpers",
),
(
zombie_wrapper_volume_name.as_ref(),
"/scripts",
)
]))
.detach(false)
.rm(),
)
.await
.map_err(|err| ProviderError::CreateNamespaceFailed(self.name.clone(), err.into()))?;
self.docker_client
.container_run(
ContainerRunOptions::new("alpine:latest", vec!["chmod", "-R", "777", "/helpers"])
.volume_mounts(HashMap::from([(
helper_binaries_volume_name.as_ref(),
"/helpers",
)]))
.rm(),
)
.await
.map_err(|err| ProviderError::CreateNamespaceFailed(self.name.clone(), err.into()))?;
Ok(())
}
pub async fn set_delete_on_drop(&self, delete_on_drop: bool) {
*self.delete_on_drop.lock().await = delete_on_drop;
}
pub async fn delete_on_drop(&self) -> bool {
if let Ok(delete_on_drop) = self.delete_on_drop.try_lock() {
*delete_on_drop
} else {
true
}
}
}
#[async_trait]
impl<FS> ProviderNamespace for DockerNamespace<FS>
where
FS: FileSystem + Send + Sync + Clone + 'static,
{
fn name(&self) -> &str {
&self.name
}
fn base_dir(&self) -> &PathBuf {
&self.base_dir
}
fn capabilities(&self) -> &ProviderCapabilities {
&self.capabilities
}
async fn detach(&self) {
self.set_delete_on_drop(false).await;
}
async fn is_detached(&self) -> bool {
self.delete_on_drop().await
}
async fn nodes(&self) -> HashMap<String, DynNode> {
self.nodes
.read()
.await
.iter()
.map(|(name, node)| (name.clone(), node.clone() as DynNode))
.collect()
}
async fn get_node_available_args(
&self,
(command, image): (String, Option<String>),
) -> Result<String, ProviderError> {
let node_image = image.expect(&format!("image should be present when getting node available args with docker provider {THIS_IS_A_BUG}"));
let temp_node = self
.spawn_node(
&SpawnNodeOptions::new(format!("temp-{}", Uuid::new_v4()), "cat".to_string())
.image(node_image.clone()),
)
.await?;
let available_args_output = temp_node
.run_command(RunCommandOptions::new(command.clone()).args(vec!["--help"]))
.await?
.map_err(|(_exit, status)| {
ProviderError::NodeAvailableArgsError(node_image, command, status)
})?;
temp_node.destroy().await?;
Ok(available_args_output)
}
async fn spawn_node(&self, options: &SpawnNodeOptions) -> Result<DynNode, ProviderError> {
debug!("spawn option {:?}", options);
if self.nodes.read().await.contains_key(&options.name) {
return Err(ProviderError::DuplicatedNodeName(options.name.clone()));
}
let node = DockerNode::new(DockerNodeOptions {
namespace: &self.weak,
namespace_base_dir: &self.base_dir,
name: &options.name,
image: options.image.as_ref(),
program: &options.program,
args: &options.args,
env: &options.env,
startup_files: &options.injected_files,
db_snapshot: options.db_snapshot.as_ref(),
docker_client: &self.docker_client,
container_name: format!("{}-{}", self.name, options.name),
filesystem: &self.filesystem,
port_mapping: options.port_mapping.as_ref().unwrap_or(&HashMap::default()),
})
.await?;
self.nodes
.write()
.await
.insert(node.name().to_string(), node.clone());
Ok(node)
}
async fn generate_files(&self, options: GenerateFilesOptions) -> Result<(), ProviderError> {
debug!("generate files options {options:#?}");
let node_name = options
.temp_name
.unwrap_or_else(|| format!("temp-{}", Uuid::new_v4()));
let node_image = options.image.expect(&format!(
"image should be present when generating files with docker provider {THIS_IS_A_BUG}"
));
let temp_node = self
.spawn_node(
&SpawnNodeOptions::new(node_name, "cat".to_string())
.injected_files(options.injected_files)
.image(node_image),
)
.await?;
for GenerateFileCommand {
program,
args,
env,
local_output_path,
} in options.commands
{
let local_output_full_path = format!(
"{}{}{}",
self.base_dir.to_string_lossy(),
if local_output_path.starts_with("/") {
""
} else {
"/"
},
local_output_path.to_string_lossy()
);
match temp_node
.run_command(RunCommandOptions { program, args, env })
.await?
{
Ok(contents) => self
.filesystem
.write(local_output_full_path, contents)
.await
.map_err(|err| ProviderError::FileGenerationFailed(err.into()))?,
Err((_, msg)) => Err(ProviderError::FileGenerationFailed(anyhow!("{msg}")))?,
};
}
temp_node.destroy().await
}
async fn static_setup(&self) -> Result<(), ProviderError> {
todo!()
}
async fn destroy(&self) -> Result<(), ProviderError> {
let _ = self
.docker_client
.namespaced_containers_rm(&self.name)
.await
.map_err(|err| ProviderError::DeleteNamespaceFailed(self.name.clone(), err.into()))?;
if let Some(provider) = self.provider.upgrade() {
provider.namespaces.write().await.remove(&self.name);
}
Ok(())
}
}
impl<FS> Drop for DockerNamespace<FS>
where
FS: FileSystem + Send + Sync + Clone,
{
fn drop(&mut self) {
let ns_name = self.name.clone();
if let Ok(delete_on_drop) = self.delete_on_drop.try_lock() {
if *delete_on_drop {
let client = self.docker_client.clone();
let provider = self.provider.upgrade();
let handler = thread::spawn(move || {
let rt = tokio::runtime::Runtime::new().unwrap();
rt.block_on(async move {
trace!("🧟 deleting ns {ns_name} from cluster");
let _ = client.namespaced_containers_rm(&ns_name).await;
trace!("✅ deleted");
});
});
if handler.join().is_ok() {
if let Some(provider) = provider {
if let Ok(mut p) = provider.namespaces.try_write() {
p.remove(&self.name);
} else {
warn!(
"⚠️ Can not acquire write lock to the provider, ns {} not removed",
self.name
);
}
}
}
} else {
trace!("⚠️ leaking ns {ns_name} in cluster");
}
};
}
}