From 433370ac3c222fe3007abc8ab8dc36ebaec51001 Mon Sep 17 00:00:00 2001 From: Filippo Casarin Date: Wed, 9 Oct 2024 22:10:29 +0200 Subject: [PATCH] Reload hostmap on sighup --- pixie-server/Cargo.lock | 10 ++++++++++ pixie-server/Cargo.toml | 2 +- pixie-server/src/dnsmasq.rs | 34 +++++++++++++++++++++++----------- pixie-server/src/http.rs | 26 +++++++++++--------------- pixie-server/src/main.rs | 9 +++++++++ pixie-server/src/state.rs | 18 +++++++++++++++++- pixie.service | 1 + 7 files changed, 72 insertions(+), 28 deletions(-) diff --git a/pixie-server/Cargo.lock b/pixie-server/Cargo.lock index 4fbfdef9..bec1abb3 100644 --- a/pixie-server/Cargo.lock +++ b/pixie-server/Cargo.lock @@ -1187,6 +1187,15 @@ dependencies = [ "digest", ] +[[package]] +name = "signal-hook-registry" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1" +dependencies = [ + "libc", +] + [[package]] name = "slab" version = "0.4.9" @@ -1297,6 +1306,7 @@ dependencies = [ "mio", "num_cpus", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.48.0", diff --git a/pixie-server/Cargo.toml b/pixie-server/Cargo.toml index bb148dee..e98a6e1e 100644 --- a/pixie-server/Cargo.toml +++ b/pixie-server/Cargo.toml @@ -18,7 +18,7 @@ postcard = { version = "1.0.8", default-features = false, features = ["alloc"] } serde = "1.0.193" serde_derive = "1.0.193" serde_yaml = "0.9" -tokio = { version = "1.34.0", features = ["macros", "fs", "rt-multi-thread", "sync"] } +tokio = { version = "1.34.0", features = ["macros", "fs", "rt-multi-thread", "sync", "signal"] } serde_json = "1.0.108" hostfile = "0.2.0" hex = "0.4.3" diff --git a/pixie-server/src/dnsmasq.rs b/pixie-server/src/dnsmasq.rs index 950149f1..843d748a 100644 --- a/pixie-server/src/dnsmasq.rs +++ b/pixie-server/src/dnsmasq.rs @@ -1,4 +1,5 @@ use std::{ + collections::HashMap, fs::File, io::{BufWriter, Error, Write}, net::Ipv4Addr, @@ -99,13 +100,16 @@ async fn write_hosts(state: &State, hosts: &[(MacAddr6, Ipv4Addr, Option Ok(()) } -fn get_hosts(state: &State, units: &[Unit]) -> Vec<(MacAddr6, Ipv4Addr, Option)> { +fn get_hosts( + hostmap: &HashMap, + units: &[Unit], +) -> Vec<(MacAddr6, Ipv4Addr, Option)> { units .iter() .map(|unit| { let mac = unit.mac; let ip = unit.static_ip(); - let hostname = state.hostmap.get(&ip).cloned(); + let hostname = hostmap.get(&ip).cloned(); (mac, ip, hostname) }) .collect() @@ -113,9 +117,13 @@ fn get_hosts(state: &State, units: &[Unit]) -> Vec<(MacAddr6, Ipv4Addr, Option) -> Result<()> { let mut units_rx = state.units.subscribe(); + let mut hostmap_rx = state.hostmap.subscribe(); write_config(&state).await?; - let mut hosts = get_hosts(&state, &units_rx.borrow_and_update()); + let mut hosts = get_hosts( + &hostmap_rx.borrow_and_update(), + &units_rx.borrow_and_update(), + ); write_hosts(&state, &hosts).await?; let dnsmasq = DnsmasqHandle { @@ -132,14 +140,18 @@ pub async fn main(state: Arc) -> Result<()> { loop { tokio::select! { - _ = units_rx.changed() => { - let hosts2 = get_hosts(&state, &units_rx.borrow_and_update()); - if hosts != hosts2 { - hosts = hosts2; - write_hosts(&state, &hosts).await?; - dnsmasq.reload()?; - } - } + ret = units_rx.changed() => ret.unwrap(), + ret = hostmap_rx.changed() => ret.unwrap(), + } + + let hosts2 = get_hosts( + &hostmap_rx.borrow_and_update(), + &units_rx.borrow_and_update(), + ); + if hosts != hosts2 { + hosts = hosts2; + write_hosts(&state, &hosts).await?; + dnsmasq.reload()?; } } } diff --git a/pixie-server/src/http.rs b/pixie-server/src/http.rs index c3d993c0..8097b239 100644 --- a/pixie-server/src/http.rs +++ b/pixie-server/src/http.rs @@ -161,23 +161,19 @@ async fn gc(extract::State(state): extract::State>) -> impl IntoRespo } async fn status(extract::State(state): extract::State>) -> impl IntoResponse { - let initial_messages = vec![ - StatusUpdate::Config(state.config.clone()), - StatusUpdate::HostMap(state.hostmap.clone()), - ]; - let mut units_rx = state.units.subscribe(); - units_rx.mark_changed(); - let units_rx = WatchStream::new(units_rx); - - let mut image_rx = state.image_stats.subscribe(); - image_rx.mark_changed(); - let image_rx = WatchStream::new(image_rx); - - let messages = - futures::stream::iter(initial_messages.into_iter()).chain(futures::stream::select( + let initial_messages = [StatusUpdate::Config(state.config.clone())]; + + let units_rx = WatchStream::new(state.units.subscribe()); + let image_rx = WatchStream::new(state.image_stats.subscribe()); + let hostmap_rx = WatchStream::new(state.hostmap.subscribe()); + + let messages = futures::stream::iter(initial_messages).chain(futures::stream::select( + futures::stream::select( image_rx.map(StatusUpdate::ImageStats), units_rx.map(StatusUpdate::Units), - )); + ), + hostmap_rx.map(StatusUpdate::HostMap), + )); let lines = messages.map(|msg| serde_json::to_string(&msg).map(|x| x + "\n")); let mut res = Response::new(Body::from_stream(lines)); diff --git a/pixie-server/src/main.rs b/pixie-server/src/main.rs index c28d6635..fad67520 100644 --- a/pixie-server/src/main.rs +++ b/pixie-server/src/main.rs @@ -127,6 +127,15 @@ async fn main() -> Result<()> { let state = Arc::new(State::load(options.storage_dir)?); + let state2 = state.clone(); + tokio::spawn(async move { + let mut signal = tokio::signal::unix::signal(tokio::signal::unix::SignalKind::hangup()) + .expect("failed to register signal handler"); + while let Some(()) = signal.recv().await { + state2.reload().unwrap(); + } + }); + async fn flatten(task: JoinHandle>) -> Result<()> { task.await??; Ok(()) diff --git a/pixie-server/src/state.rs b/pixie-server/src/state.rs index 17e592de..5c4d86fe 100644 --- a/pixie-server/src/state.rs +++ b/pixie-server/src/state.rs @@ -35,7 +35,7 @@ fn get_image_csize(image: &Image) -> u64 { pub struct State { pub storage_dir: PathBuf, pub config: Config, - pub hostmap: HashMap, + pub hostmap: watch::Sender>, pub units: watch::Sender>, // TODO: use an Option @@ -64,6 +64,7 @@ impl State { } } } + let hostmap = watch::Sender::new(hostmap); let units_path = storage_dir.join("registered.json"); let units = watch::Sender::new({ @@ -152,6 +153,21 @@ impl State { }) } + pub fn reload(&self) -> Result<()> { + let mut hostmap = HashMap::new(); + if let Some(hostsfile) = &self.config.hosts.hostsfile { + let hosts = hostfile::parse_file(hostsfile) + .map_err(|e| anyhow!("Error parsing host file: {e}"))?; + for host in hosts { + if let IpAddr::V4(ip) = host.ip { + hostmap.insert(ip, host.names[0].clone()); + } + } + } + self.hostmap.send_replace(hostmap); + Ok(()) + } + pub fn gc_chunks(&self) -> Result<()> { self.image_stats.send_modify(|image_stats| { let mut chunk_stats = self.chunk_stats.lock().unwrap(); diff --git a/pixie.service b/pixie.service index 419fa511..31ff5c41 100644 --- a/pixie.service +++ b/pixie.service @@ -9,6 +9,7 @@ Group=root ExecStart=/usr/local/bin/pixie-server -s /var/local/lib/pixie Environment=RUST_LOG=info WorkingDirectory=/var/local/lib/pixie +ExecReload=kill -HUP $MAINPID [Install] WantedBy=multi-user.target