diff --git a/program-runtime/src/log_collector.rs b/program-runtime/src/log_collector.rs index 0618291db2..1cb09259c1 100644 --- a/program-runtime/src/log_collector.rs +++ b/program-runtime/src/log_collector.rs @@ -3,17 +3,36 @@ use std::{cell::RefCell, rc::Rc}; const LOG_MESSAGES_BYTES_LIMIT: usize = 10 * 1000; -#[derive(Default)] pub struct LogCollector { messages: Vec, bytes_written: usize, + bytes_limit: Option, limit_warning: bool, } +impl Default for LogCollector { + fn default() -> Self { + Self { + messages: Vec::new(), + bytes_written: 0, + bytes_limit: Some(LOG_MESSAGES_BYTES_LIMIT), + limit_warning: false, + } + } +} + impl LogCollector { pub fn log(&mut self, message: &str) { + let limit = match self.bytes_limit { + Some(limit) => limit, + None => { + self.messages.push(message.to_string()); + return; + } + }; + let bytes_written = self.bytes_written.saturating_add(message.len()); - if bytes_written >= LOG_MESSAGES_BYTES_LIMIT { + if bytes_written >= limit { if !self.limit_warning { self.limit_warning = true; self.messages.push(String::from("Log truncated")); @@ -31,6 +50,13 @@ impl LogCollector { pub fn new_ref() -> Rc> { Rc::new(RefCell::new(Self::default())) } + + pub fn new_ref_with_limit(bytes_limit: Option) -> Rc> { + Rc::new(RefCell::new(Self { + bytes_limit, + ..Self::default() + })) + } } impl From for Vec {