From 9241df18e2cddd0ce14aeff82526cfdb6c59730c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jose=20Fern=C3=A1ndez?= <joseferper@gmail.com>
Date: Tue, 28 Feb 2023 22:10:40 +0100
Subject: [PATCH] Allow sending logs to stdout by using STDOUT_LOG env var
 (#334)

* Allow sending logs to stdout by using STDOUT_LOG env var

* Increase stats buffer size
---
 src/lib.rs          |  1 +
 src/main.rs         |  5 +--
 src/multi_logger.rs | 80 +++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 84 insertions(+), 2 deletions(-)
 create mode 100644 src/multi_logger.rs

diff --git a/src/lib.rs b/src/lib.rs
index e9a683f..63eae59 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -2,6 +2,7 @@ pub mod config;
 pub mod constants;
 pub mod errors;
 pub mod messages;
+pub mod multi_logger;
 pub mod pool;
 pub mod scram;
 pub mod server;
diff --git a/src/main.rs b/src/main.rs
index 33536b0..b3ef77c 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -66,6 +66,7 @@ mod config;
 mod constants;
 mod errors;
 mod messages;
+mod multi_logger;
 mod pool;
 mod prometheus;
 mod query_router;
@@ -81,7 +82,7 @@ use crate::prometheus::start_metric_server;
 use crate::stats::{Collector, Reporter, REPORTER};
 
 fn main() -> Result<(), Box<dyn std::error::Error>> {
-    env_logger::builder().format_timestamp_micros().init();
+    multi_logger::MultiLogger::init().unwrap();
 
     info!("Welcome to PgCat! Meow. (Version {})", VERSION);
 
@@ -160,7 +161,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         let client_server_map: ClientServerMap = Arc::new(Mutex::new(HashMap::new()));
 
         // Statistics reporting.
-        let (stats_tx, stats_rx) = mpsc::channel(100_000);
+        let (stats_tx, stats_rx) = mpsc::channel(500_000);
         REPORTER.store(Arc::new(Reporter::new(stats_tx.clone())));
 
         // Connection pool that allows to query all shards and replicas.
diff --git a/src/multi_logger.rs b/src/multi_logger.rs
new file mode 100644
index 0000000..901db02
--- /dev/null
+++ b/src/multi_logger.rs
@@ -0,0 +1,80 @@
+use log::{Level, Log, Metadata, Record, SetLoggerError};
+
+// This is a special kind of logger that allows sending logs to different
+// targets depending on the log level.
+//
+// By default, if nothing is set, it acts as a regular env_log logger,
+// it sends everything to standard error.
+//
+// If the Env variable `STDOUT_LOG` is defined, it will be used for
+// configuring the standard out logger.
+//
+// The behavior is:
+//   - If it is an error, the message is written to standard error.
+//   - If it is not, and it matches the log level of the standard output logger (`STDOUT_LOG` env var), it will be send to standard output.
+//   - If the above is not true, it is sent to the stderr logger that will log it or not depending on the value
+//     of the RUST_LOG env var.
+//
+// So to summarize, if no `STDOUT_LOG` env var is present, the logger is the default logger. If `STDOUT_LOG` is set, everything
+// but errors, that matches the log level set in the `STDOUT_LOG` env var is sent to stdout. You can have also some esoteric configuration
+// where you set `RUST_LOG=debug` and `STDOUT_LOG=info`, in here, erros will go to stderr, warns and infos to stdout and debugs to stderr.
+//
+pub struct MultiLogger {
+    stderr_logger: env_logger::Logger,
+    stdout_logger: env_logger::Logger,
+}
+
+impl MultiLogger {
+    fn new() -> Self {
+        let stderr_logger = env_logger::builder().format_timestamp_micros().build();
+        let stdout_logger = env_logger::Builder::from_env("STDOUT_LOG")
+            .format_timestamp_micros()
+            .target(env_logger::Target::Stdout)
+            .build();
+
+        Self {
+            stderr_logger,
+            stdout_logger,
+        }
+    }
+
+    pub fn init() -> Result<(), SetLoggerError> {
+        let logger = Self::new();
+
+        log::set_max_level(logger.stderr_logger.filter());
+        log::set_boxed_logger(Box::new(logger))
+    }
+}
+
+impl Log for MultiLogger {
+    fn enabled(&self, metadata: &Metadata) -> bool {
+        self.stderr_logger.enabled(metadata) && self.stdout_logger.enabled(metadata)
+    }
+
+    fn log(&self, record: &Record) {
+        if record.level() == Level::Error {
+            self.stderr_logger.log(record);
+        } else {
+            if self.stdout_logger.matches(record) {
+                self.stdout_logger.log(record);
+            } else {
+                self.stderr_logger.log(record);
+            }
+        }
+    }
+
+    fn flush(&self) {
+        self.stderr_logger.flush();
+        self.stdout_logger.flush();
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn test_init() {
+        MultiLogger::init().unwrap();
+    }
+}