Skip to main content

agent_sessions/reader/
stream.rs

1use super::{ReadStatus, SessionReader, TailMode, input::read_record};
2use crate::parser::tally;
3use crate::statistics::{DecodeError, decode};
4use crate::{Event, LineErrorKind, Located, Location, StreamError};
5use std::io::BufRead;
6
7impl<R: BufRead> Iterator for SessionReader<R> {
8    type Item = Result<Located<Event>, StreamError>;
9    fn next(&mut self) -> Option<Self::Item> {
10        if let Some(event) = self.deliver() {
11            return Some(Ok(event));
12        }
13        if self.ended {
14            return None;
15        }
16        loop {
17            let start = self.summary.bytes_read;
18            let record = match read_record(
19                &mut self.source,
20                &self.opts,
21                &mut self.summary,
22                std::mem::take(&mut self.scratch),
23            ) {
24                Ok(Some(record)) => record,
25                Ok(None) => {
26                    self.ended = true;
27                    self.summary.status = if self.summary.line_errors == 0 {
28                        ReadStatus::Complete
29                    } else {
30                        ReadStatus::CompleteWithErrors
31                    };
32                    return None;
33                }
34                Err(error) => {
35                    self.ended = true;
36                    self.summary.status = ReadStatus::Failed;
37                    return Some(Err(error));
38                }
39            };
40            self.summary.lines += 1;
41            let error = if record.too_long {
42                Some(LineErrorKind::TooLong)
43            } else if std::str::from_utf8(&record.bytes).is_err() {
44                Some(LineErrorKind::InvalidUtf8)
45            } else {
46                None
47            };
48            if let Some(kind) = error {
49                self.scratch = record.bytes;
50                return Some(Err(self.line_error(start, kind)));
51            }
52            if record.bytes.iter().all(u8::is_ascii_whitespace) {
53                self.scratch = record.bytes;
54                self.checkpoint_record();
55                continue;
56            }
57            let decoded = decode(self.agent, &record.bytes, &mut self.state, &self.opts);
58            self.scratch = record.bytes;
59            let mut parsed = match decoded {
60                Ok(parsed) => parsed,
61                Err(DecodeError::Json(e))
62                    if !record.newline
63                        && e.is_eof()
64                        && self.opts.tail == TailMode::AllowIncomplete =>
65                {
66                    self.ended = true;
67                    self.summary.truncated_tail = true;
68                    self.summary.status = ReadStatus::IncompleteTail;
69                    return None;
70                }
71                Err(DecodeError::Json(_)) => {
72                    return Some(Err(self.line_error(start, LineErrorKind::InvalidJson)));
73                }
74                Err(DecodeError::Fields(kind)) => return Some(Err(self.line_error(start, kind))),
75            };
76            for unknown in &parsed.unknown {
77                tally(&mut self.summary.unknown_types, unknown);
78            }
79            for ignored in &parsed.ignored {
80                tally(&mut self.summary.ignored_types, ignored);
81            }
82            let event_count = parsed.events.len();
83            for (ordinal, (event_index, mut event)) in parsed.events.into_iter().enumerate() {
84                if let Event::Message(m) = &mut event {
85                    m.is_sidechain |= self.state.sidechain || self.file_sidechain;
86                }
87                if !self.opts.include.includes(&event) {
88                    continue;
89                }
90                self.pending.push_back(Located {
91                    location: Location {
92                        record_index: self.summary.lines - 1,
93                        line_no: self.summary.lines,
94                        byte_start: start,
95                        byte_end: self.summary.bytes_read,
96                        event_index,
97                    },
98                    at: parsed.at,
99                    timestamp_text: if ordinal + 1 == event_count {
100                        parsed.timestamp_text.take()
101                    } else {
102                        parsed.timestamp_text.clone()
103                    },
104                    record_id: if ordinal + 1 == event_count {
105                        parsed.record_id.take()
106                    } else {
107                        parsed.record_id.clone()
108                    },
109                    record_type: if ordinal + 1 == event_count {
110                        parsed.record_type.take()
111                    } else {
112                        parsed.record_type.clone()
113                    },
114                    session_id: self.state.session_id.clone(),
115                    message_id: if ordinal + 1 == event_count {
116                        parsed.message_id.take()
117                    } else {
118                        parsed.message_id.clone()
119                    },
120                    value: event,
121                });
122            }
123            self.checkpoint_record();
124            if let Some(event) = self.deliver() {
125                return Some(Ok(event));
126            }
127        }
128    }
129}
130
131impl<R: BufRead> SessionReader<R> {
132    fn checkpoint_record(&mut self) {
133        if self.summary.line_errors == 0 {
134            self.checkpoint = Some(self.summary.bytes_read);
135            if self.pending.is_empty() {
136                self.summary.last_complete_byte = self.summary.bytes_read;
137                self.checkpoint = None;
138            }
139        }
140    }
141    fn line_error(&mut self, byte_start: u64, kind: LineErrorKind) -> StreamError {
142        self.summary.line_errors += 1;
143        if !matches!(
144            kind,
145            LineErrorKind::CounterRegression | LineErrorKind::LostUsageBaseline
146        ) {
147            self.state.broken_usage = true;
148        }
149        StreamError::Line {
150            line_no: self.summary.lines,
151            byte_start,
152            kind,
153        }
154    }
155}