/*--------------------------------------------------------------------------------------------- * 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, TracerProvider}, trace::{SpanBuilder, Tracer as TraitTracer, TracerProvider as TracerProviderTrait}, }; 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>, prefix: Option, } // Copy trick from https://stackoverflow.com/a/30353928 pub trait LogSinkClone { fn clone_box(&self) -> Box; } impl LogSinkClone for T where T: 'static + LogSink + Clone, { fn clone_box(&self) -> Box { 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 { fn clone(&self) -> Box { 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>, } impl FileLogSink { pub fn new(level: Level, path: &Path) -> std::io::Result { 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 test() -> Self { Self { tracer: TracerProvider::builder().build().tracer("codeclitest"), sink: vec![], prefix: None, } } 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: impl AsRef) { for sink in &self.sink { sink.write_result(message.as_ref()); } } 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(&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 }}; }