timberfs 0.6.0

//! `timberfs grep`: entry-aware grep. A log ENTRY is a line carrying a
//! timestamp plus all its continuation lines (stack traces, wrapped
//! output) — the pattern is matched against the whole entry, and matching
//! entries are printed whole. Entry boundaries are detected with the same
//! timestamp auto-detection as import (--timestamp-regex/--format for
//! exotic formats).
//!
//!     cat any.log | timberfs grep 'tenantId=FOO' | timberfs grep -v DEBUG
//!     timberfs grep ERROR backing/app.log --from 13:00 --to 14:00
//!     timberfs grep 'req-8f3a' incident.timber --has req-8f3a
//!
//! Input is stdin (raw log bytes), a plain log file, or a timberfs
//! log/bundle — in the timberfs case --from/--to/--has pre-select chunks
//! first (time index + .grain Bloom filters), then entries are matched
//! exactly. Piping several greps gives entry-level AND, as with grep.

use std::fs::File;
use std::io::{self, BufRead, BufReader, Read, Write};
use std::os::unix::fs::FileExt;
use std::path::Path;

use anyhow::{bail, Context};
use regex::bytes::{Regex, RegexBuilder};

use crate::format::ChunkRecord;
use crate::import::Extractor;
use crate::query::{is_bundle, open_source, parse_time, select_chunks};

/// A timestamp-less flood can't balloon memory: entries are split here.
const ENTRY_CAP: usize = 16 << 20;

/// Streams decompressed content of the selected chunks in order.
struct ChunkStream {
    file: File,
    chunks: Vec<ChunkRecord>,
    idx: usize,
    buf: Vec<u8>,
    pos: usize,
}

impl Read for ChunkStream {
    fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
        while self.pos == self.buf.len() {
            if self.idx == self.chunks.len() {
                return Ok(0);
            }
            let c = self.chunks[self.idx];
            self.idx += 1;
            let mut comp = vec![0u8; c.comp_len as usize];
            self.file.read_exact_at(&mut comp, c.comp_start)?;
            self.buf = zstd::stream::decode_all(&comp[..])?;
            self.pos = 0;
        }
        let n = (self.buf.len() - self.pos).min(out.len());
        out[..n].copy_from_slice(&self.buf[self.pos..self.pos + n]);
        self.pos += n;
        Ok(n)
    }
}

struct Entries<R: BufRead> {
    reader: R,
    extractor: Extractor,
    pending: Option<Vec<u8>>,
    warned_cap: bool,
}

impl<R: BufRead> Entries<R> {
    fn next_entry(&mut self) -> io::Result<Option<Vec<u8>>> {
        let mut entry = match self.pending.take() {
            Some(line) => line,
            None => {
                let mut line = Vec::new();
                if self.reader.read_until(b'\n', &mut line)? == 0 {
                    return Ok(None);
                }
                line
            }
        };
        loop {
            let mut line = Vec::new();
            if self.reader.read_until(b'\n', &mut line)? == 0 {
                break;
            }
            let head = String::from_utf8_lossy(&line[..line.len().min(256)]);
            if self.extractor.extract(&head).is_some() {
                self.pending = Some(line);
                break;
            }
            if entry.len() + line.len() > ENTRY_CAP {
                if !self.warned_cap {
                    eprintln!("timberfs: entry exceeds 16 MiB; splitting");
                    self.warned_cap = true;
                }
                self.pending = Some(line);
                break;
            }
            entry.extend_from_slice(&line);
        }
        Ok(Some(entry))
    }
}

fn run<R: BufRead>(
    mut entries: Entries<R>,
    re: &Regex,
    invert: bool,
    count: bool,
    prefix: Option<&[u8]>,
) -> anyhow::Result<u64> {
    let stdout = io::stdout();
    let mut out = stdout.lock();
    let mut matched: u64 = 0;
    while let Some(entry) = entries.next_entry()? {
        if re.is_match(&entry) != invert {
            matched += 1;
            if !count {
                match prefix {
                    None => out.write_all(&entry)?,
                    Some(label) => {
                        for line in entry.split_inclusive(|&b| b == b'\n') {
                            out.write_all(label)?;
                            out.write_all(b":")?;
                            out.write_all(line)?;
                        }
                    }
                }
            }
        }
    }
    out.flush()?;
    Ok(matched)
}

#[allow(clippy::too_many_arguments)]
fn grep_one(
    p: &Path,
    re: &Regex,
    extractor: Extractor,
    from: Option<&str>,
    to: Option<&str>,
    has: &[String],
    invert: bool,
    count: bool,
    prefix: Option<&[u8]>,
) -> anyhow::Result<u64> {
    let is_timberfs_source = is_bundle(p)
        || matches!(
            p.extension().and_then(|e| e.to_str()),
            Some(crate::format::TRUNK_EXT) | Some(crate::format::RINGS_EXT)
        )
        || !p.is_file();

    if !is_timberfs_source {
        // a plain log file at the exact path
        if from.is_some() || to.is_some() || !has.is_empty() {
            bail!(
                "--from/--to/--has need a timberfs log or bundle; {} is a plain file",
                p.display()
            );
        }
        return run(
            Entries {
                reader: BufReader::new(
                    File::open(p).with_context(|| format!("opening {}", p.display()))?,
                ),
                extractor,
                pending: None,
                warned_cap: false,
            },
            re,
            invert,
            count,
            prefix,
        );
    }
    let source = open_source(p)?;
    let from_ms = from.map(parse_time).transpose()?.unwrap_or(0);
    let to_ms = to.map(parse_time).transpose()?.unwrap_or(u64::MAX);
    if from_ms > to_ms {
        bail!("--from is after --to");
    }
    let (selected, _) = select_chunks(p, &source.records, from_ms, to_ms, has)?;
    // Pad the selection by one chunk on each side: an entry whose
    // timestamped first line sits at the tail of the previous chunk
    // arrives whole. (Entries spanning further are subject to the usual
    // chunk-granularity slop.)
    let mut padded = std::collections::BTreeSet::new();
    for (i, _) in &selected {
        padded.insert(i.saturating_sub(1));
        padded.insert(*i);
        padded.insert(i + 1);
    }
    let chunks: Vec<ChunkRecord> = padded
        .into_iter()
        .filter_map(|i| source.records.get(i).copied())
        .collect();
    let stream = ChunkStream {
        file: source.file,
        chunks,
        idx: 0,
        buf: Vec::new(),
        pos: 0,
    };
    run(
        Entries {
            reader: BufReader::with_capacity(1 << 20, stream),
            extractor,
            pending: None,
            warned_cap: false,
        },
        re,
        invert,
        count,
        prefix,
    )
}

#[allow(clippy::too_many_arguments)]
pub fn cmd_grep(
    pattern: &str,
    files: &[std::path::PathBuf],
    from: Option<&str>,
    to: Option<&str>,
    has: &[String],
    ignore_case: bool,
    invert: bool,
    fixed: bool,
    count: bool,
    no_filename: bool,
    ts_regex: Option<&str>,
    ts_format: Option<&str>,
) -> anyhow::Result<()> {
    let pat = if fixed {
        regex::escape(pattern)
    } else {
        pattern.to_string()
    };
    let re = RegexBuilder::new(&pat)
        .case_insensitive(ignore_case)
        .multi_line(true)
        .build()
        .with_context(|| format!("bad pattern {pattern:?}"))?;

    if files.is_empty() {
        if from.is_some() || to.is_some() || !has.is_empty() {
            bail!("--from/--to/--has need a timberfs log or bundle, not stdin");
        }
        let extractor = Extractor::new(ts_regex, ts_format, false)?;
        let stdin = io::stdin();
        let matched = run(
            Entries {
                reader: stdin.lock(),
                extractor,
                pending: None,
                warned_cap: false,
            },
            &re,
            invert,
            count,
            None,
        )?;
        if count {
            println!("{matched}");
        }
        return Ok(());
    }

    // Files process in argument order (grep semantics); multiple files get
    // grep-style "path:" line prefixes and per-file -c counts.
    let multi = files.len() > 1 && !no_filename;
    let mut total: u64 = 0;
    for p in files {
        let extractor = Extractor::new(ts_regex, ts_format, false)?;
        let label = p.display().to_string().into_bytes();
        let matched = grep_one(
            p,
            &re,
            extractor,
            from,
            to,
            has,
            invert,
            count,
            if multi { Some(&label) } else { None },
        )?;
        total += matched;
        if count && multi {
            println!("{}:{matched}", p.display());
        }
    }
    if count && !multi {
        println!("{total}");
    }
    Ok(())
}