From 670bfcc1bc10792f0fe7e06141bad4684c3a6f95 Mon Sep 17 00:00:00 2001
From: Jiahao XU <Jiahao_XU@outlook.com>
Date: Sun, 19 Jun 2022 01:20:46 +1000
Subject: [PATCH] Impl new type `UIThreadLogger`

Signed-off-by: Jiahao XU <Jiahao_XU@outlook.com>
---
 src/helpers.rs                  |  6 +--
 src/helpers/ui_thread_logger.rs | 86 +++++++++++++++++++++++++++++++++
 2 files changed, 89 insertions(+), 3 deletions(-)
 create mode 100644 src/helpers/ui_thread_logger.rs

diff --git a/src/helpers.rs b/src/helpers.rs
index 5d6d350a..f6233a16 100644
--- a/src/helpers.rs
+++ b/src/helpers.rs
@@ -1,6 +1,4 @@
-use std::{
-    path::{Path, PathBuf},
-};
+use std::path::{Path, PathBuf};
 
 use cargo_toml::Manifest;
 use log::debug;
@@ -20,6 +18,8 @@ pub use auto_abort_join_handle::AutoAbortJoinHandle;
 mod ui_thread;
 pub use ui_thread::UIThread;
 
+mod ui_thread_logger;
+
 mod extracter;
 mod readable_rx;
 
diff --git a/src/helpers/ui_thread_logger.rs b/src/helpers/ui_thread_logger.rs
new file mode 100644
index 00000000..9b7db7e5
--- /dev/null
+++ b/src/helpers/ui_thread_logger.rs
@@ -0,0 +1,86 @@
+use std::cell::Cell;
+use std::fmt::Write;
+
+use bytes::BytesMut;
+use log::{set_boxed_logger, set_max_level, LevelFilter, Log, Metadata, Record};
+use tokio::{runtime::Handle, sync::mpsc::Sender};
+
+use super::ui_thread::UIRequest;
+
+#[derive(Debug)]
+pub(super) struct UIThreadLogger {
+    tx: Sender<UIRequest>,
+    level: LevelFilter,
+    filter_ignore: &'static [&'static str],
+}
+
+impl UIThreadLogger {
+    pub(super) fn init(
+        tx: Sender<UIRequest>,
+        level: LevelFilter,
+        filter_ignore: &'static [&'static str],
+    ) {
+        set_max_level(level);
+        set_boxed_logger(Self::new(tx, level, filter_ignore)).unwrap()
+    }
+
+    fn new(
+        tx: Sender<UIRequest>,
+        level: LevelFilter,
+        filter_ignore: &'static [&'static str],
+    ) -> Box<Self> {
+        Box::new(Self {
+            tx,
+            level,
+            filter_ignore,
+        })
+    }
+
+    fn send_request(&self, request: UIRequest) {
+        // TODO: Use another mpsc type.
+        // Tokio's mpsc requires the async send to be used
+        // in async context.
+        if let Ok(handle) = Handle::try_current() {
+            let tx = self.tx.clone();
+            handle.spawn(async move { tx.send(request).await.unwrap() });
+        } else {
+            self.tx.blocking_send(request).unwrap();
+        }
+    }
+
+    thread_local! {
+        static BUFFER: Cell<BytesMut> = Cell::new(BytesMut::new());
+    }
+}
+
+impl Log for UIThreadLogger {
+    fn enabled(&self, metadata: &Metadata<'_>) -> bool {
+        metadata.level() <= self.level
+    }
+
+    fn log(&self, record: &Record<'_>) {
+        let target = record.target();
+        if self.enabled(record.metadata())
+            && !self
+                .filter_ignore
+                .iter()
+                .any(|filter| target.starts_with(filter))
+        {
+            let output = Self::BUFFER.with(|cell| {
+                let mut buffer = cell.take();
+                write!(&mut buffer, "{}", record.args()).unwrap();
+
+                let output = buffer.split().freeze();
+                cell.set(buffer);
+
+                output
+            });
+
+            self.send_request(UIRequest::PrintToStdout(output));
+        }
+    }
+
+    fn flush(&self) {
+        self.send_request(UIRequest::FlushStdout);
+    }
+}