Skip to main content

byte_budgeted_export/
byte-budgeted-export.rs

1//! Caller-owned UTF-8 export with a body budget and lossless cursor continuation.
2
3use ic_logging::{EntryInput, Filter, Level, Limits, Log, PageRequest};
4use std::{num::NonZeroUsize, ops::ControlFlow};
5
6#[derive(Debug, Eq, PartialEq)]
7struct ExportPage {
8    body: String,
9    emitted: usize,
10    scanned: usize,
11    next_before: Option<u64>,
12}
13
14#[derive(Debug, Eq, PartialEq)]
15enum ExportError {
16    EntryTooLarge {
17        sequence: u64,
18        required_bytes: usize,
19        max_bytes: usize,
20    },
21}
22
23// This budget bounds returned body bytes, not temporary formatting allocations
24// or an enclosing transport envelope. The host owns the format and those bounds.
25fn export(
26    log: &Log,
27    filter: Filter<'_>,
28    request: PageRequest,
29    max_bytes: usize,
30) -> Result<ExportPage, ExportError> {
31    let mut body = String::new();
32    let mut emitted = 0;
33    let mut last_emitted = None;
34    let mut stopped_for_budget = false;
35    let mut error = None;
36    let scan = log.visit_page(filter, request, |entry| {
37        let encoded = format!("{}: {}\n", entry.sequence, entry.message);
38        if encoded.len() > max_bytes.saturating_sub(body.len()) {
39            stopped_for_budget = true;
40            if emitted == 0 {
41                error = Some(ExportError::EntryTooLarge {
42                    sequence: entry.sequence,
43                    required_bytes: encoded.len(),
44                    max_bytes,
45                });
46            }
47            return ControlFlow::Break(());
48        }
49        body.push_str(&encoded);
50        emitted += 1;
51        last_emitted = Some(entry.sequence);
52        ControlFlow::Continue(())
53    });
54    if let Some(error) = error {
55        return Err(error);
56    }
57    Ok(ExportPage {
58        body,
59        emitted,
60        scanned: scan.scanned,
61        // Break consumes its candidate in the engine. The candidate that did
62        // not fit has not been emitted, so resume from the last delivered row.
63        next_before: if stopped_for_budget {
64            last_emitted
65        } else {
66            scan.next_before
67        },
68    })
69}
70
71const fn request(before: Option<u64>) -> PageRequest {
72    PageRequest {
73        before,
74        max_results: NonZeroUsize::new(3).unwrap(),
75        max_scanned: NonZeroUsize::new(3).unwrap(),
76    }
77}
78
79fn history() -> Log {
80    let mut log = Log::new(Limits {
81        max_entries: NonZeroUsize::new(3).unwrap(),
82        max_message_bytes: 32,
83        max_label_bytes: 16,
84        max_text_bytes: 128,
85    });
86    for (time, message) in ["a", "é", "🙂"].into_iter().enumerate() {
87        log.append(EntryInput {
88            created_at_secs: u64::try_from(time).unwrap(),
89            source: "worker",
90            topic: None,
91            level: Level::Info,
92            message,
93        })
94        .unwrap();
95    }
96    log
97}
98
99fn main() {
100    let log = history();
101    let first = export(&log, Filter::default(), request(None), 8).unwrap();
102    assert_eq!(first.body, "2: 🙂\n");
103    let second = export(&log, Filter::default(), request(first.next_before), 8).unwrap();
104    assert_eq!(second.body, "1: é\n");
105    let third = export(&log, Filter::default(), request(second.next_before), 8).unwrap();
106    assert_eq!(third.body, "0: a\n");
107    assert_eq!(third.next_before, None);
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113
114    #[test]
115    fn budget_stop_does_not_skip_the_unemitted_utf8_candidate() {
116        let log = history();
117        let first = export(&log, Filter::default(), request(None), 8).unwrap();
118        assert_eq!(first.body, "2: 🙂\n");
119        assert_eq!(first.emitted, 1);
120        assert_eq!(first.scanned, 2);
121        assert_eq!(first.next_before, Some(2));
122        let second = export(&log, Filter::default(), request(first.next_before), 8).unwrap();
123        assert_eq!(second.body, "1: é\n");
124        assert_eq!(second.next_before, Some(1));
125        let last = export(&log, Filter::default(), request(second.next_before), 8).unwrap();
126        assert_eq!(last.body, "0: a\n");
127        assert_eq!(last.next_before, None);
128    }
129
130    #[test]
131    fn oversized_first_row_returns_a_typed_error_instead_of_a_retry_loop() {
132        let log = history();
133        for max_bytes in [0, 7] {
134            assert_eq!(
135                export(&log, Filter::default(), request(None), max_bytes),
136                Err(ExportError::EntryTooLarge {
137                    sequence: 2,
138                    required_bytes: 8,
139                    max_bytes,
140                })
141            );
142        }
143        assert_eq!(log.checkpoint().next_sequence(), Some(3));
144    }
145
146    #[test]
147    fn empty_filtered_scans_use_the_engine_continuation() {
148        let log = history();
149        let filter = Filter {
150            min_level: Some(Level::Error),
151            ..Filter::default()
152        };
153        let mut before = None;
154        for expected in [Some(2), Some(1), None] {
155            let mut request = request(before);
156            request.max_scanned = NonZeroUsize::new(1).unwrap();
157            let page = export(&log, filter, request, 0).unwrap();
158            assert_eq!(page.body, "");
159            assert_eq!(page.emitted, 0);
160            assert_eq!(page.scanned, 1);
161            assert_eq!(page.next_before, expected);
162            before = page.next_before;
163        }
164    }
165
166    #[test]
167    fn count_and_scan_limits_still_apply_with_a_large_body_budget() {
168        let log = history();
169        for (results, scanned) in [(1, 3), (3, 1)] {
170            let mut request = request(None);
171            request.max_results = NonZeroUsize::new(results).unwrap();
172            request.max_scanned = NonZeroUsize::new(scanned).unwrap();
173            let page = export(&log, Filter::default(), request, 128).unwrap();
174            assert_eq!(page.body, "2: 🙂\n");
175            assert_eq!(page.emitted, 1);
176            assert_eq!(page.scanned, 1);
177            assert_eq!(page.next_before, Some(2));
178        }
179    }
180}