kumo_jsonl/
registry.rs

1//! A registry of where live [`LogWriter`](crate::LogWriter)s are currently
2//! producing segments.
3
4use camino::Utf8PathBuf;
5use std::collections::HashMap;
6use std::sync::{LazyLock, Mutex};
7
8static WRITER_LOCATIONS: LazyLock<Mutex<HashMap<WriterLocation, usize>>> =
9    LazyLock::new(Default::default);
10
11/// The location of the segments produced by a writer.
12#[derive(Hash, PartialEq, Eq, Debug, Clone)]
13pub struct WriterLocation {
14    /// The segment output directory.
15    pub directory: Utf8PathBuf,
16    /// Glob that identifies segments produced by this writer among unrelated
17    /// files and segments from other writers in the same directory.
18    pub pattern: String,
19}
20
21/// An entry in the registry for a writer, held for the lifetime of the writer.
22/// While held, the location of the writer is active. Dropping it releases the
23/// entry, and the location is removed once the last writer to it is dropped.
24pub struct WriterRegistration {
25    location: WriterLocation,
26}
27
28impl WriterRegistration {
29    /// Record that a writer is producing segments (whose names end
30    /// with `suffix`) in `directory`.
31    pub fn new(directory: Utf8PathBuf, suffix: Option<&str>) -> Self {
32        let location = WriterLocation {
33            directory,
34            pattern: pattern_for_suffix(suffix),
35        };
36        *WRITER_LOCATIONS
37            .lock()
38            .expect("jsonl writer registry mutex poisoned")
39            .entry(location.clone())
40            .or_insert(0) += 1;
41        Self { location }
42    }
43}
44
45impl Drop for WriterRegistration {
46    fn drop(&mut self) {
47        let mut locations = WRITER_LOCATIONS
48            .lock()
49            .expect("jsonl writer registry mutex poisoned");
50        if let Some(count) = locations.get_mut(&self.location) {
51            *count -= 1;
52            if *count == 0 {
53                locations.remove(&self.location);
54            }
55        }
56    }
57}
58
59/// Returns the distinct locations where live writers are currently producing
60/// segments.
61pub fn active_writer_locations() -> Vec<WriterLocation> {
62    WRITER_LOCATIONS
63        .lock()
64        .expect("jsonl writer registry mutex poisoned")
65        .keys()
66        .cloned()
67        .collect()
68}
69
70/// Build the glob matching the segments written with the given `suffix`.
71/// Returns the bare `*` when a suffix isn't provided.
72fn pattern_for_suffix(suffix: Option<&str>) -> String {
73    const GLOB_META: &[char] = &['\\', '*', '?', '[', ']', '{', '}'];
74    match suffix {
75        None => "*".to_string(),
76        Some(suffix) => {
77            let mut pattern = String::with_capacity(1 + suffix.len());
78            pattern.push('*');
79            for c in suffix.chars() {
80                if GLOB_META.contains(&c) {
81                    pattern.push('\\');
82                }
83                pattern.push(c);
84            }
85            pattern
86        }
87    }
88}
89
90#[cfg(test)]
91mod test {
92    use super::*;
93
94    #[test]
95    fn suffix_to_pattern() {
96        k9::assert_equal!(pattern_for_suffix(None), "*");
97        k9::assert_equal!(pattern_for_suffix(Some(".log")), "*.log");
98        // Glob metacharacters in the suffix become literals to match the real
99        // segment names rather than acting as wildcards.
100        k9::assert_equal!(pattern_for_suffix(Some("-v1[a]*")), "*-v1\\[a\\]\\*");
101    }
102}