From a5d17b7723d0b002d68681b0d2fae32fec0783cd Mon Sep 17 00:00:00 2001 From: namse Date: Sat, 2 Dec 2023 15:54:23 +0000 Subject: [PATCH] I am ready to go!/ --- oioi/agent/Cargo.lock | 2 + oioi/agent/Cargo.toml | 2 + oioi/agent/src/docker_cli.rs | 30 +++++ oioi/agent/src/docker_engine.rs | 151 +++++++++++++++++++++ oioi/agent/src/envs.rs | 29 ++++ oioi/agent/src/main.rs | 225 +++++++------------------------- 6 files changed, 260 insertions(+), 179 deletions(-) create mode 100644 oioi/agent/src/docker_cli.rs create mode 100644 oioi/agent/src/docker_engine.rs create mode 100644 oioi/agent/src/envs.rs diff --git a/oioi/agent/Cargo.lock b/oioi/agent/Cargo.lock index 8b7b2db0c..f943d381b 100644 --- a/oioi/agent/Cargo.lock +++ b/oioi/agent/Cargo.lock @@ -989,6 +989,8 @@ dependencies = [ "bollard", "futures-util", "lazy_static", + "serde", + "serde_json", "tokio", ] diff --git a/oioi/agent/Cargo.toml b/oioi/agent/Cargo.toml index e2748a6e3..ddff09743 100644 --- a/oioi/agent/Cargo.toml +++ b/oioi/agent/Cargo.toml @@ -12,4 +12,6 @@ aws-sdk-ssm = "1.3.0" bollard = "0.15.0" futures-util = "0.3.29" lazy_static = "1.4.0" +serde = { version = "1.0.193", features = ["derive"] } +serde_json = "1.0.108" tokio = { version = "1.34.0", features = ["rt", "macros"] } diff --git a/oioi/agent/src/docker_cli.rs b/oioi/agent/src/docker_cli.rs new file mode 100644 index 000000000..c96b7e5e6 --- /dev/null +++ b/oioi/agent/src/docker_cli.rs @@ -0,0 +1,30 @@ +use anyhow::Result; + +pub(crate) async fn pull_image(image: &str) -> Result<()> { + run_docker_command(&["pull", image]).await?; + Ok(()) +} + +async fn run_docker_command(args: &[&str]) -> Result> { + let output = tokio::process::Command::new("docker") + .args(args) + .output() + .await + .map_err(|e| { + anyhow::anyhow!( + "Failed to run docker {args}: {e}", + args = args.join(" "), + e = e + ) + })?; + + if !output.status.success() { + return Err(anyhow::anyhow!( + "Failed to run docker {args}: {stderr}", + args = args.join(" "), + stderr = String::from_utf8_lossy(&output.stderr), + )); + } + + Ok(output.stdout) +} diff --git a/oioi/agent/src/docker_engine.rs b/oioi/agent/src/docker_engine.rs new file mode 100644 index 000000000..47a726537 --- /dev/null +++ b/oioi/agent/src/docker_engine.rs @@ -0,0 +1,151 @@ +use crate::envs::*; +use anyhow::Result; +use bollard::{container::ListContainersOptions, Docker}; + +pub(crate) struct DockerEngine { + docker: Docker, +} + +impl DockerEngine { + pub(crate) fn new() -> Result { + Ok(Self { + docker: Docker::connect_with_local_defaults()?, + }) + } + + pub(crate) async fn get_local_image_digest(&self, image: &str) -> Result { + let output = self.docker.inspect_image(image).await?; + + Ok(output.id.unwrap()) + } + + pub(crate) async fn get_running_container_image_digest( + &self, + container_name: &str, + ) -> Result> { + Ok(self + .docker + .list_containers(Some(ListContainersOptions:: { + all: true, + ..Default::default() + })) + .await? + .into_iter() + .find_map(|container| { + container + .names + .unwrap_or_default() + .contains(&format!("/{container_name}")) + .then(|| container.image_id.unwrap()) + })) + } + + pub(crate) async fn stop_running_container(&self) -> Result<()> { + self.docker + .stop_container( + CONTAINER_NAME, + Some(bollard::container::StopContainerOptions { + t: GRACEFUL_SHUTDOWN_TIMEOUT_SECS, + }), + ) + .await?; + Ok(()) + } + + pub(crate) async fn run_new_container(&self, image: &str) -> Result<()> { + println!("removing container {CONTAINER_NAME}"); + let remove_container_result = self + .docker + .remove_container( + CONTAINER_NAME, + Some(bollard::container::RemoveContainerOptions { + force: true, + ..Default::default() + }), + ) + .await; + + if let Err(err) = remove_container_result { + match err { + bollard::errors::Error::DockerResponseServerError { + status_code, + message, + } => if status_code == 404 && message.contains("No such container") {}, + _ => return Err(err.into()), + } + }; + + println!("creating container {CONTAINER_NAME}"); + self.docker + .create_container( + Some(bollard::container::CreateContainerOptions { + name: CONTAINER_NAME, + platform: None, + }), + bollard::container::Config { + image: Some(image.to_string()), + host_config: Some(bollard::models::HostConfig { + log_config: Some(bollard::models::HostConfigLogConfig { + typ: Some("awslogs".to_string()), + config: Some(std::collections::HashMap::from_iter([ + ("awslogs-group".to_string(), format!("oioi-{}", *GROUP_NAME)), + ( + "awslogs-stream".to_string(), + format!("oioi-{}-{}", *GROUP_NAME, *EC2_INSTANCE_ID), + ), + ("awslogs-create-group".to_string(), "true".to_string()), + ])), + }), + port_bindings: Some(std::collections::HashMap::from_iter( + PORT_MAPPINGS + .iter() + .map(|mapping| { + println!("mapping: {:?}", mapping); + ( + format!("{}/{}", mapping.container_port, mapping.protocol), + Some(vec![bollard::models::PortBinding { + host_ip: None, + host_port: Some(mapping.host_port.to_string()), + }]), + ) + }) + .collect::>(), + )), + ..Default::default() + }), + ..Default::default() + }, + ) + .await?; + + println!("starting container {CONTAINER_NAME}"); + self.docker + .start_container( + CONTAINER_NAME, + None::>, + ) + .await?; + + Ok(()) + } + + pub(crate) async fn docker_prune(&self) -> Result<()> { + self.docker + .prune_containers(None::>) + .await?; + + self.docker + .prune_images(None::>) + .await?; + + self.docker + .prune_networks(None::>) + .await?; + + self.docker + .prune_volumes(None::>) + .await?; + + Ok(()) + } +} diff --git a/oioi/agent/src/envs.rs b/oioi/agent/src/envs.rs new file mode 100644 index 000000000..548dfcf1f --- /dev/null +++ b/oioi/agent/src/envs.rs @@ -0,0 +1,29 @@ +pub(crate) const INTERVAL: std::time::Duration = std::time::Duration::from_secs(5); +pub(crate) const CONTAINER_NAME: &str = "oioi"; +pub(crate) const GRACEFUL_SHUTDOWN_TIMEOUT_SECS: i64 = 30; + +lazy_static::lazy_static! { + pub(crate) static ref GROUP_NAME: String = std::env::var("GROUP_NAME").expect("GROUP_NAME env var not set"); + pub(crate) static ref EC2_INSTANCE_ID: String = std::env::var("EC2_INSTANCE_ID").expect("EC2_INSTANCE_ID env var not set"); + pub(crate) static ref PORT_MAPPINGS: Vec = std::env::var("PORT_MAPPINGS").map(|env_string| { + env_string.split(',').map(|mapping| { + let mut parts = mapping.split(&[':', '/']); + let container_port = parts.next().expect("container port not found").parse::().expect("container port is not a number"); + let host_port = parts.next().expect("host port not found").parse::().expect("host port is not a number"); + let protocol = parts.next().expect("protocol not found").to_string(); + + PortMapping { + container_port, + host_port, + protocol, + } + }).collect() + }).expect("PORT_MAPPINGS env var not set"); +} + +#[derive(Debug)] +pub(crate) struct PortMapping { + pub(crate) container_port: u16, + pub(crate) host_port: u16, + pub(crate) protocol: String, +} diff --git a/oioi/agent/src/main.rs b/oioi/agent/src/main.rs index 0cbfa8d04..d55648545 100644 --- a/oioi/agent/src/main.rs +++ b/oioi/agent/src/main.rs @@ -1,217 +1,84 @@ +mod docker_cli; +mod docker_engine; +mod envs; + use anyhow::Result; -use bollard::{container::ListContainersOptions, Docker}; -use futures_util::TryStreamExt; +use envs::*; #[tokio::main] async fn main() -> Result<()> { real_main().await } -const INTERVAL: std::time::Duration = std::time::Duration::from_secs(5); -const CONTAINER_NAME: &str = "oioi"; -const GRACEFUL_SHUTDOWN_TIMEOUT_SECS: i64 = 30; -lazy_static::lazy_static! { - static ref GROUP_NAME: String = std::env::var("GROUP_NAME").expect("GROUP_NAME env var not set"); - static ref EC2_INSTANCE_ID: String = std::env::var("EC2_INSTANCE_ID").expect("EC2_INSTANCE_ID env var not set"); - static ref PORT_MAPPINGS: Vec = std::env::var("PORT_MAPPINGS").map(|env_string| { - env_string.split(',').map(|mapping| { - let mut parts = mapping.split(&[':', '/']); - let container_port = parts.next().expect("container port not found").parse::().expect("container port is not a number"); - let host_port = parts.next().expect("host port not found").parse::().expect("host port is not a number"); - let protocol = parts.next().expect("protocol not found").to_string(); - - PortMapping { - container_port, - host_port, - protocol, - } - }).collect() - }).expect("PORT_MAPPINGS env var not set"); -} - async fn real_main() -> Result<()> { println!("Environment variables: {:?}", std::env::vars()); - let config = aws_config::load_defaults(aws_config::BehaviorVersion::v2023_11_09()).await; - let aws_ssm_client = aws_sdk_ssm::Client::new(&config); - let docker = Docker::connect_with_local_defaults()?; + let docker_engine = docker_engine::DockerEngine::new()?; + let mut running_image_digest_cache = None; loop { - let Some(image) = get_image(&aws_ssm_client).await? else { + let Some(image) = get_target_image().await? else { println!("No image found for group {}.", *GROUP_NAME); tokio::time::sleep(INTERVAL).await; continue; }; - let running_image = get_running_image(&docker).await?; + docker_cli::pull_image(&image).await?; + + let target_image_digest = docker_engine.get_local_image_digest(&image).await?; + + let running_image_digest = { + if running_image_digest_cache.is_none() { + running_image_digest_cache = docker_engine + .get_running_container_image_digest(CONTAINER_NAME) + .await?; + } + &running_image_digest_cache + }; + + println!("Target image: {:?}", image); + println!("Target image digest: {:?}", target_image_digest); + println!("Running image digest: {:?}", running_image_digest); - if let Some(running_image) = running_image { - if running_image == image { + if let Some(running_image_digest) = running_image_digest { + if running_image_digest == &target_image_digest { println!("Good! Image {} is already running.", image); tokio::time::sleep(INTERVAL).await; continue; } - stop_running_container(&docker).await?; + + docker_engine.stop_running_container().await?; } - run_new_container(&docker, &image).await?; + docker_engine.run_new_container(&image).await?; - docker_prune(&docker).await?; + running_image_digest_cache = Some(target_image_digest); + + docker_engine.docker_prune().await?; tokio::time::sleep(INTERVAL).await; } } -async fn get_image(aws_ssm_client: &aws_sdk_ssm::Client) -> Result> { +async fn get_target_image() -> Result> { + static AWS_SSM_CLIENT: tokio::sync::OnceCell = + tokio::sync::OnceCell::const_new(); + + let aws_ssm_client = AWS_SSM_CLIENT + .get_or_init(|| async { + let config = + aws_config::load_defaults(aws_config::BehaviorVersion::v2023_11_09()).await; + aws_sdk_ssm::Client::new(&config) + }) + .await; + let parameter_path = format!("/oioi/{}/image", *GROUP_NAME); - let image = aws_ssm_client + Ok(aws_ssm_client .get_parameter() .name(¶meter_path) .send() .await? .parameter - .and_then(|p| p.value); - - Ok(image) -} - -async fn get_running_image(docker: &Docker) -> Result> { - let containers = docker - .list_containers(None::>) - .await?; - - for container in containers { - let Some(names) = container.names else { - continue; - }; - - if !names.contains(&CONTAINER_NAME.to_string()) { - continue; - } - - let Some(image) = container.image else { - anyhow::bail!("Container image should not be empty"); - }; - - return Ok(Some(image)); - } - - Ok(None) -} - -async fn stop_running_container(docker: &Docker) -> Result<()> { - docker - .stop_container( - CONTAINER_NAME, - Some(bollard::container::StopContainerOptions { - t: GRACEFUL_SHUTDOWN_TIMEOUT_SECS, - }), - ) - .await?; - - docker - .remove_container( - CONTAINER_NAME, - Some(bollard::container::RemoveContainerOptions { - force: true, - ..Default::default() - }), - ) - .await?; - - Ok(()) -} - -async fn run_new_container(docker: &Docker, image: &str) -> Result<()> { - println!("creating image {}", image); - docker - .create_image( - Some(bollard::image::CreateImageOptions { - from_image: image, - ..Default::default() - }), - None, - None, - ) - .try_collect::>() - .await?; - - println!("creating container {}", CONTAINER_NAME); - docker - .create_container( - Some(bollard::container::CreateContainerOptions { - name: CONTAINER_NAME, - platform: None, - }), - bollard::container::Config { - image: Some(image.to_string()), - host_config: Some(bollard::models::HostConfig { - log_config: Some(bollard::models::HostConfigLogConfig { - typ: Some("awslogs".to_string()), - config: Some(std::collections::HashMap::from_iter([ - ("awslogs-group".to_string(), format!("oioi-{}", *GROUP_NAME)), - ( - "awslogs-stream".to_string(), - format!("oioi-{}-{}", *GROUP_NAME, *EC2_INSTANCE_ID), - ), - ("awslogs-create-group".to_string(), "true".to_string()), - ])), - }), - port_bindings: Some(std::collections::HashMap::from_iter( - PORT_MAPPINGS - .iter() - .map(|mapping| { - ( - format!("{}/{}", mapping.container_port, mapping.protocol), - Some(vec![bollard::models::PortBinding { - host_ip: None, - host_port: Some(mapping.host_port.to_string()), - }]), - ) - }) - .collect::>(), - )), - ..Default::default() - }), - ..Default::default() - }, - ) - .await?; - - println!("starting container {}", CONTAINER_NAME); - docker - .start_container( - CONTAINER_NAME, - None::>, - ) - .await?; - - Ok(()) -} - -async fn docker_prune(docker: &Docker) -> Result<()> { - docker - .prune_containers(None::>) - .await?; - - docker - .prune_images(None::>) - .await?; - - docker - .prune_networks(None::>) - .await?; - - docker - .prune_volumes(None::>) - .await?; - - Ok(()) -} - -struct PortMapping { - container_port: u16, - host_port: u16, - protocol: String, + .and_then(|p| p.value)) }