move cli to top level

cli/src/log.rs

@@ -0,0 +1,389 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+use chrono::Local;
+use opentelemetry::{
+    sdk::trace::Tracer,
+    trace::{SpanBuilder, Tracer as TraitTracer},
+};
+use std::fmt;
+use std::{env, path::Path, sync::Arc};
+use std::{
+    io::Write,
+    sync::atomic::{AtomicU32, Ordering},
+};
+
+const NO_COLOR_ENV: &str = "NO_COLOR";
+
+static INSTANCE_COUNTER: AtomicU32 = AtomicU32::new(0);
+
+// Gets a next incrementing number that can be used in logs
+pub fn next_counter() -> u32 {
+    INSTANCE_COUNTER.fetch_add(1, Ordering::SeqCst)
+}
+
+// Log level
+#[derive(clap::ArgEnum, PartialEq, Eq, PartialOrd, Clone, Copy, Debug)]
+pub enum Level {
+    Trace = 0,
+    Debug,
+    Info,
+    Warn,
+    Error,
+    Critical,
+    Off,
+}
+
+impl Default for Level {
+    fn default() -> Self {
+        Level::Info
+    }
+}
+
+impl fmt::Display for Level {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        match self {
+            Level::Critical => write!(f, "critical"),
+            Level::Debug => write!(f, "debug"),
+            Level::Error => write!(f, "error"),
+            Level::Info => write!(f, "info"),
+            Level::Off => write!(f, "off"),
+            Level::Trace => write!(f, "trace"),
+            Level::Warn => write!(f, "warn"),
+        }
+    }
+}
+
+impl Level {
+    pub fn name(&self) -> Option<&str> {
+        match self {
+            Level::Trace => Some("trace"),
+            Level::Debug => Some("debug"),
+            Level::Info => Some("info"),
+            Level::Warn => Some("warn"),
+            Level::Error => Some("error"),
+            Level::Critical => Some("critical"),
+            Level::Off => None,
+        }
+    }
+
+    pub fn color_code(&self) -> Option<&str> {
+        if env::var(NO_COLOR_ENV).is_ok() || !atty::is(atty::Stream::Stdout) {
+            return None;
+        }
+
+        match self {
+            Level::Trace => None,
+            Level::Debug => Some("\x1b[36m"),
+            Level::Info => Some("\x1b[35m"),
+            Level::Warn => Some("\x1b[33m"),
+            Level::Error => Some("\x1b[31m"),
+            Level::Critical => Some("\x1b[31m"),
+            Level::Off => None,
+        }
+    }
+
+    pub fn to_u8(self) -> u8 {
+        self as u8
+    }
+}
+
+pub fn new_tunnel_prefix() -> String {
+    format!("[tunnel.{}]", next_counter())
+}
+
+pub fn new_code_server_prefix() -> String {
+    format!("[codeserver.{}]", next_counter())
+}
+
+pub fn new_rpc_prefix() -> String {
+    format!("[rpc.{}]", next_counter())
+}
+
+// Base logger implementation
+#[derive(Clone)]
+pub struct Logger {
+    tracer: Tracer,
+    sink: Vec<Box<dyn LogSink>>,
+    prefix: Option<String>,
+}
+
+// Copy trick from https://stackoverflow.com/a/30353928
+pub trait LogSinkClone {
+    fn clone_box(&self) -> Box<dyn LogSink>;
+}
+
+impl<T> LogSinkClone for T
+where
+    T: 'static + LogSink + Clone,
+{
+    fn clone_box(&self) -> Box<dyn LogSink> {
+        Box::new(self.clone())
+    }
+}
+
+pub trait LogSink: LogSinkClone + Sync + Send {
+    fn write_log(&self, level: Level, prefix: &str, message: &str);
+    fn write_result(&self, message: &str);
+}
+
+impl Clone for Box<dyn LogSink> {
+    fn clone(&self) -> Box<dyn LogSink> {
+        self.clone_box()
+    }
+}
+
+#[derive(Clone)]
+pub struct StdioLogSink {
+    level: Level,
+}
+
+impl LogSink for StdioLogSink {
+    fn write_log(&self, level: Level, prefix: &str, message: &str) {
+        if level < self.level {
+            return;
+        }
+
+        emit(level, prefix, message);
+    }
+
+    fn write_result(&self, message: &str) {
+        println!("{}", message);
+    }
+}
+
+#[derive(Clone)]
+pub struct FileLogSink {
+    level: Level,
+    file: Arc<std::sync::Mutex<std::fs::File>>,
+}
+
+impl FileLogSink {
+    pub fn new(level: Level, path: &Path) -> std::io::Result<Self> {
+        let file = std::fs::File::create(path)?;
+        Ok(Self {
+            level,
+            file: Arc::new(std::sync::Mutex::new(file)),
+        })
+    }
+}
+
+impl LogSink for FileLogSink {
+    fn write_log(&self, level: Level, prefix: &str, message: &str) {
+        if level < self.level {
+            return;
+        }
+
+        let line = format(level, prefix, message);
+
+        // ignore any errors, not much we can do if logging fails...
+        self.file.lock().unwrap().write_all(line.as_bytes()).ok();
+    }
+
+    fn write_result(&self, _message: &str) {}
+}
+
+impl Logger {
+    pub fn new(tracer: Tracer, level: Level) -> Self {
+        Self {
+            tracer,
+            sink: vec![Box::new(StdioLogSink { level })],
+            prefix: None,
+        }
+    }
+
+    pub fn span(&self, name: &str) -> SpanBuilder {
+        self.tracer.span_builder(format!("serverlauncher/{}", name))
+    }
+
+    pub fn tracer(&self) -> &Tracer {
+        &self.tracer
+    }
+
+    pub fn emit(&self, level: Level, message: &str) {
+        let prefix = self.prefix.as_deref().unwrap_or("");
+        for sink in &self.sink {
+            sink.write_log(level, prefix, message);
+        }
+    }
+
+    pub fn result(&self, message: &str) {
+        for sink in &self.sink {
+            sink.write_result(message);
+        }
+    }
+
+    pub fn prefixed(&self, prefix: &str) -> Logger {
+        Logger {
+            prefix: Some(match &self.prefix {
+                Some(p) => format!("{}{} ", p, prefix),
+                None => format!("{} ", prefix),
+            }),
+            ..self.clone()
+        }
+    }
+
+    /// Creates a new logger with the additional log sink added.
+    pub fn tee<T>(&self, sink: T) -> Logger
+    where
+        T: LogSink + 'static,
+    {
+        let mut new_sinks = self.sink.clone();
+        new_sinks.push(Box::new(sink));
+
+        Logger {
+            sink: new_sinks,
+            ..self.clone()
+        }
+    }
+
+    pub fn get_download_logger<'a>(&'a self, prefix: &'static str) -> DownloadLogger<'a> {
+        DownloadLogger {
+            prefix,
+            logger: self,
+        }
+    }
+}
+
+pub struct DownloadLogger<'a> {
+    prefix: &'static str,
+    logger: &'a Logger,
+}
+
+impl<'a> crate::util::io::ReportCopyProgress for DownloadLogger<'a> {
+    fn report_progress(&mut self, bytes_so_far: u64, total_bytes: u64) {
+        if total_bytes > 0 {
+            self.logger.emit(
+                Level::Trace,
+                &format!(
+                    "{} {}/{} ({:.0}%)",
+                    self.prefix,
+                    bytes_so_far,
+                    total_bytes,
+                    (bytes_so_far as f64 / total_bytes as f64) * 100.0,
+                ),
+            );
+        } else {
+            self.logger.emit(
+                Level::Trace,
+                &format!("{} {}/{}", self.prefix, bytes_so_far, total_bytes,),
+            );
+        }
+    }
+}
+
+pub fn format(level: Level, prefix: &str, message: &str) -> String {
+    let current = Local::now();
+    let timestamp = current.format("%Y-%m-%d %H:%M:%S").to_string();
+
+    let name = level.name().unwrap();
+
+    if let Some(c) = level.color_code() {
+        format!(
+            "\x1b[2m[{}]\x1b[0m {}{}\x1b[0m {}{}\n",
+            timestamp, c, name, prefix, message
+        )
+    } else {
+        format!("[{}] {} {}{}\n", timestamp, name, prefix, message)
+    }
+}
+
+pub fn emit(level: Level, prefix: &str, message: &str) {
+    let line = format(level, prefix, message);
+    if level == Level::Trace {
+        print!("\x1b[2m{}\x1b[0m", line);
+    } else {
+        print!("{}", line);
+    }
+}
+
+#[macro_export]
+macro_rules! error {
+    ($logger:expr, $str:expr) => {
+        $logger.emit(log::Level::Error, $str)
+     };
+     ($logger:expr, $($fmt:expr),+) => {
+        $logger.emit(log::Level::Error, &format!($($fmt),+))
+     };
+ }
+
+#[macro_export]
+macro_rules! trace {
+     ($logger:expr, $str:expr) => {
+         $logger.emit(log::Level::Trace, $str)
+     };
+     ($logger:expr, $($fmt:expr),+) => {
+         $logger.emit(log::Level::Trace, &format!($($fmt),+))
+     };
+ }
+
+#[macro_export]
+macro_rules! debug {
+     ($logger:expr, $str:expr) => {
+         $logger.emit(log::Level::Debug, $str)
+     };
+     ($logger:expr, $($fmt:expr),+) => {
+         $logger.emit(log::Level::Debug, &format!($($fmt),+))
+     };
+ }
+
+#[macro_export]
+macro_rules! info {
+     ($logger:expr, $str:expr) => {
+         $logger.emit(log::Level::Info, $str)
+     };
+     ($logger:expr, $($fmt:expr),+) => {
+         $logger.emit(log::Level::Info, &format!($($fmt),+))
+     };
+ }
+
+#[macro_export]
+macro_rules! warning {
+     ($logger:expr, $str:expr) => {
+         $logger.emit(log::Level::Warn, $str)
+     };
+     ($logger:expr, $($fmt:expr),+) => {
+         $logger.emit(log::Level::Warn, &format!($($fmt),+))
+     };
+ }
+
+#[macro_export]
+macro_rules! span {
+    ($logger:expr, $span:expr, $func:expr) => {{
+        use opentelemetry::trace::TraceContextExt;
+
+        let span = $span.start($logger.tracer());
+        let cx = opentelemetry::Context::current_with_span(span);
+        let guard = cx.clone().attach();
+        let t = $func;
+
+        if let Err(e) = &t {
+            cx.span().record_error(e);
+        }
+
+        std::mem::drop(guard);
+
+        t
+    }};
+}
+
+#[macro_export]
+macro_rules! spanf {
+    ($logger:expr, $span:expr, $func:expr) => {{
+        use opentelemetry::trace::{FutureExt, TraceContextExt};
+
+        let span = $span.start($logger.tracer());
+        let cx = opentelemetry::Context::current_with_span(span);
+        let t = $func.with_context(cx.clone()).await;
+
+        if let Err(e) = &t {
+            cx.span().record_error(e);
+        }
+
+        cx.span().end();
+
+        t
+    }};
+}