Skip to content

Commit

Permalink
Better line handling when input is chunked
Browse files Browse the repository at this point in the history
Fixes #6
  • Loading branch information
snoyberg committed Jun 19, 2024
1 parent b7c44fc commit cb87d36
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 5 deletions.
19 changes: 14 additions & 5 deletions src/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ use anyhow::{Context, Result};

use clap::{arg, Parser};

use crate::slack::SlackApp;
use crate::{line_helper::LineHelper, slack::SlackApp};

#[derive(Parser)]
pub(crate) struct Cli {
Expand Down Expand Up @@ -217,6 +217,8 @@ fn process_std_handle(
max_recent_output: usize,
) {
let mut buffer = [0u8; 4096];
let mut line_helper = LineHelper::new();

loop {
match reader
.read(&mut buffer)
Expand All @@ -243,13 +245,12 @@ fn process_std_handle(
break;
}

let mut guard = recent_output.lock();
for line in buffer.split(|x| *x == b'\n') {
for line in line_helper.append(&buffer[..size]) {
let mut guard = recent_output.lock();
if guard.len() >= max_recent_output {
guard.pop_front();
}
let line = line.strip_suffix(&[b'\r']).unwrap_or(line);
guard.push_back(String::from_utf8_lossy(line).into_owned());
guard.push_back(line);
}
}
Err(e) => {
Expand All @@ -258,6 +259,14 @@ fn process_std_handle(
}
}
}

if let Some(line) = line_helper.finish() {
let mut guard = recent_output.lock();
if guard.len() >= max_recent_output {
guard.pop_front();
}
guard.push_back(line);
}
}

fn detect_deadlock(
Expand Down
101 changes: 101 additions & 0 deletions src/line_helper.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
const BUFFER_SIZE: usize = 8192;

pub(crate) struct LineHelper {
buffer: [u8; BUFFER_SIZE],
len: usize,
}

impl LineHelper {
pub(crate) fn new() -> Self {
LineHelper {
buffer: [0; BUFFER_SIZE],
len: 0,
}
}
pub(crate) fn append(&mut self, new_data: &[u8]) -> impl Iterator<Item = String> {
let mut res = vec![];

if new_data.len() + self.len > BUFFER_SIZE {
res.push(String::from_utf8_lossy(&self.buffer[..self.len]).into_owned());
self.len = 0;
}

self.buffer[self.len..self.len + new_data.len()].copy_from_slice(new_data);
self.len += new_data.len();

while let Some(idx) = find_newline(&self.buffer[..self.len]) {
let end = if idx > 0 && self.buffer[idx - 1] == b'\r' {
idx - 1
} else {
idx
};
let line = String::from_utf8_lossy(&self.buffer[..end]).into_owned();
res.push(line);
self.buffer.rotate_left(idx + 1);
self.len -= idx + 1;
}

res.into_iter()
}

pub(crate) fn finish(self) -> Option<String> {
if self.len == 0 {
None
} else {
Some(String::from_utf8_lossy(&self.buffer[..self.len]).into_owned())
}
}
}

fn find_newline(s: &[u8]) -> Option<usize> {
s.iter()
.enumerate()
.find_map(|(idx, c)| if *c == b'\n' { Some(idx) } else { None })
}

#[cfg(test)]
mod tests {
use super::*;

#[test]
fn one_line() {
let mut helper = LineHelper::new();
assert_eq!(
helper.append(b"helloworld").collect::<Vec<_>>(),
Vec::<String>::new()
);
assert_eq!(helper.finish(), Some("helloworld".to_owned()))
}

#[test]
fn basic() {
let mut helper = LineHelper::new();
assert_eq!(
helper.append(b"hello\nworld\r\n").collect::<Vec<_>>(),
vec!["hello".to_owned(), "world".to_owned()]
);
assert_eq!(helper.finish(), None)
}

#[test]
fn chunked() {
let mut helper = LineHelper::new();
assert_eq!(
helper.append(b"he").collect::<Vec<_>>(),
Vec::<String>::new()
);
assert_eq!(
helper.append(b"ll").collect::<Vec<_>>(),
Vec::<String>::new()
);
assert_eq!(
helper.append(b"o\r\n").collect::<Vec<_>>(),
vec!["hello".to_owned()]
);
assert_eq!(
helper.append(b"world").collect::<Vec<_>>(),
Vec::<String>::new()
);
assert_eq!(helper.finish(), Some("world".to_owned()))
}
}
1 change: 1 addition & 0 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ use clap::Parser;
use pid1::Pid1Settings;

mod cli;
mod line_helper;
mod slack;

fn main() -> Result<()> {
Expand Down

0 comments on commit cb87d36

Please sign in to comment.