start cheating systemdcheat-systemd

Signed-off-by: Christine Dodrill <me@christine.website>

7 files changed, 312 insertions, 46 deletions

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",
@@ -226,6 +228,12 @@ dependencies = [
 
 [[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"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
@@ -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",
 ]
@@ -974,6 +982,17 @@ 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"
 source = "registry+https://github.com/rust-lang/crates.io-index"
@@ -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",
@@ -2377,6 +2394,15 @@ 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"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "bc5cf98d8186244414c848017f0e2676b3fcb46807f6668a97dfe67359a3c4b7"
@@ -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 <plhk@sdf.org>", "Xe <me@christine.website>"]
+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<Self, std::io::Error> {
+        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<env::VarError> for Error {
+    fn from(_: env::VarError) -> Error {
+        Error::NoSocket
+    }
+}
+
+impl From<std::io::Error> 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<Self, Error> {
+        let sockname = env::var("NOTIFY_SOCKET")?;
+        Self::from_path(sockname)
+    }
+
+    pub fn from_path<P: AsRef<Path>>(path: P) -> Result<Self, Error> {
+        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::<u32>()?;
+        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 server = warp::serve(site);
+        let mut lfd = listenfd::ListenFd::from_env();
 
-    match std::env::var("SOCKPATH") {
-        Ok(sockpath) => {
-            let _ = std::fs::remove_file(&sockpath);
-            let listener = UnixListener::bind(sockpath)?;
-            let incoming = UnixListenerStream::new(listener);
+        if let Some(lis) = lfd.take_unix_listener(0)? {
+            let incoming = UnixListenerStream::new(UnixListener::from_std(lis)?);
             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::<u16>()?,
-                ))
-                .await;
-
-            Ok(())
+        Ok(())
+    }
+
+    #[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;
+
+                Ok(())
+            }
+            Err(_) => {
+                server
+                    .run((
+                        IpAddr::from_str(&std::env::var("HOST").unwrap_or("::".into()))?,
+                        std::env::var("PORT")
+                            .unwrap_or("3030".into())
+                            .parse::<u16>()?,
+                    ))
+                    .await;
+
+                Ok(())
+            }
         }
     }
 }