vmbase: Temporarily disable logging in exception handlers

Add `logger::block` which temporarily disables logging. This should be used
to prevent log messages being sent to UART from exception handlers. It
allows calling external code in exception handlers which may produce log
messages. Also stop exporting `println` macros from vmbase for the same
reason and use log macros instead.

Bug: 245267332
Test: atest vmbase_example.integration_test

Change-Id: I99e5ea1fa15e7371fd680e4086a3685ac489111d
diff --git a/vmbase/src/console.rs b/vmbase/src/console.rs
index fabea91..7c8ddf6 100644
--- a/vmbase/src/console.rs
+++ b/vmbase/src/console.rs
@@ -40,14 +40,14 @@
 /// Writes a string to the console.
 ///
 /// Panics if [`init`] was not called first.
-pub fn write_str(s: &str) {
+pub(crate) fn write_str(s: &str) {
     CONSOLE.lock().as_mut().unwrap().write_str(s).unwrap();
 }
 
 /// Writes a formatted string to the console.
 ///
 /// Panics if [`init`] was not called first.
-pub fn write_args(format_args: Arguments) {
+pub(crate) fn write_args(format_args: Arguments) {
     write(CONSOLE.lock().as_mut().unwrap(), format_args).unwrap();
 }
 
@@ -69,20 +69,10 @@
     let _ = write(&mut uart, format_args);
 }
 
-/// Prints the given string to the console.
-///
-/// Panics if the console has not yet been initialised. May hang if used in an exception context;
-/// use `eprint!` instead.
-#[macro_export]
-macro_rules! print {
-    ($($arg:tt)*) => ($crate::console::write_args(format_args!($($arg)*)));
-}
-
 /// Prints the given formatted string to the console, followed by a newline.
 ///
 /// Panics if the console has not yet been initialised. May hang if used in an exception context;
 /// use `eprintln!` instead.
-#[macro_export]
 macro_rules! println {
     () => ($crate::console::write_str("\n"));
     ($($arg:tt)*) => ({
@@ -91,6 +81,8 @@
     );
 }
 
+pub(crate) use println; // Make it available in this crate.
+
 /// Prints the given string to the console in an emergency, such as an exception handler.
 ///
 /// Never panics.
diff --git a/vmbase/src/entry.rs b/vmbase/src/entry.rs
index 1510ae2..8cdfe77 100644
--- a/vmbase/src/entry.rs
+++ b/vmbase/src/entry.rs
@@ -36,12 +36,14 @@
 /// Example:
 ///
 /// ```rust
-/// use vmbase::main;
+/// use vmbase::{logger, main};
+/// use log::{info, LevelFilter};
 ///
 /// main!(my_main);
 ///
 /// fn my_main() {
-///     println!("Hello world");
+///     logger::init(LevelFilter::Info).unwrap();
+///     info!("Hello world");
 /// }
 /// ```
 #[macro_export]
diff --git a/vmbase/src/logger.rs b/vmbase/src/logger.rs
index 5f0f1c2..94dc880 100644
--- a/vmbase/src/logger.rs
+++ b/vmbase/src/logger.rs
@@ -20,27 +20,71 @@
 
 extern crate log;
 
-use super::println;
+use crate::console::println;
+use core::sync::atomic::{AtomicBool, Ordering};
 use log::{LevelFilter, Log, Metadata, Record, SetLoggerError};
 
-struct Logger;
-static LOGGER: Logger = Logger;
+struct Logger {
+    is_enabled: AtomicBool,
+}
+static mut LOGGER: Logger = Logger::new();
+
+impl Logger {
+    const fn new() -> Self {
+        Self { is_enabled: AtomicBool::new(true) }
+    }
+
+    fn swap_enabled(&mut self, enabled: bool) -> bool {
+        self.is_enabled.swap(enabled, Ordering::Relaxed)
+    }
+}
 
 impl Log for Logger {
     fn enabled(&self, _metadata: &Metadata) -> bool {
-        true
+        self.is_enabled.load(Ordering::Relaxed)
     }
 
     fn log(&self, record: &Record) {
-        println!("[{}] {}", record.level(), record.args());
+        if self.enabled(record.metadata()) {
+            println!("[{}] {}", record.level(), record.args());
+        }
     }
 
     fn flush(&self) {}
 }
 
+/// An RAII implementation of a log suppressor. When the instance is dropped, logging is re-enabled.
+pub struct SuppressGuard {
+    old_enabled: bool,
+}
+
+impl SuppressGuard {
+    fn new() -> Self {
+        // Safe because it modifies an atomic.
+        unsafe { Self { old_enabled: LOGGER.swap_enabled(false) } }
+    }
+}
+
+impl Drop for SuppressGuard {
+    fn drop(&mut self) {
+        // Safe because it modifies an atomic.
+        unsafe {
+            LOGGER.swap_enabled(self.old_enabled);
+        }
+    }
+}
+
 /// Initialize vmbase logger with a given max logging level.
 pub fn init(max_level: LevelFilter) -> Result<(), SetLoggerError> {
-    log::set_logger(&LOGGER)?;
+    // Safe because it only sets the global logger.
+    unsafe {
+        log::set_logger(&LOGGER)?;
+    }
     log::set_max_level(max_level);
     Ok(())
 }
+
+/// Suppress logging until the return value goes out of scope.
+pub fn suppress() -> SuppressGuard {
+    SuppressGuard::new()
+}