From 874e4d706c31cd9ce03efd8f03b6fec7c19612a9 Mon Sep 17 00:00:00 2001
From: SeanOMik <seanomik@gmail.com>
Date: Sun, 21 Jan 2024 10:46:39 -0500
Subject: [PATCH] add listen addr and port to config, listen for ctrl+c and
 terminate signals

---
 Cargo.lock  | 22 ++++++++++++++++++++++
 Cargo.toml  |  1 +
 README.md   | 23 +++++++++++++++++++++++
 src/main.rs | 50 ++++++++++++++++++++++++++++++++++++++++++++------
 4 files changed, 90 insertions(+), 6 deletions(-)
 create mode 100644 README.md

diff --git a/Cargo.lock b/Cargo.lock
index 4a481d5..90f755c 100755
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -307,6 +307,16 @@ version = "0.8.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "06ea2b9bc92be3c2baa9334a323ebca2d6f074ff852cd1d7b11064035cd3868f"
 
+[[package]]
+name = "ctrlc"
+version = "3.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b467862cc8610ca6fc9a1532d7777cee0804e678ab45410897b9396495994a0b"
+dependencies = [
+ "nix",
+ "windows-sys 0.52.0",
+]
+
 [[package]]
 name = "encoding_rs"
 version = "0.8.33"
@@ -785,6 +795,17 @@ dependencies = [
  "tempfile",
 ]
 
+[[package]]
+name = "nix"
+version = "0.27.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2eb04e9c688eff1c89d72b407f168cf79bb9e867a9d3323ed6c01519eb9cc053"
+dependencies = [
+ "bitflags 2.4.2",
+ "cfg-if",
+ "libc",
+]
+
 [[package]]
 name = "nu-ansi-term"
 version = "0.46.0"
@@ -1304,6 +1325,7 @@ dependencies = [
  "anyhow",
  "axum",
  "clap",
+ "ctrlc",
  "figment",
  "prometheus",
  "reqwest",
diff --git a/Cargo.toml b/Cargo.toml
index 49a4888..0a22986 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,6 +9,7 @@ edition = "2021"
 anyhow = "1.0.79"
 axum = "0.7.4"
 clap = { version = "4.4.18", features = ["derive"] }
+ctrlc = "3.4.2"
 figment = { version = "0.10.14", features = ["toml", "env"] }
 prometheus = "0.13.3"
 reqwest = { version = "0.11.23", features = ["json"] }
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..cecbe58
--- /dev/null
+++ b/README.md
@@ -0,0 +1,23 @@
+# Tautulli Prometheus Exporter
+This is a small application that can collect information from Tautulli and exports it in a Prometheus format.
+
+## Config
+```toml
+# The full url of tautulli
+tautulli_url = "https://tautulli.example.com/"
+tautulli_apikey = ""
+# The address to listen on. Default is 0.0.0.0
+listen_address = "0.0.0.0"
+# The port to listen on. Default is 3000
+listen_port = "3000"
+```
+
+## Docker image
+Docker images are published [here](https://git.seanomik.net/SeanOMik/-/packages/container/tautulli-exporter/v0.1.0). 
+
+You can run the following command to run a docker container:
+```shell
+$ docker run -it --rm -v $PWD/config.toml:/app/config.toml -p 3000:3000 git.seanomik.net/seanomik/tautulli-exporter:v0.1.0
+```
+
+## Exported metrics
diff --git a/src/main.rs b/src/main.rs
index cc856ee..6394dd0 100755
--- a/src/main.rs
+++ b/src/main.rs
@@ -23,7 +23,7 @@ use figment::{
     providers::{Env, Format, Toml},
     Figment,
 };
-use tokio::sync::Mutex;
+use tokio::{sync::{Mutex, mpsc::{Receiver, self}}, select, signal};
 
 mod dto;
 use dto::*;
@@ -183,12 +183,24 @@ struct Config {
     tautulli_apikey: String,
     #[serde(default = "metrics_prefix_default")]
     metrics_prefix: String,
+    #[serde(default = "listen_address_default")]
+    listen_address: String,
+    #[serde(default = "listen_port_default")]
+    listen_port: String,
 }
 
 fn metrics_prefix_default() -> String {
     "tautulli".to_string()
 }
 
+fn listen_address_default() -> String {
+    "0.0.0.0".to_string()
+}
+
+fn listen_port_default() -> String {
+    "3000".to_string()
+}
+
 impl Config {
     /// Constructs an api endpoint for you
     fn api_endpoint(&self, command: &str) -> String {
@@ -236,6 +248,30 @@ struct AppState {
     metrics: Arc<Mutex<Metrics>>,
 }
 
+async fn shutdown_signal() {
+    let ctrl_c = async {
+        signal::ctrl_c()
+            .await
+            .expect("failed to install Ctrl+C handler");
+    };
+
+    #[cfg(unix)]
+    let terminate = async {
+        signal::unix::signal(signal::unix::SignalKind::terminate())
+            .expect("failed to install signal handler")
+            .recv()
+            .await;
+    };
+
+    #[cfg(not(unix))]
+    let terminate = std::future::pending::<()>();
+
+    tokio::select! {
+        _ = ctrl_c => { info!("Received ctrl+c, exiting...") },
+        _ = terminate => {  info!("Received terminate signal, exiting...") },
+    }
+}
+
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     tracing_subscriber::registry()
@@ -254,6 +290,7 @@ async fn main() -> anyhow::Result<()> {
         .merge(Toml::file(args.config))
         .merge(Env::prefixed("TAUTEXP_"))
         .extract()?;
+    let full_listen_address = format!("{}:{}", config.listen_address, config.listen_port);
 
     // if the url ends with a `/` it can cause some issues
     if config.tautulli_url.ends_with("/") {
@@ -275,17 +312,18 @@ async fn main() -> anyhow::Result<()> {
         .with_state(state)
         .layer(TraceLayer::new_for_http());
 
-    let bind = "0.0.0.0:3000";
-    let listener = tokio::net::TcpListener::bind(bind).await?;
-    info!("Starting http server, listening on {}", bind);
-    axum::serve(listener, app).await?;
+    let listener = tokio::net::TcpListener::bind(&full_listen_address).await?;
+    info!("Starting http server, listening on {}", full_listen_address);
+    axum::serve(listener, app)
+        .with_graceful_shutdown(shutdown_signal())
+        .await?;
 
     Ok(())
 }
 
 // basic handler that responds with a static string
 async fn root() -> &'static str {
-    "Hello, World!"
+    "Up"
 }
 
 async fn metrics(State(state): State<AppState>) -> Result<String, CollectionError> {