tracing: Simplify initialization (#1026)

olix0r · web-flow · commit 6a9c2135238d · 2021-06-02T17:18:42.000-07:00
The `Settings::build` method is fairly large and complex. This change splits into a few smaller units and reduces boilerplate--especially around reading the settings from the environment. This change is being made in service of #601 to help avoid making the build method even more complicated.
diff --git a/linkerd/tracing/src/lib.rs b/linkerd/tracing/src/lib.rs
@@ -12,27 +12,48 @@ use std::{env, str};
 pub use tokio_trace::tasks::TaskList;
 use tokio_trace::tasks::TasksLayer;
 use tracing::Dispatch;
-use tracing_subscriber::{fmt::format, prelude::*};
+use tracing_subscriber::{
+    fmt::format::{self, DefaultFields},
+    layer::Layered,
+    prelude::*,
+    reload, EnvFilter,
+};
+
+type Registry =
+    Layered<reload::Layer<EnvFilter, tracing_subscriber::Registry>, tracing_subscriber::Registry>;
 
 const ENV_LOG_LEVEL: &str = "LINKERD2_PROXY_LOG";
 const ENV_LOG_FORMAT: &str = "LINKERD2_PROXY_LOG_FORMAT";
 
 const DEFAULT_LOG_LEVEL: &str = "warn,linkerd=info";
 const DEFAULT_LOG_FORMAT: &str = "PLAIN";
 
+#[derive(Debug, Default)]
+pub struct Settings {
+    filter: Option<String>,
+    format: Option<String>,
+    is_test: bool,
+}
+
+#[derive(Clone)]
+pub struct Handle(Inner);
+
+#[derive(Clone)]
+enum Inner {
+    Disabled,
+    Enabled {
+        level: level::Handle,
+        tasks: TaskList,
+    },
+}
+
 /// Initialize tracing and logging with the value of the `ENV_LOG`
 /// environment variable as the verbosity-level filter.
 pub fn init() -> Result<Handle, Error> {
-    let log_level = env::var(ENV_LOG_LEVEL).unwrap_or_else(|_| DEFAULT_LOG_LEVEL.to_string());
-    if let "OFF" = log_level.to_uppercase().trim() {
-        return Ok(Handle(Inner::Disabled));
-    }
-
-    let log_format = env::var(ENV_LOG_FORMAT).unwrap_or_else(|_| DEFAULT_LOG_FORMAT.to_string());
-    let (dispatch, handle) = Settings::default()
-        .filter(log_level)
-        .format(log_format)
-        .build();
+    let (dispatch, handle) = match Settings::from_env() {
+        Some(s) => s.build(),
+        None => return Ok(Handle(Inner::Disabled)),
+    };
 
     // Set the default subscriber.
     tracing::dispatcher::set_global_default(dispatch)?;
@@ -57,110 +78,104 @@ pub fn init_log_compat() -> Result<(), Error> {
     Ok(())
 }
 
-#[derive(Debug, Default)]
-pub struct Settings {
-    filter: Option<String>,
-    format: Option<String>,
-    test: bool,
-}
+// === impl Settings ===
 
 impl Settings {
-    pub fn from_env() -> Self {
-        let mut settings = Settings::default();
-        if let Ok(filter) = env::var(ENV_LOG_LEVEL) {
-            settings = settings.filter(filter);
-        }
-        if let Ok(format) = env::var(ENV_LOG_FORMAT) {
-            settings = settings.format(format);
+    pub fn from_env() -> Option<Self> {
+        let filter = std::env::var(ENV_LOG_LEVEL).ok();
+        if let Some(level) = filter.as_ref() {
+            if level.to_uppercase().trim() == "OFF" {
+                return None;
+            }
         }
-        settings
+
+        Some(Self {
+            filter,
+            format: std::env::var(ENV_LOG_FORMAT).ok(),
+            is_test: false,
+        })
     }
 
-    pub fn filter(self, filter: impl Into<String>) -> Self {
+    fn for_test(filter: String, format: String) -> Self {
         Self {
-            filter: Some(filter.into()),
-            ..self
+            filter: Some(filter),
+            format: Some(format),
+            is_test: true,
         }
     }
 
-    pub fn format(self, format: impl Into<String>) -> Self {
-        Self {
-            format: Some(format.into()),
-            ..self
-        }
+    fn format(&self) -> String {
+        self.format
+            .as_deref()
+            .unwrap_or(DEFAULT_LOG_FORMAT)
+            .to_uppercase()
     }
 
-    pub fn test(self, test: bool) -> Self {
-        Self { test, ..self }
+    fn mk_registry(&self) -> (Registry, level::Handle) {
+        let log_level = self.filter.as_deref().unwrap_or(DEFAULT_LOG_LEVEL);
+        let (filter, level) = reload::Layer::new(EnvFilter::new(log_level));
+        let reg = tracing_subscriber::registry().with(filter);
+        (reg, level::Handle::new(level))
     }
 
-    pub fn build(self) -> (Dispatch, Handle) {
-        let filter = self.filter.unwrap_or_else(|| DEFAULT_LOG_LEVEL.to_string());
-        let format = self
-            .format
-            .unwrap_or_else(|| DEFAULT_LOG_FORMAT.to_string());
+    fn mk_json(&self, registry: Registry) -> (Dispatch, TaskList) {
+        let (tasks, tasks_layer) = TasksLayer::<format::JsonFields>::new();
+        let registry = registry.with(tasks_layer);
 
-        // Set up the subscriber
         let fmt = tracing_subscriber::fmt::format()
             .with_timer(Uptime::starting_now())
-            .with_thread_ids(!self.test);
-        let filter = tracing_subscriber::EnvFilter::new(filter);
-        let (filter, level) = tracing_subscriber::reload::Layer::new(filter);
-        let level = level::Handle::new(level);
-        let registry = tracing_subscriber::registry().with(filter);
-
-        let (dispatch, tasks) = match format.to_uppercase().as_ref() {
-            "JSON" => {
-                let (tasks, tasks_layer) = TasksLayer::<format::JsonFields>::new();
-                let fmt = fmt
-                    // Configure the formatter to output JSON logs.
-                    .json()
-                    // Output the current span context as a JSON list.
-                    .with_span_list(true)
-                    // Don't output a field for the current span, since this
-                    // would duplicate information already in the span list.
-                    .with_current_span(false);
-                let fmt = tracing_subscriber::fmt::layer()
-                    // Use the JSON event formatter.
-                    .event_format(fmt)
-                    // Since we're using the JSON event formatter, we must also
-                    // use the JSON field formatter.
-                    .fmt_fields(format::JsonFields::default());
-                let registry = registry.with(tasks_layer);
-                let dispatch = if self.test {
-                    registry.with(fmt.with_test_writer()).into()
-                } else {
-                    registry.with(fmt).into()
-                };
-                (dispatch, tasks)
-            }
-            _ => {
-                let (tasks, tasks_layer) = TasksLayer::<format::DefaultFields>::new();
-                let registry = registry.with(tasks_layer);
-                let fmt = tracing_subscriber::fmt::layer().event_format(fmt);
-                let dispatch = if self.test {
-                    registry.with(fmt.with_test_writer()).into()
-                } else {
-                    registry.with(fmt).into()
-                };
-                (dispatch, tasks)
-            }
+            .with_thread_ids(!self.is_test)
+            // Configure the formatter to output JSON logs.
+            .json()
+            // Output the current span context as a JSON list.
+            .with_span_list(true)
+            // Don't output a field for the current span, since this
+            // would duplicate information already in the span list.
+            .with_current_span(false);
+
+        let fmt = tracing_subscriber::fmt::layer()
+            // Use the JSON event formatter.
+            .event_format(fmt)
+            // Since we're using the JSON event formatter, we must also
+            // use the JSON field formatter.
+            .fmt_fields(format::JsonFields::default());
+
+        let dispatch = if self.is_test {
+            registry.with(fmt.with_test_writer()).into()
+        } else {
+            registry.with(fmt).into()
         };
 
-        (dispatch, Handle(Inner::Enabled { level, tasks }))
+        (dispatch, tasks)
     }
-}
 
-#[derive(Clone)]
-pub struct Handle(Inner);
+    fn mk_plain(&self, registry: Registry) -> (Dispatch, TaskList) {
+        let (tasks, tasks_layer) = TasksLayer::<DefaultFields>::new();
+        let registry = registry.with(tasks_layer);
 
-#[derive(Clone)]
-enum Inner {
-    Disabled,
-    Enabled {
-        level: level::Handle,
-        tasks: TaskList,
-    },
+        let fmt = tracing_subscriber::fmt::format()
+            .with_timer(Uptime::starting_now())
+            .with_thread_ids(!self.is_test);
+        let fmt = tracing_subscriber::fmt::layer().event_format(fmt);
+        let dispatch = if self.is_test {
+            registry.with(fmt.with_test_writer()).into()
+        } else {
+            registry.with(fmt).into()
+        };
+
+        (dispatch, tasks)
+    }
+
+    pub fn build(self) -> (Dispatch, Handle) {
+        let (registry, level) = self.mk_registry();
+
+        let (dispatch, tasks) = match self.format().as_ref() {
+            "JSON" => self.mk_json(registry),
+            _ => self.mk_plain(registry),
+        };
+
+        (dispatch, Handle(Inner::Enabled { level, tasks }))
+    }
 }
 
 // === impl Handle ===
diff --git a/linkerd/tracing/src/test.rs b/linkerd/tracing/src/test.rs
@@ -15,11 +15,7 @@ pub fn trace_subscriber(default: impl ToString) -> (Dispatch, Handle) {
     // This may fail, since the global log compat layer may have been
     // initialized by another test.
     let _ = init_log_compat();
-    Settings::default()
-        .filter(log_level)
-        .format(log_format)
-        .test(true)
-        .build()
+    Settings::for_test(log_level, log_format).build()
 }
 
 pub fn with_default_filter(default: impl ToString) -> tracing::dispatcher::DefaultGuard {
