zlog: Implement better scope map (#27408)

Ben Kunkle created

Another step towards having `zlog` as the default logging solution in
Zed.
The new ScopeMap replaces the previous HashMap based implementation used
for scope lookups to:

A. Reduce complexity
B. Increase speed at which non-enabled logs can be filtered out
C. Provide more granular control over how scopes are determined to be
enabled/disabled,
   and what caching/other speed increase opportunities are available

Release Notes:

- N/A

Change summary

crates/zlog/src/zlog.rs | 410 ++++++++++++++++++++++++++++++++++++------
1 file changed, 350 insertions(+), 60 deletions(-)

Detailed changes

crates/zlog/src/zlog.rs 🔗

@@ -272,17 +272,17 @@ impl Timer {
 
 pub mod scope_map {
     use std::{
-        collections::HashMap,
+        collections::{HashMap, VecDeque},
         hash::{DefaultHasher, Hasher},
         sync::{
             atomic::{AtomicU64, Ordering},
             RwLock,
         },
+        usize,
     };
 
     use super::*;
 
-    type ScopeMap = HashMap<ScopeAlloc, log_impl::Level>;
     static SCOPE_MAP: RwLock<Option<ScopeMap>> = RwLock::new(None);
     static SCOPE_MAP_HASH: AtomicU64 = AtomicU64::new(0);
 
@@ -309,33 +309,24 @@ pub mod scope_map {
             // if no scopes are enabled, default to enabled detection done by `log` crate
             return (true, level);
         }
-        let mut scope_alloc = private::scope_to_alloc(scope);
-        let mut level_enabled = map.get(&scope_alloc);
-        if level_enabled.is_none() {
-            for i in (0..SCOPE_DEPTH_MAX).rev() {
-                if scope_alloc[i] == "" {
-                    continue;
-                }
-                scope_alloc[i].clear();
-                if let Some(level) = map.get(&scope_alloc) {
-                    level_enabled = Some(level);
-                    break;
-                }
+        let enabled_status = map.is_enabled(&scope, level);
+        match enabled_status {
+            EnabledStatus::NotConfigured => {
+                // if this scope isn't configured, default to enabled detection done by `log` crate
+                return (true, level);
+            }
+            EnabledStatus::Enabled => {
+                // if this scope is enabled, enable logging
+                // note: bumping level to min level that will be printed
+                // to work around log crate limitations
+                return (true, level_min);
+            }
+            EnabledStatus::Disabled => {
+                // if the configured level is lower than the requested level, disable logging
+                // note: err = 0, warn = 1, etc.
+                return (false, level);
             }
         }
-        let Some(level_enabled) = level_enabled else {
-            // if this scope isn't configured, default to enabled detection done by `log` crate
-            return (true, level);
-        };
-        if level_enabled < &level {
-            // if the configured level is lower than the requested level, disable logging
-            // note: err = 0, warn = 1, etc.
-            return (false, level);
-        }
-
-        // note: bumping level to min level that will be printed
-        // to work around log crate limitations
-        return (true, level_min);
     }
 
     fn hash_scope_map_settings(map: &HashMap<String, String>) -> u64 {
@@ -355,37 +346,7 @@ pub mod scope_map {
         if hash_old == hash_new && hash_old != 0 {
             return;
         }
-        // compute new scope map then atomically swap it, instead of
-        // updating in place to reduce contention
-        let mut map_new = ScopeMap::with_capacity(settings.len());
-        'settings: for (key, value) in settings {
-            let level = match value.to_ascii_lowercase().as_str() {
-                "" => log_impl::Level::Trace,
-                "trace" => log_impl::Level::Trace,
-                "debug" => log_impl::Level::Debug,
-                "info" => log_impl::Level::Info,
-                "warn" => log_impl::Level::Warn,
-                "error" => log_impl::Level::Error,
-                "off" | "disable" | "no" | "none" | "disabled" => {
-                    crate::warn!("Invalid log level \"{value}\", set to error to disable non-error logging. Defaulting to error");
-                    log_impl::Level::Error
-                }
-                _ => {
-                    crate::warn!("Invalid log level \"{value}\", ignoring");
-                    continue 'settings;
-                }
-            };
-            let mut scope_buf = [""; SCOPE_DEPTH_MAX];
-            for (index, scope) in key.split(SCOPE_STRING_SEP).enumerate() {
-                let Some(scope_ptr) = scope_buf.get_mut(index) else {
-                    crate::warn!("Invalid scope key, too many nested scopes: '{key}'");
-                    continue 'settings;
-                };
-                *scope_ptr = scope;
-            }
-            let scope = scope_buf.map(|s| s.to_string());
-            map_new.insert(scope, level);
-        }
+        let map_new = ScopeMap::new_from_settings(settings);
 
         if let Ok(_) = SCOPE_MAP_HASH.compare_exchange(
             hash_old,
@@ -397,8 +358,337 @@ pub mod scope_map {
                 SCOPE_MAP.clear_poison();
                 err.into_inner()
             });
-            *map = Some(map_new.clone());
-            // note: hash update done here to ensure consistency with scope map
+            *map = Some(map_new);
+        }
+    }
+
+    fn level_from_level_str(level_str: &String) -> Option<log_impl::Level> {
+        let level = match level_str.to_ascii_lowercase().as_str() {
+            "" => log_impl::Level::Trace,
+            "trace" => log_impl::Level::Trace,
+            "debug" => log_impl::Level::Debug,
+            "info" => log_impl::Level::Info,
+            "warn" => log_impl::Level::Warn,
+            "error" => log_impl::Level::Error,
+            "off" | "disable" | "no" | "none" | "disabled" => {
+                crate::warn!("Invalid log level \"{level_str}\", set to error to disable non-error logging. Defaulting to error");
+                log_impl::Level::Error
+            }
+            _ => {
+                crate::warn!("Invalid log level \"{level_str}\", ignoring");
+                return None;
+            }
+        };
+        return Some(level);
+    }
+
+    fn scope_alloc_from_scope_str(scope_str: &String) -> Option<ScopeAlloc> {
+        let mut scope_buf = [""; SCOPE_DEPTH_MAX];
+        let mut index = 0;
+        let mut scope_iter = scope_str.split(SCOPE_STRING_SEP);
+        while index < SCOPE_DEPTH_MAX {
+            let Some(scope) = scope_iter.next() else {
+                break;
+            };
+            if scope == "" {
+                continue;
+            }
+            scope_buf[index] = scope;
+            index += 1;
+        }
+        if index == 0 {
+            return None;
+        }
+        if let Some(_) = scope_iter.next() {
+            crate::warn!(
+                "Invalid scope key, too many nested scopes: '{scope_str}'. Max depth is {SCOPE_DEPTH_MAX}",
+            );
+            return None;
+        }
+        let scope = scope_buf.map(|s| s.to_string());
+        return Some(scope);
+    }
+
+    pub struct ScopeMap {
+        entries: Vec<ScopeMapEntry>,
+        root_count: usize,
+    }
+
+    pub struct ScopeMapEntry {
+        scope: String,
+        enabled: Option<log_impl::Level>,
+        descendants: std::ops::Range<usize>,
+    }
+
+    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
+    pub enum EnabledStatus {
+        Enabled,
+        Disabled,
+        NotConfigured,
+    }
+
+    impl ScopeMap {
+        pub fn new_from_settings(items_input_map: &HashMap<String, String>) -> Self {
+            let mut items = items_input_map
+                .into_iter()
+                .filter_map(|(scope_str, level_str)| {
+                    let scope = scope_alloc_from_scope_str(&scope_str)?;
+                    let level = level_from_level_str(&level_str)?;
+                    return Some((scope, level));
+                })
+                .collect::<Vec<_>>();
+
+            items.sort_by(|a, b| a.0.cmp(&b.0));
+
+            let mut this = Self {
+                entries: Vec::with_capacity(items.len() * SCOPE_DEPTH_MAX),
+                root_count: 0,
+            };
+
+            let items_count = items.len();
+
+            struct ProcessQueueEntry {
+                parent_index: usize,
+                depth: usize,
+                items_range: std::ops::Range<usize>,
+            }
+            let mut process_queue = VecDeque::new();
+            process_queue.push_back(ProcessQueueEntry {
+                parent_index: usize::MAX,
+                depth: 0,
+                items_range: 0..items_count,
+            });
+
+            let empty_range = 0..0;
+
+            while let Some(process_entry) = process_queue.pop_front() {
+                let ProcessQueueEntry {
+                    items_range,
+                    depth,
+                    parent_index,
+                } = process_entry;
+                let mut cursor = items_range.start;
+                let res_entries_start = this.entries.len();
+                while cursor < items_range.end {
+                    let sub_items_start = cursor;
+                    cursor += 1;
+                    let scope_name = &items[sub_items_start].0[depth];
+                    while cursor < items_range.end && &items[cursor].0[depth] == scope_name {
+                        cursor += 1;
+                    }
+                    let sub_items_end = cursor;
+                    if scope_name == "" {
+                        assert_eq!(sub_items_start + 1, sub_items_end);
+                        assert_ne!(depth, 0);
+                        assert_ne!(parent_index, usize::MAX);
+                        assert!(this.entries[parent_index].enabled.is_none());
+                        this.entries[parent_index].enabled = Some(items[sub_items_start].1);
+                        continue;
+                    }
+                    let is_valid_scope = scope_name != "";
+                    let is_last = depth + 1 == SCOPE_DEPTH_MAX || !is_valid_scope;
+                    let mut enabled = None;
+                    if is_last {
+                        assert_eq!(sub_items_start + 1, sub_items_end);
+                        enabled = Some(items[sub_items_start].1);
+                    } else {
+                        let entry_index = this.entries.len();
+                        process_queue.push_back(ProcessQueueEntry {
+                            items_range: sub_items_start..sub_items_end,
+                            parent_index: entry_index,
+                            depth: depth + 1,
+                        });
+                    }
+                    this.entries.push(ScopeMapEntry {
+                        scope: scope_name.to_owned(),
+                        enabled,
+                        descendants: empty_range.clone(),
+                    });
+                }
+                let res_entries_end = this.entries.len();
+                if parent_index != usize::MAX {
+                    this.entries[parent_index].descendants = res_entries_start..res_entries_end;
+                } else {
+                    this.root_count = res_entries_end;
+                }
+            }
+
+            return this;
+        }
+
+        pub fn is_empty(&self) -> bool {
+            self.entries.is_empty()
+        }
+
+        pub fn is_enabled<S>(
+            &self,
+            scope: &[S; SCOPE_DEPTH_MAX],
+            level: log_impl::Level,
+        ) -> EnabledStatus
+        where
+            S: AsRef<str>,
+        {
+            let mut enabled = None;
+            let mut cur_range = &self.entries[0..self.root_count];
+            let mut depth = 0;
+
+            'search: while !cur_range.is_empty()
+                && depth < SCOPE_DEPTH_MAX
+                && scope[depth].as_ref() != ""
+            {
+                for entry in cur_range {
+                    if entry.scope == scope[depth].as_ref() {
+                        // note:
+                        enabled = entry.enabled.or(enabled);
+                        cur_range = &self.entries[entry.descendants.clone()];
+                        depth += 1;
+                        continue 'search;
+                    }
+                }
+                break 'search;
+            }
+
+            return enabled.map_or(EnabledStatus::NotConfigured, |level_enabled| {
+                if level <= level_enabled {
+                    EnabledStatus::Enabled
+                } else {
+                    EnabledStatus::Disabled
+                }
+            });
+        }
+    }
+
+    #[cfg(test)]
+    mod tests {
+        use super::*;
+
+        fn scope_map_from_keys(kv: &[(&str, &str)]) -> ScopeMap {
+            let hash_map: HashMap<String, String> = kv
+                .iter()
+                .map(|(k, v)| (k.to_string(), v.to_string()))
+                .collect();
+            ScopeMap::new_from_settings(&hash_map)
+        }
+
+        #[test]
+        fn test_initialization() {
+            let map = scope_map_from_keys(&[("a.b.c.d", "trace")]);
+            assert_eq!(map.root_count, 1);
+            assert_eq!(map.entries.len(), 4);
+
+            let map = scope_map_from_keys(&[]);
+            assert_eq!(map.root_count, 0);
+            assert_eq!(map.entries.len(), 0);
+
+            let map = scope_map_from_keys(&[("", "trace")]);
+            assert_eq!(map.root_count, 0);
+            assert_eq!(map.entries.len(), 0);
+
+            let map = scope_map_from_keys(&[("foo..bar", "trace")]);
+            assert_eq!(map.root_count, 1);
+            assert_eq!(map.entries.len(), 2);
+
+            let map = scope_map_from_keys(&[
+                ("a.b.c.d", "trace"),
+                ("e.f.g.h", "debug"),
+                ("i.j.k.l", "info"),
+                ("m.n.o.p", "warn"),
+                ("q.r.s.t", "error"),
+            ]);
+            assert_eq!(map.root_count, 5);
+            assert_eq!(map.entries.len(), 20);
+            assert_eq!(map.entries[0].scope, "a");
+            assert_eq!(map.entries[1].scope, "e");
+            assert_eq!(map.entries[2].scope, "i");
+            assert_eq!(map.entries[3].scope, "m");
+            assert_eq!(map.entries[4].scope, "q");
+        }
+
+        fn scope_from_scope_str(scope_str: &'static str) -> Scope {
+            let mut scope_buf = [""; SCOPE_DEPTH_MAX];
+            let mut index = 0;
+            let mut scope_iter = scope_str.split(SCOPE_STRING_SEP);
+            while index < SCOPE_DEPTH_MAX {
+                let Some(scope) = scope_iter.next() else {
+                    break;
+                };
+                if scope == "" {
+                    continue;
+                }
+                scope_buf[index] = scope;
+                index += 1;
+            }
+            assert_ne!(index, 0);
+            assert!(scope_iter.next().is_none());
+            return scope_buf;
+        }
+
+        #[test]
+        fn test_is_enabled() {
+            let map = scope_map_from_keys(&[
+                ("a.b.c.d", "trace"),
+                ("e.f.g.h", "debug"),
+                ("i.j.k.l", "info"),
+                ("m.n.o.p", "warn"),
+                ("q.r.s.t", "error"),
+            ]);
+            use log_impl::Level;
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("a.b.c.d"), Level::Trace),
+                EnabledStatus::Enabled
+            );
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("a.b.c.d"), Level::Debug),
+                EnabledStatus::Enabled
+            );
+
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("e.f.g.h"), Level::Debug),
+                EnabledStatus::Enabled
+            );
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("e.f.g.h"), Level::Info),
+                EnabledStatus::Enabled
+            );
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("e.f.g.h"), Level::Trace),
+                EnabledStatus::Disabled
+            );
+
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("i.j.k.l"), Level::Info),
+                EnabledStatus::Enabled
+            );
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("i.j.k.l"), Level::Warn),
+                EnabledStatus::Enabled
+            );
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("i.j.k.l"), Level::Debug),
+                EnabledStatus::Disabled
+            );
+
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("m.n.o.p"), Level::Warn),
+                EnabledStatus::Enabled
+            );
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("m.n.o.p"), Level::Error),
+                EnabledStatus::Enabled
+            );
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("m.n.o.p"), Level::Info),
+                EnabledStatus::Disabled
+            );
+
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("q.r.s.t"), Level::Error),
+                EnabledStatus::Enabled
+            );
+            assert_eq!(
+                map.is_enabled(&scope_from_scope_str("q.r.s.t"), Level::Warn),
+                EnabledStatus::Disabled
+            );
         }
     }
 }