Merge pull request #789 from vlovich/fix-log-panic

Fix panic introduced upstream if debug log level is disabled

Author: MarcusDunn

Cargo.lock

@@ -17,6 +17,7 @@ tracing-core = { workspace = true }
 
 [dev-dependencies]
 encoding_rs = { workspace = true }
+tracing-subscriber = { workspace = true }
 
 [features]
 default = ["openmp", "android-shared-stdcxx"]

llama-cpp-2/Cargo.toml

@@ -375,19 +375,20 @@ extern "C" fn logs_to_trace(
 
     let log_state = unsafe { &*(data as *const log::State) };
 
+    if log_state.options.disabled {
+        return;
+    }
+
     // If the log level is disabled, we can just return early
     if !log_state.is_enabled_for_level(level) {
+        log_state.update_previous_level_for_disabled_log(level);
         return;
     }
 
     let text = unsafe { std::ffi::CStr::from_ptr(text) };
     let text = text.to_string_lossy();
     let text: &str = text.borrow();
 
-    if log_state.options.disabled {
-        return;
-    }
-
     // As best I can tell llama.cpp / ggml require all log format strings at call sites to have the '\n'.
     // If it's missing, it means that you expect more logs via CONT (or there's a typo in the codebase). To
     // distinguish typo from intentional support for CONT, we have to buffer until the next message comes in

llama-cpp-2/src/lib.rs

@@ -252,16 +252,140 @@ impl State {
         }
     }
 
+    pub(super) fn update_previous_level_for_disabled_log(
+        &self,
+        level: llama_cpp_sys_2::ggml_log_level,
+    ) {
+        if level != llama_cpp_sys_2::GGML_LOG_LEVEL_CONT {
+            self.previous_level
+                .store(level as i32, std::sync::atomic::Ordering::Release);
+        }
+    }
+
     /// Checks whether the given log level is enabled by the current tracing subscriber.
     pub(super) fn is_enabled_for_level(&self, level: llama_cpp_sys_2::ggml_log_level) -> bool {
         // CONT logs do not need to check if they are enabled.
-        if level == llama_cpp_sys_2::GGML_LOG_LEVEL_CONT {
-            return true;
-        }
+        let level = if level == llama_cpp_sys_2::GGML_LOG_LEVEL_CONT {
+            self.previous_level
+                .load(std::sync::atomic::Ordering::Relaxed)
+                as llama_cpp_sys_2::ggml_log_level
+        } else {
+            level
+        };
         let (meta, _) = meta_for_level(level);
         tracing::dispatcher::get_default(|dispatcher| dispatcher.enabled(meta))
     }
 }
 
 pub(super) static LLAMA_STATE: OnceLock<Box<State>> = OnceLock::new();
 pub(super) static GGML_STATE: OnceLock<Box<State>> = OnceLock::new();
+
+#[cfg(test)]
+mod tests {
+    use crate::logs_to_trace;
+    use std::sync::{Arc, Mutex};
+    use tracing::subscriber::DefaultGuard;
+    use tracing_subscriber::util::SubscriberInitExt;
+
+    use super::*;
+
+    struct Logger {
+        #[allow(unused)]
+        guard: DefaultGuard,
+        logs: Arc<Mutex<Vec<String>>>,
+    }
+
+    #[derive(Clone)]
+    struct VecWriter(Arc<Mutex<Vec<String>>>);
+
+    impl std::io::Write for VecWriter {
+        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
+            let log_line = String::from_utf8(buf.to_vec()).map_err(|_| {
+                std::io::Error::new(std::io::ErrorKind::InvalidData, "Invalid UTF-8")
+            })?;
+            self.0.lock().unwrap().push(log_line);
+            Ok(buf.len())
+        }
+
+        fn flush(&mut self) -> std::io::Result<()> {
+            Ok(())
+        }
+    }
+
+    fn create_logger(max_level: tracing::Level) -> Logger {
+        let logs = Arc::new(Mutex::new(vec![]));
+        let writer = VecWriter(logs.clone());
+
+        Logger {
+            guard: tracing_subscriber::fmt()
+                .with_max_level(max_level)
+                .with_ansi(false)
+                .without_time()
+                .with_file(false)
+                .with_line_number(false)
+                .with_level(false)
+                .with_target(false)
+                .with_writer(move || writer.clone())
+                .finish()
+                .set_default(),
+            logs,
+        }
+    }
+
+    #[test]
+    fn cont_disabled_log() {
+        let logger = create_logger(tracing::Level::INFO);
+        let mut log_state = Box::new(State::new(Module::LlamaCpp, LogOptions::default()));
+        let log_ptr = log_state.as_mut() as *mut State as *mut std::os::raw::c_void;
+
+        logs_to_trace(
+            llama_cpp_sys_2::GGML_LOG_LEVEL_DEBUG,
+            c"Hello ".as_ptr(),
+            log_ptr,
+        );
+        logs_to_trace(
+            llama_cpp_sys_2::GGML_LOG_LEVEL_CONT,
+            c"world\n".as_ptr(),
+            log_ptr,
+        );
+
+        assert!(logger.logs.lock().unwrap().is_empty());
+
+        logs_to_trace(
+            llama_cpp_sys_2::GGML_LOG_LEVEL_DEBUG,
+            c"Hello ".as_ptr(),
+            log_ptr,
+        );
+        logs_to_trace(
+            llama_cpp_sys_2::GGML_LOG_LEVEL_CONT,
+            c"world".as_ptr(),
+            log_ptr,
+        );
+        logs_to_trace(
+            llama_cpp_sys_2::GGML_LOG_LEVEL_CONT,
+            c"\n".as_ptr(),
+            log_ptr,
+        );
+    }
+
+    #[test]
+    fn cont_enabled_log() {
+        let logger = create_logger(tracing::Level::INFO);
+        let mut log_state = Box::new(State::new(Module::LlamaCpp, LogOptions::default()));
+        let log_ptr = log_state.as_mut() as *mut State as *mut std::os::raw::c_void;
+
+        logs_to_trace(
+            llama_cpp_sys_2::GGML_LOG_LEVEL_INFO,
+            c"Hello ".as_ptr(),
+            log_ptr,
+        );
+        logs_to_trace(
+            llama_cpp_sys_2::GGML_LOG_LEVEL_CONT,
+            c"world\n".as_ptr(),
+            log_ptr,
+        );
+
+        // Not sure where the extra \n comes from.
+        assert_eq!(*logger.logs.lock().unwrap(), vec!["Hello world\n\n"]);
+    }
+}