aft/bash_background/
buffer.rs

1use std::fs::{self, File};
2use std::io::{self, Read, Seek, SeekFrom};
3use std::path::{Path, PathBuf};
4
5pub const DISK_LIMIT_BYTES: u64 = 100 * 1024 * 1024;
6
7#[derive(Debug, Clone, Copy)]
8pub enum StreamKind {
9    Stdout,
10    Stderr,
11}
12
13#[derive(Debug, Clone)]
14pub struct BgBuffer {
15    stdout_path: PathBuf,
16    stderr_path: PathBuf,
17    rotated: bool,
18}
19
20impl BgBuffer {
21    pub fn new(stdout_path: PathBuf, stderr_path: PathBuf) -> Self {
22        Self {
23            stdout_path,
24            stderr_path,
25            rotated: false,
26        }
27    }
28
29    pub fn stdout_path(&self) -> &Path {
30        &self.stdout_path
31    }
32
33    pub fn stderr_path(&self) -> &Path {
34        &self.stderr_path
35    }
36
37    pub fn read_tail(&self, max_bytes: usize) -> (String, bool) {
38        let stdout = read_file_tail(&self.stdout_path, max_bytes);
39        let stderr = read_file_tail(&self.stderr_path, max_bytes);
40        match (stdout, stderr) {
41            (Ok((stdout, stdout_truncated)), Ok((stderr, stderr_truncated))) => {
42                let mut output = Vec::with_capacity(stdout.len().saturating_add(stderr.len()));
43                output.extend_from_slice(&stdout);
44                output.extend_from_slice(&stderr);
45                if output.len() > max_bytes {
46                    let keep_from = output.len().saturating_sub(max_bytes);
47                    output.drain(..keep_from);
48                }
49                (
50                    String::from_utf8_lossy(&output).into_owned(),
51                    self.rotated
52                        || stdout_truncated
53                        || stderr_truncated
54                        || output.len() >= max_bytes && (stdout.len() + stderr.len()) > max_bytes,
55                )
56            }
57            (Ok((stdout, stdout_truncated)), Err(_)) => (
58                String::from_utf8_lossy(&stdout).into_owned(),
59                self.rotated || stdout_truncated,
60            ),
61            (Err(_), Ok((stderr, stderr_truncated))) => (
62                String::from_utf8_lossy(&stderr).into_owned(),
63                self.rotated || stderr_truncated,
64            ),
65            (Err(_), Err(_)) => (String::new(), self.rotated),
66        }
67    }
68
69    pub fn read_stream_tail(&self, stream: StreamKind, max_bytes: usize) -> (String, bool) {
70        let path = match stream {
71            StreamKind::Stdout => &self.stdout_path,
72            StreamKind::Stderr => &self.stderr_path,
73        };
74        match read_file_tail(path, max_bytes) {
75            Ok((bytes, truncated)) => (
76                String::from_utf8_lossy(&bytes).into_owned(),
77                self.rotated || truncated,
78            ),
79            Err(_) => (String::new(), self.rotated),
80        }
81    }
82
83    pub fn output_path(&self) -> Option<PathBuf> {
84        Some(self.stdout_path.clone())
85    }
86
87    pub fn enforce_terminal_cap(&mut self) {
88        if truncate_front(&self.stdout_path, DISK_LIMIT_BYTES).unwrap_or(false) {
89            self.rotated = true;
90        }
91        if truncate_front(&self.stderr_path, DISK_LIMIT_BYTES).unwrap_or(false) {
92            self.rotated = true;
93        }
94    }
95
96    pub fn cleanup(&self) {
97        let _ = fs::remove_file(&self.stdout_path);
98        let _ = fs::remove_file(&self.stderr_path);
99    }
100}
101
102fn read_file_tail(path: &Path, max_bytes: usize) -> io::Result<(Vec<u8>, bool)> {
103    if max_bytes == 0 {
104        return Ok((
105            Vec::new(),
106            path.metadata()
107                .map(|metadata| metadata.len() > 0)
108                .unwrap_or(false),
109        ));
110    }
111
112    let mut file = File::open(path)?;
113    let len = file.metadata()?.len();
114    let read_len = len.min(max_bytes as u64);
115    if read_len > 0 {
116        file.seek(SeekFrom::End(-(read_len as i64)))?;
117    }
118    let mut bytes = Vec::with_capacity(read_len as usize);
119    file.read_to_end(&mut bytes)?;
120    Ok((bytes, len > max_bytes as u64))
121}
122
123fn truncate_front(path: &Path, retain_bytes: u64) -> io::Result<bool> {
124    let len = match path.metadata() {
125        Ok(metadata) => metadata.len(),
126        Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(false),
127        Err(error) => return Err(error),
128    };
129    if len <= retain_bytes {
130        return Ok(false);
131    }
132
133    let mut file = File::open(path)?;
134    file.seek(SeekFrom::End(-(retain_bytes as i64)))?;
135    let mut tail = Vec::with_capacity(retain_bytes as usize);
136    file.read_to_end(&mut tail)?;
137    let tmp = path.with_extension(format!(
138        "{}.tmp",
139        path.extension()
140            .and_then(|extension| extension.to_str())
141            .unwrap_or("out")
142    ));
143    fs::write(&tmp, tail)?;
144    fs::rename(&tmp, path)?;
145    Ok(true)
146}
aft/bash_background/buffer.rs

aft/bash_background/
buffer.rs
