Create proper logging infrastructure

2 files changed, 496 insertions, 10 deletions

diff --git a/zluda_dump/src/lib.rs b/zluda_dump/src/lib.rs
index 6425768..2887409 100644
--- a/zluda_dump/src/lib.rs
+++ b/zluda_dump/src/lib.rs
@@ -4,10 +4,10 @@ use std::{
     error::Error,
     ffi::{c_void, CStr},
     fs,
-    io::prelude::*,
+    io::{self, prelude::*},
     mem,
     os::raw::{c_int, c_uint, c_ulong, c_ushort},
-    path::PathBuf,
+    path::{Path, PathBuf},
     rc::Rc,
     slice,
 };
@@ -20,10 +20,6 @@ use cuda::{
 use ptx::ast;
 use regex::Regex;
 
-#[cfg_attr(windows, path = "os_win.rs")]
-#[cfg_attr(not(windows), path = "os_unix.rs")]
-mod os;
-
 const CU_LAUNCH_PARAM_END: *mut c_void = 0 as *mut _;
 const CU_LAUNCH_PARAM_BUFFER_POINTER: *mut c_void = 1 as *mut _;
 const CU_LAUNCH_PARAM_BUFFER_SIZE: *mut c_void = 2 as *mut _;
@@ -32,7 +28,7 @@ macro_rules! extern_redirect {
     (pub fn $fn_name:ident ( $($arg_id:ident: $arg_type:ty),* $(,)? ) -> $ret_type:ty ;) => {
         #[no_mangle]
         pub extern "system" fn $fn_name ( $( $arg_id : $arg_type),* ) -> $ret_type {
-            unsafe { $crate::init_libcuda_handle() };
+            unsafe { $crate::init_libcuda_handle(stringify!($fn_name)) };
             let name = std::ffi::CString::new(stringify!($fn_name)).unwrap();
             let fn_ptr = unsafe { crate::os::get_proc_address($crate::LIBCUDA_HANDLE, &name) };
             if fn_ptr == std::ptr::null_mut() {
@@ -51,7 +47,7 @@ macro_rules! extern_redirect_with {
     ) => {
         #[no_mangle]
         pub extern "system" fn $fn_name ( $( $arg_id : $arg_type),* ) -> $ret_type {
-            unsafe { $crate::init_libcuda_handle() };
+            unsafe { $crate::init_libcuda_handle(stringify!($fn_name)) };
             let continuation = |$( $arg_id : $arg_type),* | {
                 let name = std::ffi::CString::new(stringify!($fn_name)).unwrap();
                 let fn_ptr = unsafe { crate::os::get_proc_address($crate::LIBCUDA_HANDLE, &name) };
@@ -68,6 +64,10 @@ macro_rules! extern_redirect_with {
 
 #[allow(warnings)]
 mod cuda;
+mod log;
+#[cfg_attr(windows, path = "os_win.rs")]
+#[cfg_attr(not(windows), path = "os_unix.rs")]
+mod os;
 
 pub static mut LIBCUDA_HANDLE: *mut c_void = ptr::null_mut();
 pub static mut PENDING_LINKING: Option<HashMap<CUlinkState, Vec<ModuleDump>>> = None;
@@ -81,6 +81,42 @@ pub static mut KERNEL_PATTERN: Option<Regex> = None;
 pub static mut OVERRIDE_COMPUTE_CAPABILITY_MAJOR: Option<i32> = None;
 pub static mut KERNEL_INDEX_MINIMUM: usize = 0;
 pub static mut KERNEL_INDEX_MAXIMUM: usize = usize::MAX;
+pub(crate) static mut LOG_FACTORY: Option<log::Factory> = None;
+
+pub(crate) struct Settings {
+    dump_dir: Option<PathBuf>,
+}
+
+impl Settings {
+    fn read_and_init(logger: &mut log::FunctionLogger) -> Self {
+        let maybe_dump_dir = Self::read_and_init_dump_dir();
+        let dump_dir = match maybe_dump_dir {
+            Ok(d) => d,
+            Err(err) => {
+                logger.log(log::LogEntry::ErrorBox(err));
+                None
+            }
+        };
+        Settings { dump_dir }
+    }
+
+    fn read_and_init_dump_dir() -> Result<Option<PathBuf>, Box<dyn Error>> {
+        let dir = match env::var("ZLUDA_DUMP_DIR") {
+            Ok(dir) => dir,
+            Err(env::VarError::NotPresent) => return Ok(None),
+            Err(err) => return Err(Box::new(err) as Box<_>),
+        };
+        Ok(Some(Self::create_dump_directory(dir)?))
+    }
+
+    fn create_dump_directory(dir: String) -> io::Result<PathBuf> {
+        let mut main_dir = PathBuf::from(dir);
+        let current_exe = env::current_exe()?;
+        main_dir.push(current_exe.file_name().unwrap());
+        fs::create_dir_all(&main_dir)?;
+        Ok(main_dir)
+    }
+}
 
 #[derive(Clone, Copy)]
 enum AllocLocation {
@@ -102,8 +138,10 @@ pub struct KernelDump {
 
 // We are doing dlopen here instead of just using LD_PRELOAD,
 // it's because CUDA Runtime API does dlopen to open libcuda.so, which ignores LD_PRELOAD
-pub unsafe fn init_libcuda_handle() {
+pub unsafe fn init_libcuda_handle(func: &'static str) {
     if LIBCUDA_HANDLE == ptr::null_mut() {
+        let mut log_factory = log::Factory::new();
+        let (logger, settings) = log_factory.get_first_logger_and_init_settings(func);
         MODULES = Some(HashMap::new());
         KERNELS = Some(HashMap::new());
         BUFFERS = Some(BTreeMap::new());
@@ -147,7 +185,7 @@ pub unsafe fn init_libcuda_handle() {
             },
             Err(_) => (),
         }
-        os_log!("Initialized");
+        drop(logger);
     }
 }
 
diff --git a/zluda_dump/src/log.rs b/zluda_dump/src/log.rs
new file mode 100644
index 0000000..2b9ef91
--- /dev/null
+++ b/zluda_dump/src/log.rs
@@ -0,0 +1,448 @@
+use super::CUresult;

+use super::Settings;

+use std::error::Error;

+use std::fmt::Display;

+use std::fs::File;

+use std::io;

+use std::io::Stderr;

+use std::io::Write;

+

+const LOG_PREFIX: &[u8] = b"[ZLUDA_DUMP] ";

+

+// This type holds all the relevant settings for logging like output path and

+// creates objects which match those settings

+pub(crate) struct Factory {

+    // Fallible emitter is optional emitter to file system, we might lack

+    // file permissions or be out of disk space

+    fallible_emitter: Option<Box<dyn WriteTrailingZeroAware>>,

+    // This is emitter that "always works" (and if it does not, then we don't

+    // care). In addition of normal logs it emits errors from fallible emitter

+    infallible_emitter: Box<dyn WriteTrailingZeroAware>,

+    write_buffer: WriteBuffer,

+    // another shared buffer, so we dont't reallocate on every function call

+    log_queue: Vec<LogEntry>,

+}

+

+// When writing out to the emitter (file, WinAPI, whatever else) instead of

+// writing piece-by-piece it's better to first concatenate everything in memory

+// then write out from memory to the slow emitter only once.

+// Additionally we might have an unprefixed and prefixed buffer, this struct

+// handles this detail

+struct WriteBuffer {

+    prefixed_buffer: Option<Vec<u8>>,

+    unprefixed_buffer: Option<Vec<u8>>,

+}

+

+impl WriteBuffer {

+    fn new() -> Self {

+        WriteBuffer {

+            prefixed_buffer: None,

+            unprefixed_buffer: None,

+        }

+    }

+

+    fn init(

+        &mut self,

+        fallible_emitter: &Option<Box<dyn WriteTrailingZeroAware>>,

+        infallible_emitter: &Box<dyn WriteTrailingZeroAware>,

+    ) {

+        if infallible_emitter.should_prefix() {

+            self.prefixed_buffer = Some(Vec::new());

+        } else {

+            self.unprefixed_buffer = Some(Vec::new());

+        }

+        if let Some(emitter) = fallible_emitter {

+            if emitter.should_prefix() {

+                self.prefixed_buffer = Some(Vec::new());

+            } else {

+                self.unprefixed_buffer = Some(Vec::new());

+            }

+        }

+    }

+

+    fn all_buffers(&mut self) -> impl Iterator<Item = &mut Vec<u8>> {

+        self.prefixed_buffer

+            .as_mut()

+            .into_iter()

+            .chain(self.unprefixed_buffer.as_mut().into_iter())

+    }

+

+    fn start_line(&mut self) {

+        if let Some(buffer) = &mut self.prefixed_buffer {

+            buffer.extend_from_slice(LOG_PREFIX);

+        }

+    }

+

+    fn end_line(&mut self) {

+        for buffer in self.all_buffers() {

+            buffer.push(b'\n');

+        }

+    }

+

+    fn write(&mut self, s: &str) {

+        for buffer in self.all_buffers() {

+            buffer.extend_from_slice(s.as_bytes());

+        }

+    }

+

+    fn finish(&mut self) {

+        for buffer in self.all_buffers() {

+            buffer.push(b'\0');

+        }

+    }

+

+    fn undo_finish(&mut self) {

+        for buffer in self.all_buffers() {

+            buffer.truncate(buffer.len() - 1);

+        }

+    }

+

+    fn send_to(&self, log_emitter: &mut Box<dyn WriteTrailingZeroAware>) -> Result<(), io::Error> {

+        if log_emitter.should_prefix() {

+            log_emitter.write_zero_aware(

+                &*self

+                    .prefixed_buffer

+                    .as_ref()

+                    .unwrap_or_else(|| unreachable!()),

+            )

+        } else {

+            log_emitter.write_zero_aware(

+                &*self

+                    .unprefixed_buffer

+                    .as_ref()

+                    .unwrap_or_else(|| unreachable!()),

+            )

+        }

+    }

+

+    fn reset(&mut self) {

+        for buffer in self.all_buffers() {

+            unsafe { buffer.set_len(0) };

+        }

+    }

+}

+

+impl Write for WriteBuffer {

+    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {

+        if let Some(buffer) = &mut self.prefixed_buffer {

+            buffer.extend_from_slice(buf);

+        }

+        if let Some(buffer) = &mut self.unprefixed_buffer {

+            buffer.extend_from_slice(buf);

+        }

+        Ok(buf.len())

+    }

+

+    fn flush(&mut self) -> io::Result<()> {

+        Ok(())

+    }

+}

+

+impl Factory {

+    pub(crate) fn new() -> Self {

+        let debug_emitter = os::new_debug_logger();

+        Factory {

+            infallible_emitter: debug_emitter,

+            fallible_emitter: None,

+            write_buffer: WriteBuffer::new(),

+            log_queue: Vec::new(),

+        }

+    }

+

+    fn initalize_fallible_emitter(

+        settings: &Settings,

+    ) -> std::io::Result<Option<Box<dyn WriteTrailingZeroAware>>> {

+        settings

+            .dump_dir

+            .as_ref()

+            .map(|path| {

+                Ok::<_, std::io::Error>(Box::new(File::create(path.to_path_buf().join("log.txt"))?)

+                    as Box<dyn WriteTrailingZeroAware>)

+            })

+            .transpose()

+    }

+

+    // We load settings during first function call, since during that time we

+    // also create one of the loggers, what do we do about errors encountered

+    // at that time? We log them to the newly created logger, but to make it

+    // "nice" we do both of those in a single function

+    // An alternative would be to have something like this:

+    //   let mut factory = Factory::new();

+    //   let mut cuInitLog = factory.get_logger("cuInit");

+    //   cuInitLog.load_settings(&settings);

+    // which is a bit nonsensical

+    pub(crate) fn get_first_logger_and_init_settings(

+        &mut self,

+        func: &'static str,

+    ) -> (FunctionLogger, Settings) {

+        let mut first_logger = self.get_logger(func);

+        let settings = Settings::read_and_init(&mut first_logger);

+        match Self::initalize_fallible_emitter(&settings) {

+            Ok(fallible_emitter) => {

+                *first_logger.fallible_emitter = fallible_emitter;

+            }

+            Err(err) => first_logger.log(LogEntry::IoError(err)),

+        }

+        first_logger.write_buffer.init(

+            first_logger.fallible_emitter,

+            first_logger.infallible_emitter,

+        );

+        (first_logger, settings)

+    }

+

+    pub(crate) fn get_logger(&mut self, func: &'static str) -> FunctionLogger {

+        FunctionLogger {

+            result: None,

+            name: func,

+            fallible_emitter: &mut self.fallible_emitter,

+            infallible_emitter: &mut self.infallible_emitter,

+            write_buffer: &mut self.write_buffer,

+            log_queue: &mut self.log_queue,

+        }

+    }

+}

+

+// This encapsulates log output for a single function call.

+// It's a separate struct and not just a plain function for two reasons:

+// * While we want to always display return code before logging errors,

+//   logging errors might come before return code is returned

+// * We want to handle panics gracefully with Drop

+pub(crate) struct FunctionLogger<'a> {

+    pub(crate) result: Option<CUresult>,

+    name: &'static str,

+    infallible_emitter: &'a mut Box<dyn WriteTrailingZeroAware>,

+    fallible_emitter: &'a mut Option<Box<dyn WriteTrailingZeroAware>>,

+    write_buffer: &'a mut WriteBuffer,

+    log_queue: &'a mut Vec<LogEntry>,

+}

+

+impl<'a> FunctionLogger<'a> {

+    pub(crate) fn log(&mut self, l: LogEntry) {

+        self.log_queue.push(l);

+    }

+

+    fn flush_log_queue_to_write_buffer(&mut self) {

+        self.write_buffer.start_line();

+        self.write_buffer.write(self.name);

+        self.write_buffer.write("(...) -> ");

+        if let Some(result) = self.result {

+            write!(self.write_buffer, "{:#X}", result.0).unwrap_or_else(|_| unreachable!());

+        } else {

+            self.write_buffer.write("(INTERNAL ERROR)");

+        };

+        self.write_buffer.end_line();

+        for entry in self.log_queue.iter() {

+            write!(self.write_buffer, "    {}", entry).unwrap_or_else(|_| unreachable!());

+            self.write_buffer.end_line();

+        }

+        self.write_buffer.finish();

+    }

+

+    // This is a dirty hack: we call it at the point where our write buffer is

+    // already finalized and squeeze the error produced by the previous emitter

+    fn hack_squeeze_in_additional_error(&mut self, entry: LogEntry) {

+        self.write_buffer.undo_finish();

+        write!(self.write_buffer, "    {}", entry).unwrap_or_else(|_| unreachable!());

+        self.write_buffer.end_line();

+        self.write_buffer.finish();

+    }

+}

+

+impl<'a> Drop for FunctionLogger<'a> {

+    fn drop(&mut self) {

+        self.flush_log_queue_to_write_buffer();

+        let error_from_writing_to_fallible_emitter = match self.fallible_emitter {

+            Some(emitter) => self.write_buffer.send_to(emitter),

+            None => Ok(()),

+        };

+        if let Err(e) = error_from_writing_to_fallible_emitter {

+            self.hack_squeeze_in_additional_error(LogEntry::IoError(e))

+        }

+        self.write_buffer.send_to(self.infallible_emitter).ok();

+        self.write_buffer.reset();

+        self.log_queue.truncate(0);

+    }

+}

+

+// Structured log type. We don't want frontend to care about log formatting

+pub(crate) enum LogEntry {

+    IoError(io::Error),

+    ErrorBox(Box<dyn Error>),

+}

+

+impl Display for LogEntry {

+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

+        match self {

+            LogEntry::IoError(e) => e.fmt(f),

+            LogEntry::ErrorBox(e) => e.fmt(f),

+        }

+    }

+}

+

+// Some of our writers want to have trailing zero (WinAPI debug logger) and some

+// don't (everything else), this trait encapsulates that logic

+pub(crate) trait WriteTrailingZeroAware {

+    fn write_zero_aware(&mut self, buf: &[u8]) -> std::io::Result<()>;

+    fn flush(&mut self) -> std::io::Result<()>;

+    fn should_prefix(&self) -> bool;

+}

+

+impl WriteTrailingZeroAware for File {

+    fn write_zero_aware(&mut self, buf: &[u8]) -> std::io::Result<()> {

+        <Self as std::io::Write>::write_all(self, buf.split_last().unwrap().1)

+    }

+

+    fn flush(&mut self) -> std::io::Result<()> {

+        <Self as std::io::Write>::flush(self)

+    }

+

+    fn should_prefix(&self) -> bool {

+        false

+    }

+}

+

+impl WriteTrailingZeroAware for Stderr {

+    fn write_zero_aware(&mut self, buf: &[u8]) -> std::io::Result<()> {

+        <Self as std::io::Write>::write_all(self, buf.split_last().unwrap().1)

+    }

+

+    fn flush(&mut self) -> std::io::Result<()> {

+        <Self as std::io::Write>::flush(self)

+    }

+

+    fn should_prefix(&self) -> bool {

+        true

+    }

+}

+

+#[cfg(windows)]

+mod os {

+    use super::WriteTrailingZeroAware;

+    use std::{os::windows::prelude::AsRawHandle, ptr};

+    use winapi::um::debugapi::OutputDebugStringA;

+

+    struct OutputDebugString {}

+

+    impl WriteTrailingZeroAware for OutputDebugString {

+        fn write_zero_aware(&mut self, buf: &[u8]) -> std::io::Result<()> {

+            unsafe { OutputDebugStringA(buf.as_ptr() as *const _) };

+            Ok(())

+        }

+

+        fn flush(&mut self) -> std::io::Result<()> {

+            Ok(())

+        }

+

+        fn should_prefix(&self) -> bool {

+            true

+        }

+    }

+

+    pub(crate) fn new_debug_logger() -> Box<dyn WriteTrailingZeroAware> {

+        let stderr = std::io::stderr();

+        let log_to_stderr = stderr.as_raw_handle() != ptr::null_mut();

+        if log_to_stderr {

+            Box::new(stderr)

+        } else {

+            Box::new(OutputDebugString {})

+        }

+    }

+}

+

+#[cfg(not(windows))]

+mod os {

+    use super::WriteTrailingZeroAware;

+

+    pub(crate) fn new_debug_logger() -> Box<dyn WriteTrailingZeroAware> {

+        Box::new(std::io::stderr())

+    }

+}

+

+#[cfg(test)]

+mod tests {

+    use std::{cell::RefCell, io, rc::Rc, str};

+

+    use super::{FunctionLogger, LogEntry, WriteTrailingZeroAware};

+    use crate::{log::WriteBuffer, CUresult};

+

+    struct FailOnNthWrite {

+        fail_on: usize,

+        counter: usize,

+    }

+

+    impl WriteTrailingZeroAware for FailOnNthWrite {

+        fn write_zero_aware(&mut self, _: &[u8]) -> std::io::Result<()> {

+            self.counter += 1;

+            if self.counter >= self.fail_on {

+                Err(io::Error::from_raw_os_error(4))

+            } else {

+                Ok(())

+            }

+        }

+

+        fn flush(&mut self) -> std::io::Result<()> {

+            panic!()

+        }

+

+        fn should_prefix(&self) -> bool {

+            false

+        }

+    }

+

+    // Custom type to not trigger trait coherence rules

+    #[derive(Clone)]

+    struct RcVec<T>(Rc<RefCell<Vec<T>>>);

+

+    impl WriteTrailingZeroAware for RcVec<u8> {

+        fn write_zero_aware(&mut self, buf: &[u8]) -> std::io::Result<()> {

+            let mut vec = self.0.borrow_mut();

+            vec.extend_from_slice(buf.split_last().unwrap().1);

+            Ok(())

+        }

+

+        fn flush(&mut self) -> std::io::Result<()> {

+            Ok(())

+        }

+

+        fn should_prefix(&self) -> bool {

+            false

+        }

+    }

+

+    #[test]

+    fn error_in_fallible_emitter_is_handled_gracefully() {

+        let result = RcVec(Rc::new(RefCell::new(Vec::<u8>::new())));

+        let mut infallible_emitter = Box::new(result.clone()) as Box<dyn WriteTrailingZeroAware>;

+        let mut fallible_emitter = Some(Box::new(FailOnNthWrite {

+            fail_on: 1,

+            counter: 0,

+        }) as Box<dyn WriteTrailingZeroAware>);

+        let mut write_buffer = WriteBuffer::new();

+        write_buffer.unprefixed_buffer = Some(Vec::new());

+        let mut log_queue = Vec::new();

+        let mut func_logger = FunctionLogger {

+            result: Some(CUresult::CUDA_SUCCESS),

+            name: "cuInit",

+            infallible_emitter: &mut infallible_emitter,

+            fallible_emitter: &mut fallible_emitter,

+            write_buffer: &mut write_buffer,

+            log_queue: &mut log_queue,

+        };

+

+        func_logger.log(LogEntry::IoError(io::Error::from_raw_os_error(1)));

+        func_logger.log(LogEntry::IoError(io::Error::from_raw_os_error(2)));

+        func_logger.log(LogEntry::IoError(io::Error::from_raw_os_error(3)));

+        drop(func_logger);

+        drop(infallible_emitter);

+

+        let result = result.0.borrow_mut();

+        let result_str = str::from_utf8(&*result).unwrap();

+        let result_lines = result_str.lines().collect::<Vec<_>>();

+        assert_eq!(result_lines.len(), 5);

+        assert_eq!(result_lines[0], "cuInit(...) -> 0x0");

+        assert!(result_lines[1].starts_with("    "));

+        assert!(result_lines[2].starts_with("    "));

+        assert!(result_lines[3].starts_with("    "));

+        assert!(result_lines[4].starts_with("    "));

+    }

+}