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 = Ok(false);
35    let scan = log.visit_page(filter, request, |entry| {
36        let encoded = format!("{}: {}\n", entry.sequence, entry.message);
37        if encoded.len() > max_bytes.saturating_sub(body.len()) {
38            stopped_for_budget = if emitted == 0 {
39                Err(ExportError::EntryTooLarge {
40                    sequence: entry.sequence,
41                    required_bytes: encoded.len(),
42                    max_bytes,
43                })
44            } else {
45                Ok(true)
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    let stopped_for_budget = stopped_for_budget?;
55    Ok(ExportPage {
56        body,
57        emitted,
58        scanned: scan.scanned,
59        // Break consumes its candidate in the engine. The candidate that did
60        // not fit has not been emitted, so resume from the last delivered row.
61        next_before: if stopped_for_budget {
62            last_emitted
63        } else {
64            scan.next_before
65        },
66    })
67}
68
69const fn request(before: Option<u64>) -> PageRequest {
70    PageRequest {
71        before,
72        max_results: NonZeroUsize::new(3).unwrap(),
73        max_scanned: NonZeroUsize::new(3).unwrap(),
74    }
75}
76
77fn history() -> Log {
78    let mut log = Log::new(Limits {
79        max_entries: NonZeroUsize::new(3).unwrap(),
80        max_message_bytes: 32,
81        max_label_bytes: 16,
82        max_text_bytes: 128,
83    });
84    for (time, message) in ["a", "é", "🙂"].into_iter().enumerate() {
85        log.append(EntryInput {
86            created_at_secs: u64::try_from(time).unwrap(),
87            source: "worker",
88            topic: None,
89            level: Level::Info,
90            message,
91        })
92        .unwrap();
93    }
94    log
95}
96
97fn main() {
98    let log = history();
99    let first = export(&log, Filter::default(), request(None), 8).unwrap();
100    assert_eq!(first.body, "2: 🙂\n");
101    let second = export(&log, Filter::default(), request(first.next_before), 8).unwrap();
102    assert_eq!(second.body, "1: é\n");
103    let third = export(&log, Filter::default(), request(second.next_before), 8).unwrap();
104    assert_eq!(third.body, "0: a\n");
105    assert_eq!(third.next_before, None);
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn budget_stop_does_not_skip_the_unemitted_utf8_candidate() {
114        let log = history();
115        let first = export(&log, Filter::default(), request(None), 8).unwrap();
116        assert_eq!(first.body, "2: 🙂\n");
117        assert_eq!(first.emitted, 1);
118        assert_eq!(first.scanned, 2);
119        assert_eq!(first.next_before, Some(2));
120        let second = export(&log, Filter::default(), request(first.next_before), 8).unwrap();
121        assert_eq!(second.body, "1: é\n");
122        assert_eq!(second.next_before, Some(1));
123        let last = export(&log, Filter::default(), request(second.next_before), 8).unwrap();
124        assert_eq!(last.body, "0: a\n");
125        assert_eq!(last.next_before, None);
126    }
127
128    #[test]
129    fn oversized_first_row_returns_a_typed_error_instead_of_a_retry_loop() {
130        let log = history();
131        for max_bytes in [0, 7] {
132            assert_eq!(
133                export(&log, Filter::default(), request(None), max_bytes),
134                Err(ExportError::EntryTooLarge {
135                    sequence: 2,
136                    required_bytes: 8,
137                    max_bytes,
138                })
139            );
140        }
141        assert_eq!(log.checkpoint().next_sequence(), Some(3));
142    }
143
144    #[test]
145    fn empty_filtered_scans_use_the_engine_continuation() {
146        let log = history();
147        let filter = Filter {
148            min_level: Some(Level::Error),
149            ..Filter::default()
150        };
151        let mut before = None;
152        for expected in [Some(2), Some(1), None] {
153            let mut request = request(before);
154            request.max_scanned = NonZeroUsize::new(1).unwrap();
155            let page = export(&log, filter, request, 0).unwrap();
156            assert_eq!(page.body, "");
157            assert_eq!(page.emitted, 0);
158            assert_eq!(page.scanned, 1);
159            assert_eq!(page.next_before, expected);
160            before = page.next_before;
161        }
162    }
163
164    #[test]
165    fn count_and_scan_limits_still_apply_with_a_large_body_budget() {
166        let log = history();
167        for (results, scanned) in [(1, 3), (3, 1)] {
168            let mut request = request(None);
169            request.max_results = NonZeroUsize::new(results).unwrap();
170            request.max_scanned = NonZeroUsize::new(scanned).unwrap();
171            let page = export(&log, Filter::default(), request, 128).unwrap();
172            assert_eq!(page.body, "2: 🙂\n");
173            assert_eq!(page.emitted, 1);
174            assert_eq!(page.scanned, 1);
175            assert_eq!(page.next_before, Some(2));
176        }
177    }
178}