From 50e57c427c8f8978d9125721c677bd935bb854e7 Mon Sep 17 00:00:00 2001 From: Christine Dodrill Date: Sat, 15 May 2021 00:11:59 -0400 Subject: [PATCH] start cheating systemd Signed-off-by: Christine Dodrill --- Cargo.lock | 71 +++++++++++----- Cargo.toml | 9 +- lib/sdnotify/Cargo.toml | 11 +++ lib/sdnotify/src/lib.rs | 184 ++++++++++++++++++++++++++++++++++++++++ shell.nix | 1 + src/bin/decoy.py | 13 +++ src/main.rs | 67 ++++++++++----- 7 files changed, 311 insertions(+), 45 deletions(-) create mode 100644 lib/sdnotify/Cargo.toml create mode 100644 lib/sdnotify/src/lib.rs create mode 100755 src/bin/decoy.py diff --git a/Cargo.lock b/Cargo.lock index 2945d6c..e963762 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1,5 +1,7 @@ # This file is automatically @generated by Cargo. # It is not intended for manual editing. +version = 3 + [[package]] name = "abnf" version = "0.6.1" @@ -102,7 +104,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9d117600f438b1707d4e4ae15d3595657288f8235a0eb593e80ecc98ab34e1bc" dependencies = [ "addr2line", - "cfg-if", + "cfg-if 1.0.0", "libc", "miniz_oxide", "object", @@ -224,6 +226,12 @@ dependencies = [ "tracing-futures", ] +[[package]] +name = "cfg-if" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4785bdd1c96b2a846b2bd7cc02e86b6b3dbf14e7e53446c4f54c92a361040822" + [[package]] name = "cfg-if" version = "1.0.0" @@ -342,7 +350,7 @@ version = "1.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "81156fece84ab6a9f2afdb109ce3ae577e42b1228441eded99bd77f627953b1a" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", ] [[package]] @@ -432,7 +440,7 @@ version = "0.8.28" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "80df024fbc5ac80f87dfef0d9f5209a252f2a497f7f42944cff24d8253cac065" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", ] [[package]] @@ -495,7 +503,7 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cd3aec53de10fe96d7d8c565eb17f2c687bb5518a2ec453b5b1252964526abe0" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "crc32fast", "libc", "miniz_oxide", @@ -655,7 +663,7 @@ version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "libc", "wasi 0.9.0+wasi-snapshot-preview1", ] @@ -666,7 +674,7 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c9495705279e7140bf035dde1f6e750c162df8b625267cd52cc44e0b156732c8" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "libc", "wasi 0.10.2+wasi-snapshot-preview1", ] @@ -884,7 +892,7 @@ version = "0.1.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "61124eeebbd69b8190558df225adf7e4caafce0d743919e5d6b19652314ec5ec" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", ] [[package]] @@ -956,7 +964,7 @@ checksum = "21f866863575d0e1d654fbeeabdc927292fdf862873dc3c96c6f753357e13374" dependencies = [ "arrayvec", "bitflags", - "cfg-if", + "cfg-if 1.0.0", "ryu", "static_assertions", ] @@ -973,6 +981,17 @@ version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7fb9b38af92608140b86b693604b9ffcc5824240a484d1ecd4795bacb2fe88f3" +[[package]] +name = "listenfd" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "492158e732f2e2de81c592f0a2427e57e12cd3d59877378fe7af624b6bbe0ca1" +dependencies = [ + "libc", + "uuid 0.6.5", + "winapi", +] + [[package]] name = "lock_api" version = "0.4.2" @@ -988,7 +1007,7 @@ version = "0.4.14" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "51b9bbe6c47d51fc3e1a9b945965946b4c44142ab8792c50835a980d362c2710" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", ] [[package]] @@ -1219,7 +1238,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "038d43985d1ddca7a9900630d8cd031b56e4794eecc2e9ea39dd17aa04399a70" dependencies = [ "bitflags", - "cfg-if", + "cfg-if 1.0.0", "foreign-types", "lazy_static", "libc", @@ -1268,7 +1287,7 @@ version = "0.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fa7a782938e745763fe6907fc6ba86946d72f49fe7e21de074e08128a99fb018" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "instant", "libc", "redox_syscall", @@ -1481,7 +1500,7 @@ version = "0.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5986aa8d62380092d2f50f8b1cdba9cb9b6731ffd4b25b51fd126b6c3e05b99c" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "fnv", "lazy_static", "libc", @@ -1731,8 +1750,6 @@ checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" [[package]] name = "sdnotify" version = "0.1.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "71ce7eac2075a4562fbcbad544cd55d72ebc760e0a5594a7c8829cf2b4b42a7a" [[package]] name = "security-framework" @@ -1854,7 +1871,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dfebf75d25bd900fd1e7d11501efab59bc846dbc76196839663e6637bba9f25f" dependencies = [ "block-buffer 0.9.0", - "cfg-if", + "cfg-if 1.0.0", "cpuid-bool", "digest 0.9.0", "opaque-debug 0.3.0", @@ -1867,7 +1884,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fa827a14b29ab7f44778d14a88d3cb76e949c45083f7dbfa507d0cb699dc12de" dependencies = [ "block-buffer 0.9.0", - "cfg-if", + "cfg-if 1.0.0", "cpuid-bool", "digest 0.9.0", "opaque-debug 0.3.0", @@ -1927,7 +1944,7 @@ version = "0.3.19" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "122e570113d28d773067fab24266b66753f6ea915758651696b6e35e49f88d6e" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "libc", "winapi", ] @@ -1983,7 +2000,7 @@ version = "3.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dac1c663cfc93810f88aed9b8941d48cabf856a1b111c29a40439018d870eb22" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "libc", "rand 0.8.3", "redox_syscall", @@ -2154,7 +2171,7 @@ version = "0.1.25" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "01ebdc2bb4498ab1ab5f5b73c5803825e60199229ccba0698170e3be0e7f959f" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "log", "pin-project-lite", "tracing-attributes", @@ -2375,6 +2392,15 @@ version = "0.7.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05e42f7c18b8f902290b009cde6d651262f956c98bc51bca4cd1d511c9cd85c7" +[[package]] +name = "uuid" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e1436e58182935dcd9ce0add9ea0b558e8a87befe01c1a301e6020aeb0876363" +dependencies = [ + "cfg-if 0.1.10", +] + [[package]] name = "uuid" version = "0.8.2" @@ -2460,7 +2486,7 @@ version = "0.2.71" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7ee1280240b7c461d6a0071313e08f34a60b0365f14260362e5a2b17d1d31aa7" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "serde", "serde_json", "wasm-bindgen-macro", @@ -2487,7 +2513,7 @@ version = "0.4.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e67a5806118af01f0d9045915676b22aaebecf4178ae7021bc171dab0b897ab" dependencies = [ - "cfg-if", + "cfg-if 1.0.0", "js-sys", "wasm-bindgen", "web-sys", @@ -2601,6 +2627,7 @@ dependencies = [ "jsonfeed", "kankyo", "lazy_static", + "listenfd", "log", "mi", "mime", @@ -2624,7 +2651,7 @@ dependencies = [ "tracing-futures", "tracing-subscriber", "url", - "uuid", + "uuid 0.8.2", "warp", "xml-rs", ] diff --git a/Cargo.toml b/Cargo.toml index 82f4e99..733e2fb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -45,9 +45,13 @@ jsonfeed = { path = "./lib/jsonfeed" } mi = { path = "./lib/mi" } patreon = { path = "./lib/patreon" } +[dependencies.listenfd] +version = "0.3" +optional = true + # os-specific dependencies [target.'cfg(target_os = "linux")'.dependencies] -sdnotify = { version = "0.1", default-features = false } +sdnotify = { path = "./lib/sdnotify", optional = true } [build-dependencies] ructe = { version = "0.13", features = ["warp02"] } @@ -62,3 +66,6 @@ pretty_env_logger = "0" members = [ "./lib/*", ] + +[features] +systemd = ["sdnotify", "listenfd"] diff --git a/lib/sdnotify/Cargo.toml b/lib/sdnotify/Cargo.toml new file mode 100644 index 0000000..48fc438 --- /dev/null +++ b/lib/sdnotify/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "sdnotify" +version = "0.1.3" +authors = ["Alexander Polakov ", "Xe "] +description = "Notify service manager about start-up completion and other daemon status changes" +license = "MIT" +homepage = "https://github.com/polachok/sdnotify" +repository = "https://github.com/polachok/sdnotify" +documentation = "https://docs.rs/sdnotify" +edition = "2018" +keywords = ["systemd"] diff --git a/lib/sdnotify/src/lib.rs b/lib/sdnotify/src/lib.rs new file mode 100644 index 0000000..9c328f3 --- /dev/null +++ b/lib/sdnotify/src/lib.rs @@ -0,0 +1,184 @@ +//! Notify service manager about start-up completion and +//! other daemon status changes. +//! +//! ### Prerequisites +//! +//! A unit file with service type `Notify` is required. +//! +//! Example: +//! ```toml +//! [Unit] +//! Description=Frobulator +//! [Service] +//! Type=notify +//! ExecStart=/usr/sbin/frobulator +//! [Install] +//! WantedBy=multi-user.target +//! ``` +//! ### Sync API +//! ```no_run +//! use sdnotify::{SdNotify, Message, Error}; +//! +//! # fn notify() -> Result<(), Error> { +//! let notifier = SdNotify::from_env()?; +//! notifier.notify_ready()?; +//! # Ok(()) +//! # } +//! ``` +//! +//! ### Async API +//! ```no_run +//! use sdnotify::{Message, Error, async_io::SdNotify}; +//! use tokio::prelude::*; +//! use tokio::runtime::current_thread::Runtime; +//! +//! # fn notify() -> Result<(), Error> { +//! let notifier = SdNotify::from_env()?; +//! let mut rt = Runtime::new().unwrap(); +//! rt.block_on(notifier.send(Message::ready())).unwrap(); +//! # Ok(()) +//! # } +//! ``` + +use std::env; +use std::os::unix::net::UnixDatagram; +use std::path::Path; + +/// Message to send to init system +#[derive(Debug)] +pub struct Message(InnerMessage); + +impl Message { + /// Tells the init system that daemon startup is finished. + pub fn ready() -> Self { + Message(InnerMessage::Ready) + } + + /// Passes a single-line status string back to the init system that describes the daemon state. + pub fn status(status: String) -> Result { + if status.as_bytes().iter().any(|x| *x == b'\n') { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidData, + "newline not allowed", + )); + } + Ok(Message(InnerMessage::Status(status))) + } + + /// Tells systemd to update the watchdog timestamp. + /// This is the keep-alive ping that services need to issue in regular + /// intervals if WatchdogSec= is enabled for it. + pub fn watchdog() -> Self { + Message(InnerMessage::Watchdog) + } + + /// Tells systemd what the main pid of this service is. + /// This is needed in order to hack up 0-downtime deployments. + pub fn main_pid(pid: u32) -> Self { + Message(InnerMessage::MainPid(pid)) + } +} + +#[derive(Debug)] +enum InnerMessage { + Ready, + Status(String), + Watchdog, + MainPid(u32), +} + +#[derive(Debug)] +pub enum Error { + NoSocket, + Io(std::io::Error), +} + +impl std::fmt::Display for Error { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Error::NoSocket => write!(f, "NOTIFY_SOCKET variable not set"), + Error::Io(err) => write!(f, "{}", err), + } + } +} + +impl From for Error { + fn from(_: env::VarError) -> Error { + Error::NoSocket + } +} + +impl From for Error { + fn from(err: std::io::Error) -> Error { + Error::Io(err) + } +} + +impl std::error::Error for Error {} + +pub struct SdNotify(UnixDatagram); + +impl SdNotify { + pub fn from_env() -> Result { + let sockname = env::var("NOTIFY_SOCKET")?; + Self::from_path(sockname) + } + + pub fn from_path>(path: P) -> Result { + let socket = UnixDatagram::unbound()?; + socket.connect(path)?; + Ok(SdNotify(socket)) + } + + /// Tells the init system that daemon startup is finished. + pub fn notify_ready(&self) -> Result<(), std::io::Error> { + self.state(Message::ready()) + } + + /// Passes a single-line status string back to the init system that describes the daemon state. + pub fn set_status(&self, status: String) -> Result<(), std::io::Error> { + self.state(Message::status(status)?) + } + + /// Tells systemd to update the watchdog timestamp. + /// This is the keep-alive ping that services need to issue in regular + /// intervals if WatchdogSec= is enabled for it. + pub fn ping_watchdog(&self) -> Result<(), std::io::Error> { + self.state(Message::watchdog()) + } + + /// Tells systemd what the main pid of this service is. + /// This is needed in order to hack up 0-downtime deployments. + pub fn set_main_pid(&self, pid: u32) -> Result<(), std::io::Error> { + self.state(Message::main_pid(pid)) + } + + fn state(&self, state: Message) -> Result<(), std::io::Error> { + match state.0 { + InnerMessage::Ready => self.0.send(b"READY=1")?, + InnerMessage::Status(status) => self.0.send(format!("STATUS={}", status).as_bytes())?, + InnerMessage::Watchdog => self.0.send(b"WATCHDOG=1")?, + InnerMessage::MainPid(pid) => self.0.send(format!("MAINPID={}", pid).as_bytes())?, + }; + Ok(()) + } +} + +#[cfg(test)] +mod tests { + #[test] + fn ok() { + use super::*; + + let path = "/tmp/kek-async.sock"; + + let _ = std::fs::remove_file(path); + + let listener = UnixDatagram::bind(path).unwrap(); + let notifier = SdNotify::from_path(path).unwrap(); + notifier.state(Message::ready()).unwrap(); + let mut buf = [0; 100]; + listener.recv(&mut buf).unwrap(); + assert_eq!(&buf[..7], b"READY=1"); + } +} diff --git a/shell.nix b/shell.nix index 59bbdc9..e9df5f4 100644 --- a/shell.nix +++ b/shell.nix @@ -28,6 +28,7 @@ mkShell { # tools ispell + systemfd ]; SITE_PREFIX = "devel."; diff --git a/src/bin/decoy.py b/src/bin/decoy.py new file mode 100755 index 0000000..f985413 --- /dev/null +++ b/src/bin/decoy.py @@ -0,0 +1,13 @@ +#!/usr/bin/env nix-shell +#! nix-shell -p python3 -i python3 + +import os +import time + +pid = os.fork() +if pid == 0: + for fd in {0, 1, 2}: + os.close(fd) + time.sleep(1) +else: + print(pid) diff --git a/src/main.rs b/src/main.rs index cac19cf..4472623 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,8 +4,6 @@ extern crate tracing; use color_eyre::eyre::Result; use hyper::{header::CONTENT_TYPE, Body, Response}; use prometheus::{Encoder, TextEncoder}; -use std::net::IpAddr; -use std::str::FromStr; use std::sync::Arc; use tokio::net::UnixListener; use tokio_stream::wrappers::UnixListenerStream; @@ -33,6 +31,16 @@ async fn main() -> Result<()> { tracing_subscriber::fmt::init(); info!("starting up commit {}", env!("GITHUB_SHA")); + #[cfg(all(feature = "systemd", target_os = "linux"))] + { + use std::process::Command; + let pid = Command::new("./src/bin/decoy.py").output()?.stdout; + let pid = String::from_utf8(pid)?.trim().parse::()?; + if let Ok(ref mut n) = sdnotify::SdNotify::from_env() { + n.set_main_pid(pid)?; + } + } + let state = Arc::new( app::init( std::env::var("CONFIG_FNAME") @@ -211,8 +219,11 @@ async fn main() -> Result<()> { .with(warp::log(APPLICATION_NAME)) .recover(handlers::rejection); - #[cfg(target_os = "linux")] + let server = warp::serve(site); + + #[cfg(feature = "systemd")] { + #[cfg(target_os = "linux")] match sdnotify::SdNotify::from_env() { Ok(ref mut n) => { // shitty heuristic for detecting if we're running in prod @@ -234,30 +245,42 @@ async fn main() -> Result<()> { } Err(why) => error!("not running under systemd with Type=notify: {}", why), } + + let mut lfd = listenfd::ListenFd::from_env(); + + if let Some(lis) = lfd.take_unix_listener(0)? { + let incoming = UnixListenerStream::new(UnixListener::from_std(lis)?); + server.run_incoming(incoming).await; + } + Ok(()) } - let server = warp::serve(site); + #[cfg(not(feature = "systemd"))] + { + use std::net::IpAddr; + use std::str::FromStr; - match std::env::var("SOCKPATH") { - Ok(sockpath) => { - let _ = std::fs::remove_file(&sockpath); - let listener = UnixListener::bind(sockpath)?; - let incoming = UnixListenerStream::new(listener); - server.run_incoming(incoming).await; + match std::env::var("SOCKPATH") { + Ok(sockpath) => { + let _ = std::fs::remove_file(&sockpath); + let listener = UnixListener::bind(sockpath)?; + let incoming = UnixListenerStream::new(listener); + server.run_incoming(incoming).await; - Ok(()) - } - Err(_) => { - server - .run(( - IpAddr::from_str(&std::env::var("HOST").unwrap_or("::".into()))?, - std::env::var("PORT") - .unwrap_or("3030".into()) - .parse::()?, - )) - .await; + Ok(()) + } + Err(_) => { + server + .run(( + IpAddr::from_str(&std::env::var("HOST").unwrap_or("::".into()))?, + std::env::var("PORT") + .unwrap_or("3030".into()) + .parse::()?, + )) + .await; - Ok(()) + Ok(()) + } } } }