use btreemap for log storage in upgrader and limit to 25k items

Cargo.toml

@@ -57,7 +57,7 @@ ic-cdk = "0.16.0"
 ic-cdk-macros = "0.16.0"
 ic-cdk-timers = "0.9.0"
 ic-ledger-types = "0.12.0"
-ic-stable-structures = "0.6.4"
+ic-stable-structures = "0.6.6"
 icrc-ledger-types = "0.1.6"
 ic-utils = "0.38"
 itertools = "0.13.0"

core/upgrader/api/spec.did

@@ -296,4 +296,5 @@ service : (InitArg) -> {
   "get_disaster_recovery_state" : () -> (GetDisasterRecoveryStateResult) query;
   "request_disaster_recovery" : (RequestDisasterRecoveryInput) -> (RequestDisasterRecoveryResult);
   "get_logs" : (GetLogsInput) -> (GetLogsResult) query;
+  "deprecated_get_logs" : (GetLogsInput) -> (GetLogsResult) query;
 };

core/upgrader/impl/src/controllers/logs.rs

@@ -28,6 +28,13 @@ fn get_logs(input: upgrader_api::GetLogsInput) -> ApiResult<upgrader_api::GetLog
     CONTROLLER.get_logs(input)
 }
 
+#[query]
+fn deprecated_get_logs(
+    input: upgrader_api::GetLogsInput,
+) -> ApiResult<upgrader_api::GetLogsResponse> {
+    CONTROLLER.deprecated_get_logs(input)
+}
+
 pub struct LogsController {
     disaster_recover_service: Arc<DisasterRecoveryService>,
     logger_service: Arc<LoggerService>,
@@ -59,4 +66,31 @@ impl LogsController {
             Err(UpgraderApiError::Unauthorized.into())
         }
     }
+
+    // Supports fetching the logs from the deprecated log storage.
+    pub fn deprecated_get_logs(
+        &self,
+        input: upgrader_api::GetLogsInput,
+    ) -> ApiResult<upgrader_api::GetLogsResponse> {
+        let caller = caller();
+
+        if is_controller(&caller) || self.disaster_recover_service.is_committee_member(&caller) {
+            let GetLogsResult {
+                logs,
+                next_offset,
+                total,
+            } = self.logger_service.deprecated_get_logs(
+                input.pagination.as_ref().and_then(|p| p.offset),
+                input.pagination.as_ref().and_then(|p| p.limit),
+            );
+
+            Ok(upgrader_api::GetLogsResponse {
+                logs: logs.into_iter().map(|l| l.into()).collect(),
+                total,
+                next_offset,
+            })
+        } else {
+            Err(UpgraderApiError::Unauthorized.into())
+        }
+    }
 }

core/upgrader/impl/src/lib.rs

@@ -33,8 +33,9 @@ type LocalRef<T> = &'static LocalKey<RefCell<T>>;
 
 const MEMORY_ID_TARGET_CANISTER_ID: u8 = 0;
 const MEMORY_ID_DISASTER_RECOVERY: u8 = 1;
-const MEMORY_ID_LOG_INDEX: u8 = 2;
-const MEMORY_ID_LOG_DATA: u8 = 3;
+const DEPRECATED_MEMORY_ID_LOG_INDEX: u8 = 2;
+const DEPRECATED_MEMORY_ID_LOG_DATA: u8 = 3;
+const MEMORY_ID_LOGS: u8 = 4;
 
 thread_local! {
     static MEMORY_MANAGER: RefCell<MemoryManager<DefaultMemoryImpl>> =

core/upgrader/impl/src/model/logging.rs

@@ -1,4 +1,4 @@
-use crate::upgrader_ic_cdk::api::time;
+use crate::upgrader_ic_cdk::next_time;
 use orbit_essentials::{storable, types::Timestamp, utils::timestamp_to_rfc3339};
 use serde::Serialize;
 
@@ -171,7 +171,7 @@ impl LogEntryType {
 impl LogEntry {
     pub fn try_from_entry_type(entry_type: LogEntryType) -> Result<Self, String> {
         Ok(LogEntry {
-            time: time(),
+            time: next_time(),
             entry_type: entry_type.to_type_string(),
             message: entry_type.to_message(),
             data_json: entry_type.to_json_string()?,

core/upgrader/impl/src/services/logger.rs

@@ -1,25 +1,32 @@
 use std::{cell::RefCell, sync::Arc};
 
-use ic_stable_structures::{memory_manager::MemoryId, Log};
+use ic_stable_structures::{memory_manager::MemoryId, BTreeMap, Log};
 use lazy_static::lazy_static;
+use orbit_essentials::types::Timestamp;
 
 use crate::{
     model::{LogEntry, LogEntryType},
-    Memory, MEMORY_ID_LOG_DATA, MEMORY_ID_LOG_INDEX, MEMORY_MANAGER,
+    Memory, DEPRECATED_MEMORY_ID_LOG_DATA, DEPRECATED_MEMORY_ID_LOG_INDEX, MEMORY_ID_LOGS,
+    MEMORY_MANAGER,
 };
 
 pub const MAX_GET_LOGS_LIMIT: u64 = 100;
 pub const DEFAULT_GET_LOGS_LIMIT: u64 = 10;
+pub const MAX_LOG_ENTRIES: u64 = 25000;
 
 thread_local! {
-
-  static STORAGE: RefCell<Log<LogEntry, Memory, Memory>> = RefCell::new(
-      Log::init(
-          MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(MEMORY_ID_LOG_INDEX))),
-          MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(MEMORY_ID_LOG_DATA))),
-      ).expect("Failed to initialize log storage")
-  );
-
+    static DEPRECATED_STORAGE: RefCell<Log<LogEntry, Memory, Memory>> = RefCell::new(
+        Log::init(
+            MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(DEPRECATED_MEMORY_ID_LOG_INDEX))),
+            MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(DEPRECATED_MEMORY_ID_LOG_DATA))),
+        ).expect("Failed to initialize deprecated log storage")
+    );
+
+    static STORAGE: RefCell<BTreeMap<Timestamp, LogEntry, Memory>> = RefCell::new(
+        BTreeMap::init(
+            MEMORY_MANAGER.with(|m| m.borrow().get(MemoryId::new(MEMORY_ID_LOGS))),
+        )
+    );
 }
 
 lazy_static! {
@@ -40,12 +47,12 @@ impl LoggerService {
     /// Tries to log an entry to the storage.
     pub fn try_log(&self, entry_type: LogEntryType) -> Result<(), String> {
         let entry = LogEntry::try_from_entry_type(entry_type)?;
-        STORAGE.with(|storage| {
-            storage
-                .borrow_mut()
-                .append(&entry)
-                .map_err(|err| format!("Failed to log entry: {:?}", err))
-        })?;
+        STORAGE.with_borrow_mut(|storage| {
+            if storage.len() >= MAX_LOG_ENTRIES {
+                let _ = storage.pop_first();
+            }
+            storage.insert(entry.time, entry);
+        });
         Ok(())
     }
 
@@ -71,6 +78,47 @@ impl LoggerService {
                 };
             }
 
+            let offset = offset.unwrap_or(0);
+            let limit = limit
+                .unwrap_or(DEFAULT_GET_LOGS_LIMIT)
+                .min(MAX_GET_LOGS_LIMIT);
+
+            let logs = borrowed
+                .iter()
+                .rev()
+                .skip(offset as usize)
+                .take(limit as usize)
+                .map(|(_, v)| v)
+                .collect::<Vec<_>>();
+
+            let next_offset = if total > offset + limit {
+                Some(offset + limit)
+            } else {
+                None
+            };
+            GetLogsResult {
+                logs,
+                total,
+                next_offset,
+            }
+        })
+    }
+
+    /// Returns logs from the deprecated storage starting from the end of the log.
+    pub fn deprecated_get_logs(&self, offset: Option<u64>, limit: Option<u64>) -> GetLogsResult {
+        DEPRECATED_STORAGE.with(|storage| {
+            let borrowed = storage.borrow();
+
+            let total = borrowed.len();
+
+            if total == 0 {
+                return GetLogsResult {
+                    logs: vec![],
+                    total,
+                    next_offset: None,
+                };
+            }
+
             let offset = offset.unwrap_or(0);
             let limit = limit
                 .unwrap_or(DEFAULT_GET_LOGS_LIMIT)
@@ -153,4 +201,43 @@ mod tests {
         assert_eq!(result.next_offset, None);
         assert_eq!(result.logs[0].entry_type, "set_committee".to_owned());
     }
+
+    #[test]
+    fn test_log_trimming() {
+        for _ in 0..MAX_LOG_ENTRIES {
+            LOGGER_SERVICE.log(LogEntryType::SetCommittee(SetCommitteeLog {
+                committee: mock_committee(),
+            }));
+        }
+
+        let result = LOGGER_SERVICE.get_logs(None, None);
+        assert_eq!(result.total, MAX_LOG_ENTRIES);
+
+        let latest_log_time = result.logs.last().unwrap().time;
+
+        LOGGER_SERVICE.log(LogEntryType::SetCommittee(SetCommitteeLog {
+            committee: mock_committee(),
+        }));
+
+        let result = LOGGER_SERVICE.get_logs(None, None);
+
+        assert_eq!(result.total, MAX_LOG_ENTRIES);
+        assert_ne!(result.logs.last().unwrap().time, latest_log_time);
+    }
+
+    #[test]
+    fn test_deprecated_storage() {
+        let logger_service = LoggerService::default();
+        logger_service.log(LogEntryType::SetCommittee(SetCommitteeLog {
+            committee: mock_committee(),
+        }));
+
+        // new logs should be in the new storage
+        let result = logger_service.get_logs(None, None);
+        assert_eq!(result.total, 1);
+
+        // deprecated logs should not get new logs
+        let result = logger_service.deprecated_get_logs(None, None);
+        assert_eq!(result.total, 0);
+    }
 }