use std::fs::{self, File, OpenOptions};
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 serde_json::Value;
use crate::format::ChunkRecord;
use crate::import::Extractor;
use crate::query::{fmt_ms, is_bundle, open_source, resolve_backing, select_chunks};
use crate::store;
const ENTRY_CAP: usize = 16 << 20;
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 is_match(res: &[Regex], entry: &[u8]) -> bool {
res.iter().all(|re| re.is_match(entry))
}
fn run<R: BufRead>(
mut entries: Entries<R>,
res: &[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 is_match(res, &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)
}
fn padded_chunks(
p: &Path,
records: &[ChunkRecord],
from_ms: u64,
to_ms: u64,
has: &[String],
) -> anyhow::Result<Vec<ChunkRecord>> {
let (selected, _) = select_chunks(p, records, from_ms, to_ms, has)?;
let mut padded = std::collections::BTreeSet::new();
for (i, _) in &selected {
padded.insert(i.saturating_sub(1));
padded.insert(*i);
padded.insert(i + 1);
}
Ok(padded
.into_iter()
.filter_map(|i| records.get(i).copied())
.collect())
}
fn names_timberfs_source(s: &str) -> bool {
let p = Path::new(s);
if is_bundle(p) {
return p.is_file();
}
match resolve_backing(p) {
Ok((dir, name)) => crate::format::rings_path(&dir, &name).exists(),
Err(_) => false,
}
}
fn word_regex(lit: &str, ignore_case: bool) -> anyhow::Result<Regex> {
let pat = format!(
r"(?:\A|(?-u:[^0-9A-Za-z])){}(?:(?-u:[^0-9A-Za-z])|\z)",
regex::escape(lit)
);
RegexBuilder::new(&pat)
.case_insensitive(ignore_case)
.build()
.with_context(|| format!("bad pattern {lit:?}"))
}
fn engage_auto_has(
p: &Path,
has: &[String],
auto_has: &[String],
scan_reason: Option<&str>,
records: &[ChunkRecord],
) -> anyhow::Result<Vec<String>> {
if !has.is_empty() || is_bundle(p) {
return Ok(has.to_vec());
}
let Ok((dir, base)) = resolve_backing(p) else {
return Ok(Vec::new());
};
if !crate::format::grain_path(&dir, &base).exists() {
return Ok(Vec::new());
}
if auto_has.is_empty() {
if let Some(reason) = scan_reason {
eprintln!(
"timberfs: full scan of {} chunk(s): {reason} cannot use the .grain \
token index",
records.len()
);
}
return Ok(Vec::new());
}
let tokens = crate::grain::tokenize_query(&auto_has.join(" "));
eprintln!(
"timberfs: .grain: pre-filtering {} chunk(s) on the pattern's tokens ({}) — \
chunk-level, like --has; --scan forces a full scan",
records.len(),
tokens
.iter()
.map(|t| String::from_utf8_lossy(t).into_owned())
.collect::<Vec<_>>()
.join(", ")
);
Ok(auto_has.to_vec())
}
fn command_line() -> String {
fn quote(a: &str) -> String {
let plain = !a.is_empty()
&& a.bytes()
.all(|b| b.is_ascii_alphanumeric() || b"-_./=:%+,@^".contains(&b));
if plain {
a.to_string()
} else {
format!("'{}'", a.replace('\'', "'\\''"))
}
}
std::iter::once("timberfs".to_string())
.chain(std::env::args().skip(1).map(|a| quote(&a)))
.collect::<Vec<_>>()
.join(" ")
}
#[allow(clippy::too_many_arguments)]
fn grep_into(
p: &Path,
res: &[Regex],
extractor: Extractor,
from: Option<u64>,
to: Option<u64>,
has: &[String],
auto_has: &[String],
scan_reason: Option<&str>,
invert: bool,
pattern: &str,
dest: &Path,
fail_on_empty: bool,
) -> anyhow::Result<()> {
let source = open_source(p)?;
let from_ms = from.unwrap_or(0);
let to_ms = to.unwrap_or(u64::MAX);
if from_ms > to_ms {
bail!("--from is after --to");
}
let has = engage_auto_has(p, has, auto_has, scan_reason, &source.records)?;
let chunks = padded_chunks(p, &source.records, from_ms, to_ms, &has)?;
let bundled = is_bundle(dest);
let (pair_dir, pair_name, _locks) = if bundled {
if dest.exists() {
bail!(
"{} already exists — grep --into always creates",
dest.display()
);
}
let parent = match dest.parent() {
Some(d) if !d.as_os_str().is_empty() => d.to_path_buf(),
_ => std::path::PathBuf::from("."),
};
let stem = dest
.file_stem()
.and_then(|s| s.to_str())
.context("bad bundle name")?
.to_string();
let tmp = parent.join(format!(".{stem}.grep-tmp.{}", std::process::id()));
fs::create_dir_all(&tmp)?;
(tmp, stem, None)
} else {
crate::query::ensure_dest_is_not_plain_file(dest, "grep")?;
let (d, n) = resolve_backing(dest)?;
fs::create_dir_all(&d)?;
if crate::format::rings_path(&d, &n).exists() || crate::format::trunk_path(&d, &n).exists()
{
bail!(
"{n} already exists in {} — grep --into always creates; merge with import",
d.display()
);
}
let dir_lock = store::lock_backing_shared(&d)?.with_context(|| {
format!(
"destination directory {} is served by a timberfs mount",
d.display()
)
})?;
let file_lock = store::lock_file_exclusive(&d, &n)?
.with_context(|| format!("{n} already has a writer"))?;
(d, n, Some((dir_lock, file_lock)))
};
let cleanup = |dir: &Path| {
if bundled {
let _ = fs::remove_dir_all(dir);
} else {
let _ = fs::remove_file(crate::format::trunk_path(dir, &pair_name));
let _ = fs::remove_file(crate::format::rings_path(dir, &pair_name));
}
};
let cfg = store::Config {
chunk_size: 256 * 1024,
level: 3,
flush_age_ms: 5000,
};
let mut st = store::Store {
dir: pair_dir.clone(),
cfg,
files: std::collections::BTreeMap::new(),
};
st.create(&pair_name)?;
let mut entries = Entries {
reader: BufReader::with_capacity(
1 << 20,
ChunkStream {
file: source.file,
chunks,
idx: 0,
buf: Vec::new(),
pos: 0,
},
),
extractor,
pending: None,
warned_cap: false,
};
let mut matched: u64 = 0;
let mut seen: u64 = 0;
let mut last_ts: Option<u64> = None;
let mut leading: Vec<Vec<u8>> = Vec::new(); let mut span: Option<(u64, u64)> = None;
while let Some(entry) = entries.next_entry()? {
seen += 1;
let head = String::from_utf8_lossy(&entry[..entry.len().min(256)]);
let ts = entries.extractor.extract(&head);
if is_match(res, &entry) != invert {
matched += 1;
match ts.or(last_ts) {
Some(t) => {
let f = st.files.get_mut(&pair_name).unwrap();
for l in leading.drain(..) {
f.append_stamped(&l, t, &cfg)?;
}
f.append_stamped(&entry, t, &cfg)?;
span = Some(match span {
None => (t, t),
Some((a, b)) => (a.min(t), b.max(t)),
});
}
None => leading.push(entry.clone()),
}
}
if let Some(t) = ts {
last_ts = Some(t);
}
}
if !leading.is_empty() {
cleanup(&pair_dir);
bail!(
"{} matching entr{} carried no timestamp and none followed; \
try --timestamp-regex/--timestamp-format",
leading.len(),
if leading.len() == 1 { "y" } else { "ies" }
);
}
if matched == 0 && fail_on_empty {
cleanup(&pair_dir);
bail!("no entries matched (--fail-on-empty)");
}
st.flush_all();
let mut derived = crate::bark::derived_map(source.bark.as_ref(), "grep");
if from.is_some() {
derived.insert(
"window_from".to_string(),
Value::String(crate::bark::ms_rfc3339(from_ms)),
);
}
if to.is_some() {
derived.insert(
"window_to".to_string(),
Value::String(crate::bark::ms_rfc3339(to_ms)),
);
}
derived.insert("pattern".to_string(), Value::String(pattern.to_string()));
derived.insert("command".to_string(), Value::String(command_line()));
let bark_map = crate::bark::with_identity(derived)?;
let bark_text = serde_json::to_string_pretty(&Value::Object(bark_map))? + "\n";
if bundled {
let res = assemble_bundle(&pair_dir, &pair_name, dest, &bark_text, span);
cleanup(&pair_dir);
res?;
} else {
fs::write(crate::format::bark_path(&pair_dir, &pair_name), bark_text)?;
}
match span {
Some((a, b)) => eprintln!(
"timberfs: grep: {matched} of {seen} entr{} matched, spanning {} .. {} -> {}{}",
if matched == 1 { "y" } else { "ies" },
fmt_ms(a),
fmt_ms(b),
dest.display(),
if bundled { " (bundle)" } else { "" }
),
None => eprintln!(
"timberfs: grep: 0 of {seen} entries matched — empty result; the artifact \
attests it (--fail-on-empty to error instead) -> {}{}",
dest.display(),
if bundled { " (bundle)" } else { "" }
),
}
Ok(())
}
fn assemble_bundle(
pair_dir: &Path,
name: &str,
dest: &Path,
bark_text: &str,
span: Option<(u64, u64)>,
) -> anyhow::Result<()> {
let mtime_secs = span.map(|(_, b)| b / 1000).unwrap_or_else(|| {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
});
let out = OpenOptions::new()
.write(true)
.create_new(true)
.open(dest)
.with_context(|| format!("creating {}", dest.display()))?;
let mut builder = tar::Builder::new(out);
let mut add = |member: String, bytes_len: u64, reader: &mut dyn Read| -> anyhow::Result<()> {
let mut header = tar::Header::new_ustar();
header.set_mode(0o644);
header.set_mtime(mtime_secs);
header.set_size(bytes_len);
builder.append_data(&mut header, member, reader)?;
Ok(())
};
let rings = fs::read(crate::format::rings_path(pair_dir, name))?;
add(format!("{name}.rings"), rings.len() as u64, &mut &rings[..])?;
add(
format!("{name}.bark"),
bark_text.len() as u64,
&mut bark_text.as_bytes(),
)?;
let trunk_path = crate::format::trunk_path(pair_dir, name);
let trunk_len = fs::metadata(&trunk_path)?.len();
let mut trunk = File::open(&trunk_path)?;
add(format!("{name}.trunk"), trunk_len, &mut trunk)?;
let out = builder.into_inner()?;
out.sync_all()?;
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn grep_one(
p: &Path,
res: &[Regex],
extractor: Extractor,
from: Option<u64>,
to: Option<u64>,
has: &[String],
auto_has: &[String],
scan_reason: Option<&str>,
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 {
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,
},
res,
invert,
count,
prefix,
);
}
let source = open_source(p)?;
let from_ms = from.unwrap_or(0);
let to_ms = to.unwrap_or(u64::MAX);
if from_ms > to_ms {
bail!("--from is after --to");
}
let has = engage_auto_has(p, has, auto_has, scan_reason, &source.records)?;
let chunks = padded_chunks(p, &source.records, from_ms, to_ms, &has)?;
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,
},
res,
invert,
count,
prefix,
)
}
#[allow(clippy::too_many_arguments)]
pub fn cmd_grep(
pattern: &str,
files: &[std::path::PathBuf],
from: Option<u64>,
to: Option<u64>,
has: &[String],
ignore_case: bool,
invert: bool,
fixed: bool,
count: bool,
no_filename: bool,
ts_regex: Option<&str>,
ts_format: Option<&str>,
into: Option<&Path>,
fail_on_empty: bool,
scan: bool,
regex_mode: bool,
) -> anyhow::Result<()> {
let mut files: Vec<std::path::PathBuf> = files.to_vec();
let mut pattern = pattern.to_string();
let selection_given = from.is_some() || to.is_some() || !has.is_empty();
let pattern_is_matchless = if selection_given && names_timberfs_source(&pattern) {
files.insert(0, std::path::PathBuf::from(&pattern));
if has.is_empty() {
eprintln!("timberfs: no PATTERN given; every entry in the window matches");
pattern = String::new();
} else {
eprintln!(
"timberfs: no PATTERN given; matching entries that contain: {}",
has.join(", ")
);
pattern = has.join(" AND ");
}
true
} else {
false
};
if files.is_empty() && !selection_given && names_timberfs_source(&pattern) {
use std::io::IsTerminal;
if io::stdin().is_terminal() {
bail!(
"{pattern} names a timberfs log, but it was read as the PATTERN \
(which comes first: timberfs grep PATTERN {pattern}); to select \
without a pattern, give --has/--from/--to"
);
}
}
let res: Vec<Regex> = if pattern_is_matchless {
has.iter()
.map(|t| word_regex(t, ignore_case))
.collect::<anyhow::Result<_>>()?
} else if regex_mode {
vec![RegexBuilder::new(&pattern)
.case_insensitive(ignore_case)
.multi_line(true)
.build()
.with_context(|| format!("bad pattern {pattern:?}"))?]
} else if fixed {
vec![RegexBuilder::new(®ex::escape(&pattern))
.case_insensitive(ignore_case)
.build()
.with_context(|| format!("bad pattern {pattern:?}"))?]
} else {
vec![word_regex(&pattern, ignore_case)?]
};
let auto_has: Vec<String> = if !scan
&& !invert
&& !ignore_case
&& !regex_mode
&& !fixed
&& !pattern_is_matchless
&& has.is_empty()
&& !crate::grain::tokenize_query(&pattern).is_empty()
{
vec![pattern.clone()]
} else {
Vec::new()
};
let scan_reason: Option<&str> = if auto_has.is_empty() && has.is_empty() && !scan {
if regex_mode {
Some("--regex")
} else if fixed {
Some("-F (raw substring)")
} else if ignore_case {
Some("-i (the index is exact-case)")
} else if invert {
Some("-v (non-matches must be read to be printed)")
} else {
None
}
} else {
None
};
if let Some(dest) = into {
let [p] = &files[..] else {
bail!(
"grep --into wants exactly one timberfs source \
(got {} — merge or import first)",
if files.is_empty() {
"stdin".to_string()
} else {
format!("{} files", files.len())
}
);
};
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 {
bail!(
"{} is a plain file — grep --into derives from a timberfs log or \
bundle (import it first)",
p.display()
);
}
let extractor = Extractor::new(ts_regex, ts_format, false)?;
return grep_into(
p,
&res,
extractor,
from,
to,
has,
&auto_has,
scan_reason,
invert,
&pattern,
dest,
fail_on_empty,
);
}
if files.is_empty() {
if selection_given {
bail!(
"--from/--to/--has need a timberfs log or bundle, not stdin \
(PATTERN comes first: timberfs grep PATTERN FILE)"
);
}
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,
},
&res,
invert,
count,
None,
)?;
if count {
println!("{matched}");
}
return Ok(());
}
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,
&res,
extractor,
from,
to,
has,
&auto_has,
scan_reason,
invert,
count,
if multi { Some(&label) } else { None },
)?;
total += matched;
if count && multi {
println!("{}:{matched}", p.display());
}
}
if count && !multi {
println!("{total}");
}
Ok(())
}