From d0c2d56fbd49daa99d61557030a679d8bd7f6c4b Mon Sep 17 00:00:00 2001
From: Christian Schwarz <christian@neon.tech>
Date: Wed, 13 Dec 2023 15:04:28 +0000
Subject: [PATCH] extract env var stuff into other utility module

---
 libs/utils/src/env_config.rs | 48 ++++++++++++++++++++++++++++++++++++
 libs/utils/src/lib.rs        |  1 +
 libs/utils/src/logging.rs    | 42 ++++++++++++++-----------------
 3 files changed, 68 insertions(+), 23 deletions(-)
 create mode 100644 libs/utils/src/env_config.rs

diff --git a/libs/utils/src/env_config.rs b/libs/utils/src/env_config.rs
new file mode 100644
index 0000000000..449c9aa9ac
--- /dev/null
+++ b/libs/utils/src/env_config.rs
@@ -0,0 +1,48 @@
+use std::{fmt::Display, str::FromStr};
+
+pub fn var<V, E, D>(varname: &str, default: D) -> V
+where
+    V: FromStr<Err = E>,
+    E: Display,
+    D: FnOnce() -> V,
+{
+    match std::env::var(varname) {
+        Ok(s) => s
+            .parse()
+            .map_err(|e| format!("failed to parse env var {varname}: {e:#}"))
+            .unwrap(),
+        Err(std::env::VarError::NotPresent) => default(),
+        Err(std::env::VarError::NotUnicode(_)) => {
+            panic!("env var {varname} is not unicode")
+        }
+    }
+}
+
+pub struct Bool(bool);
+
+impl Bool {
+    pub const fn new_const<const V: bool>() -> Self {
+        Bool(V)
+    }
+}
+
+impl FromStr for Bool {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        if let Ok(b) = s.parse() {
+            return Ok(b);
+        }
+        Ok(Bool(match s {
+            "0" => false,
+            "1" => true,
+            _ => return Err(format!("not a bool, accepting 0|1|{}|{}", false, true)),
+        }))
+    }
+}
+
+impl Into<bool> for Bool {
+    fn into(self) -> bool {
+        self.0
+    }
+}
diff --git a/libs/utils/src/lib.rs b/libs/utils/src/lib.rs
index bb6c848bf4..3511eaa3ca 100644
--- a/libs/utils/src/lib.rs
+++ b/libs/utils/src/lib.rs
@@ -27,6 +27,7 @@ pub mod auth;
 pub mod id;
 
 mod hex;
+pub mod env_config;
 pub use hex::Hex;
 
 // http endpoint utils
diff --git a/libs/utils/src/logging.rs b/libs/utils/src/logging.rs
index 844c526c96..ab2489ad09 100644
--- a/libs/utils/src/logging.rs
+++ b/libs/utils/src/logging.rs
@@ -1,9 +1,15 @@
-use std::{io::BufWriter, str::FromStr, sync::{Arc, Mutex}};
+use std::{
+    io::BufWriter,
+    str::FromStr,
+    sync::{Arc, Mutex},
+};
 
 use anyhow::Context;
 use once_cell::sync::Lazy;
 use strum_macros::{EnumString, EnumVariantNames};
 
+use super::env_config;
+
 #[derive(EnumString, EnumVariantNames, Eq, PartialEq, Debug, Clone, Copy)]
 #[strum(serialize_all = "snake_case")]
 pub enum LogFormat {
@@ -100,26 +106,6 @@ pub fn init(
             .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info"))
     };
 
-    // WIP: lift it up as an argument
-    let enable_tracing_chrome = match std::env::var("NEON_PAGESERVER_ENABLE_TRACING_CHROME") {
-        Ok(s) if s != "0" => true,
-        Ok(_s) => false,
-        Err(std::env::VarError::NotPresent) => false,
-        Err(std::env::VarError::NotUnicode(_)) => {
-            panic!("env var NEON_PAGESERVER_ENABLE_TRACING_CHROME not unicode")
-        }
-    };
-
-    // WIP: lift it up as an argument
-    let enable_tracing_flame = match std::env::var("NEON_PAGESERVER_ENABLE_TRACING_FLAME") {
-        Ok(s) if s != "0" => true,
-        Ok(_s) => false,
-        Err(std::env::VarError::NotPresent) => false,
-        Err(std::env::VarError::NotUnicode(_)) => {
-            panic!("env var NEON_PAGESERVER_ENABLE_TRACING_FLAME not unicode")
-        }
-    };
-
     // NB: the order of the with() calls does not matter.
     // See https://docs.rs/tracing-subscriber/0.3.16/tracing_subscriber/layer/index.html#per-layer-filtering
     use tracing_subscriber::prelude::*;
@@ -165,7 +151,12 @@ pub fn init(
     layers
         .add_layer(TracingEventCountLayer(&TRACING_EVENT_COUNT).with_filter(rust_log_env_filter()));
 
-    let tracing_chrome_layer_flush_guard = if enable_tracing_chrome {
+    let tracing_chrome_layer_flush_guard = if env_config::var(
+        "NEON_UTILS_LOGGING_ENABLE_TRACING_CHROME",
+        env_config::Bool::new_const::<false>,
+    )
+    .into()
+    {
         let (layer, guard) = tracing_chrome::ChromeLayerBuilder::new()
             .trace_style(tracing_chrome::TraceStyle::Async)
             .build();
@@ -175,7 +166,12 @@ pub fn init(
         None
     };
 
-    let tracing_flame_flush_guard = if enable_tracing_flame {
+    let tracing_flame_flush_guard = if env_config::var(
+        "NEON_UTILS_LOGGING_ENABLE_TRACING_FLAME",
+        env_config::Bool::new_const::<false>,
+    )
+    .into()
+    {
         let (layer, guard) = tracing_flame::FlameLayer::with_file("./tracing.folded").unwrap();
         let layer = layer
             .with_empty_samples(false)