johnandersen777.bsky.social/knotless
git that lives in your Atmosphere account
Forked from chadtmiller.com/knotless
0
0
//! Searching the code of public repositories, with Zoekt
//! (github.com/sourcegraph/zoekt): a trigram index, its shards on disk,
//! that answers substrings and regular expressions across every
//! repository at once.
//!
//! knotless-serve runs `zoekt-webserver` over the shards' directory, on a
//! port of its own on loopback, and asks it in JSON
//! ([`Code::search`]); `zoekt-git-index` writes a repository's shards
//! from its mirror's default branch ([`Code::index`]). The webserver
//! loads the shards as they are written, so a push is searched once its
//! repository is indexed again. Neither is reached from outside.
//!
//! Which repositories are indexed follows the index ([`follow`]): each
//! public one with no shards yet once at start, then each again as its
//! refs or its record move. One gone from the index, or taken down, has
//! its shards deleted as it is indexed again, and is never answered
//! meanwhile: what Zoekt finds is shown only for repositories the index
//! lists ([`crate::index::Index::listed`]).
use crate::{App, AppError};
use anyhow::{Context, Result};
use base64::Engine;
use futures::StreamExt;
use knotless::status::CodeIndex;
use knotless::telemetry::{self, Fields, Level};
use serde::{Deserialize, Serialize};
use serde_json::json;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::sync::Arc;
use std::time::Duration;
/// How long a search may take in Zoekt.
const WALL: Duration = Duration::from_secs(10);
/// The most files counted by language for a search.
const FACETED: usize = 1000;
/// How long `zoekt-webserver` has to start answering.
const START: Duration = Duration::from_secs(30);
/// Zoekt, over the shards in `dir`.
pub struct Code {
dir: PathBuf,
url: String,
http: reqwest::Client,
_server: tokio::process::Child,
}
/// A file Zoekt found, in the repository with `key` of `did`.
#[derive(Debug)]
pub struct Hit {
pub did: String,
pub key: String,
pub path: String,
pub language: String,
/// The commit of its default branch indexed, which its lines are of.
pub commit: String,
pub chunks: Vec<Chunk>,
}
/// Lines of a file around what was found in them.
#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct Chunk {
/// The first line's number, from 1.
pub start_line: u32,
/// Its lines, joined by newlines, with none at the end.
pub content: String,
pub ranges: Vec<Found>,
/// What its highlighting paints over `content`, as a file's is
/// encoded ([`crate::highlight::Ranges::encode`]): the lines read
/// alone, without the rest of their file.
#[serde(skip_serializing_if = "Option::is_none")]
pub highlights: Option<String>,
}
/// What was found on one line, by UTF-8 bytes into it.
#[derive(Serialize, Debug, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct Found {
pub line: u32,
pub byte_start: usize,
pub byte_end: usize,
}
/// What a search found: the files shown, and how many files and matches
/// there were in all.
#[derive(Debug, Default)]
pub struct Results {
pub hits: Vec<Hit>,
pub file_count: u64,
pub match_count: u64,
}
// What Zoekt's JSON API answers, as its Go types are named.
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct Reply {
result: Option<ZResult>,
error: Option<String>,
}
#[derive(Deserialize, Default)]
#[serde(rename_all = "PascalCase")]
struct ZResult {
file_count: u64,
match_count: u64,
#[serde(default)]
files: Option<Vec<ZFile>>,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct ZFile {
#[serde(default)]
score: f64,
file_name: String,
repository: String,
language: String,
#[serde(default)]
version: String,
#[serde(default)]
chunk_matches: Option<Vec<ZChunk>>,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct ZChunk {
content: String,
content_start: ZLocation,
ranges: Vec<ZRange>,
file_name: bool,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct ZLocation {
byte_offset: usize,
line_number: u32,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct ZRange {
start: ZLocation,
end: ZLocation,
}
impl Code {
/// Starts `zoekt-webserver` over `dir`, made if it is missing, and
/// waits for it to answer.
pub async fn start(dir: &Path) -> Result<Code> {
std::fs::create_dir_all(dir.join("meta"))?;
let port = std::net::TcpListener::bind("127.0.0.1:0")?
.local_addr()?
.port();
let listen = format!("127.0.0.1:{port}");
let server = tokio::process::Command::new("zoekt-webserver")
.arg("-index")
.arg(dir)
.args(["-listen", &listen, "-rpc", "-html=false"])
.stdout(Stdio::null())
.stderr(Stdio::null())
.kill_on_drop(true)
.spawn()
.context("zoekt-webserver, which code search runs, is not on PATH")?;
let url = format!("http://{listen}");
let http = reqwest::Client::new();
let health = format!("{url}/healthz");
let up = async {
while http.get(&health).send().await.is_err() {
tokio::time::sleep(Duration::from_millis(50)).await;
}
};
tokio::time::timeout(START, up)
.await
.context("zoekt-webserver did not answer")?;
Ok(Code {
dir: dir.to_path_buf(),
url,
http,
_server: server,
})
}
/// Writes the shards of `did`'s repository with `key` from its mirror
/// at `git`, of its default branch, over any it had: a commit indexed
/// already is left as it is.
pub async fn index(&self, did: &str, key: &str, git: &Path) -> Result<()> {
let prefix = prefix(did, key);
let meta = self.dir.join("meta").join(format!("{prefix}.json"));
std::fs::write(&meta, json!({ "Name": format!("{did}/{key}") }).to_string())?;
let out = tokio::process::Command::new("zoekt-git-index")
.arg("-index")
.arg(&self.dir)
.args(["-shard_prefix_override", &prefix, "-meta"])
.arg(&meta)
.args(["-branches", "HEAD", "-submodules=false", "-disable_ctags"])
.arg(git)
.output()
.await
.context("zoekt-git-index, which code search runs, is not on PATH")?;
let said = String::from_utf8_lossy(&out.stderr);
anyhow::ensure!(out.status.success(), "zoekt-git-index: {said}");
Ok(())
}
/// The shards of `did`'s repository with `key`.
fn shards(&self, did: &str, key: &str) -> Vec<PathBuf> {
let start = format!("{}_v", prefix(did, key));
let entries = std::fs::read_dir(&self.dir).into_iter().flatten().flatten();
let named = |p: &PathBuf| {
p.file_name()
.is_some_and(|n| n.to_string_lossy().starts_with(&start))
};
entries.map(|e| e.path()).filter(named).collect()
}
/// The index as `/_status` tells of it: whether Zoekt answers, and
/// the repositories, shards and bytes it holds.
pub async fn check(&self) -> CodeIndex {
let mut repos = std::collections::BTreeSet::new();
let (mut shards, mut bytes) = (0, 0);
for entry in std::fs::read_dir(&self.dir).into_iter().flatten().flatten() {
let name = entry.file_name().to_string_lossy().into_owned();
// Its `meta` directory is no shard.
let Some((repo, _)) = name
.strip_suffix(".zoekt")
.and_then(|n| n.rsplit_once("_v"))
else {
continue;
};
repos.insert(repo.to_string());
shards += 1;
bytes += entry.metadata().map_or(0, |m| m.len());
}
let health = format!("{}/healthz", self.url);
let answered = self.http.get(health).send().await;
let error = answered.and_then(|r| r.error_for_status()).err();
CodeIndex {
ok: error.is_none(),
error: error.as_ref().map(ToString::to_string),
repos: repos.len() as u64,
shards,
bytes,
}
}
/// Whether `did`'s repository with `key` has shards.
pub fn has(&self, did: &str, key: &str) -> bool {
!self.shards(did, key).is_empty()
}
/// Deletes the shards of `did`'s repository with `key`.
pub fn forget(&self, did: &str, key: &str) {
for shard in self.shards(did, key) {
let _ = std::fs::remove_file(shard);
}
let _ = std::fs::remove_file(
self.dir
.join("meta")
.join(format!("{}.json", prefix(did, key))),
);
}
/// The files at most `limit` that `q`, in Zoekt's query language,
/// finds, in `did`'s repository with `key` alone if one is named: 400
/// for a query Zoekt cannot read.
pub async fn search(
&self,
q: &str,
repo: Option<(&str, &str)>,
limit: usize,
) -> Result<Results, AppError> {
let opts = json!({
"ChunkMatches": true,
"NumContextLines": 1,
"MaxDocDisplayCount": limit,
"MaxWallTime": WALL.as_nanos() as u64,
});
let found = self.ask(within(q, repo), opts).await?;
let mut files = found.files.unwrap_or_default();
// The best first, and files Zoekt scores alike, as its shards
// answered in parallel, by repository and path: so each page
// after the first goes on from it rather than showing it again.
files.sort_by(|a, b| {
let named = a
.repository
.cmp(&b.repository)
.then(a.file_name.cmp(&b.file_name));
b.score.total_cmp(&a.score).then(named)
});
let hits = files.into_iter().filter_map(hit).collect();
Ok(Results {
hits,
file_count: found.file_count,
match_count: found.match_count,
})
}
/// The files `q` finds in any language, [`FACETED`] at most, with no
/// lines: what its `lang:` filters leave out counted too, so each
/// language can be picked from them.
pub async fn faceted(&self, q: &str, repo: Option<(&str, &str)>) -> Result<Vec<Hit>, AppError> {
let opts = json!({
"MaxDocDisplayCount": FACETED,
"MaxWallTime": WALL.as_nanos() as u64,
});
let q = format!("type:file ({})", unlanguaged(q));
let found = self.ask(within(&q, repo), opts).await?;
let files = found.files.unwrap_or_default().into_iter();
let bare = files.map(|f| ZFile {
chunk_matches: None,
..f
});
Ok(bare.filter_map(hit).collect())
}
/// What Zoekt answers `q` with `opts`.
async fn ask(&self, q: String, opts: serde_json::Value) -> Result<ZResult, AppError> {
let asked = json!({ "Q": q, "Opts": opts });
let url = format!("{}/api/search", self.url);
let res = self.http.post(url).json(&asked).send().await;
let res = res
.map_err(anyhow::Error::from)
.map_err(AppError::internal)?;
// Zoekt's status, as a query it cannot read is 400.
let status = res.status();
let reply = res.json::<Reply>().await;
let reply = reply
.map_err(anyhow::Error::from)
.map_err(AppError::internal)?;
if let Some(error) = reply.error {
return Err(AppError(status, format!("code search: {error}")));
}
Ok(reply.result.unwrap_or_default())
}
}
/// A file as Zoekt found it, if its repository is named as [`Code::index`]
/// names one.
fn hit(file: ZFile) -> Option<Hit> {
let (did, key) = file.repository.rsplit_once('/')?;
let chunks = file.chunk_matches.unwrap_or_default();
let mut chunks: Vec<Chunk> = chunks
.into_iter()
.filter(|c| !c.file_name)
.map(|c| chunk(c, &file.file_name))
.collect();
// Zoekt answers them best first: shown as the file reads.
chunks.sort_by_key(|c| c.start_line);
Some(Hit {
did: did.to_string(),
key: key.to_string(),
path: file.file_name,
language: file.language,
commit: file.version,
chunks,
})
}
/// A chunk's lines, and what was found, by line, out of Zoekt's offsets
/// into the file.
fn chunk(c: ZChunk, path: &str) -> Chunk {
let bytes = base64::engine::general_purpose::STANDARD
.decode(&c.content)
.unwrap_or_default();
let bytes = bytes.strip_suffix(b"\n").unwrap_or(&bytes);
let base = c.content_start.byte_offset;
let mut starts = Vec::new();
let mut at = 0;
let lines: Vec<&[u8]> = bytes.split(|b| *b == b'\n').collect();
for line in &lines {
starts.push(at);
at += line.len() + 1;
}
let first = c.content_start.line_number;
let mut ranges = Vec::new();
for r in &c.ranges {
let (s, e) = (r.start.byte_offset - base, r.end.byte_offset - base);
for (i, (start, line)) in starts.iter().zip(&lines).enumerate() {
let (a, b) = (s.max(*start), e.min(start + line.len()));
if a < b {
ranges.push(Found {
line: first + i as u32,
byte_start: a - start,
byte_end: b - start,
});
}
}
}
let content = String::from_utf8_lossy(bytes).into_owned();
let painted = crate::highlight::code(path, &content);
Chunk {
start_line: first,
highlights: painted.as_deref().map(crate::highlight::Ranges::encode),
content,
ranges,
}
}
/// `q` kept to `did`'s repository with `key`, if one is named.
fn within(q: &str, repo: Option<(&str, &str)>) -> String {
match repo {
Some((did, key)) => format!("repo:^{}$ ({q})", escape(&format!("{did}/{key}"))),
None => q.to_string(),
}
}
/// `q` without its `lang:` and `-lang:` filters, its words joined by
/// single spaces, or `q` itself when nothing else is left of it. An
/// operator joining a filter to the rest goes with it, and a `)` it ends
/// with stays, so what is left still reads as a query.
fn unlanguaged(q: &str) -> String {
let operator = |w: &str| matches!(w, "or" | "and" | "not");
let mut kept: Vec<&str> = Vec::new();
// A filter was just taken out with no operator before it: one after
// it goes instead.
let mut after = false;
let mut words = q.split_whitespace();
while let Some(word) = words.next() {
let filter = word.strip_prefix('-').unwrap_or(word).strip_prefix("lang:");
let Some(name) = filter else {
if !(after && operator(word)) {
kept.push(word);
}
after = false;
continue;
};
// A quoted name runs on to the word its quote closes in.
let mut end = name;
if name.starts_with('"') && !name[1..].contains('"') {
end = words.by_ref().find(|w| w.contains('"')).unwrap_or("");
}
let closing = end.len() - end.trim_end_matches(')').len();
after = !kept.last().is_some_and(|w| operator(w));
if !after {
kept.pop();
}
kept.extend(std::iter::repeat_n(")", closing));
}
if kept.is_empty() {
return q.to_string();
}
kept.join(" ")
}
/// The prefix of the shards of `did`'s repository with `key`: `did/key`,
/// each byte but letters, digits, `.` and `-` written `_` and its hex, so
/// that it is a file name Zoekt keeps as it is, and no two are one.
fn prefix(did: &str, key: &str) -> String {
format!("{did}/{key}")
.bytes()
.map(|b| match b {
b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'.' | b'-' => (b as char).to_string(),
_ => format!("_{b:02x}"),
})
.collect()
}
/// `s` as a regular expression matching it alone.
fn escape(s: &str) -> String {
s.chars()
.map(|c| match c.is_ascii_alphanumeric() {
true => c.to_string(),
false => format!("\\{c}"),
})
.collect()
}
/// What indexes the public repositories as the index asks
/// ([`Index::reindex`]), when serve keeps a code index: it asks first for
/// each with no shards yet, then answers each ask, for as long as serve
/// runs. Set up here, before it is spawned, so serve without a code
/// index starts nothing.
///
/// [`Index::reindex`]: crate::index::Index::reindex
pub fn follow(app: Arc<App>) -> Option<impl Future<Output = ()>> {
let (code, index) = app.code.clone().zip(app.index.clone())?;
let asked = index.indexing()?;
for (did, key) in index.code_repos().unwrap_or_default() {
if !code.has(&did, &key) {
index.reindex(&did, &key);
}
}
let asked = futures::stream::unfold(asked, |mut a| async { a.recv().await.map(|r| (r, a)) });
Some(asked.for_each(move |(did, key)| indexed(app.clone(), code.clone(), did, key)))
}
/// Indexes `did`'s repository with `key` again, as [`follow`] is asked,
/// and notes how it went.
async fn indexed(app: Arc<App>, code: Arc<Code>, did: String, key: String) {
let error = reindex(&app, &code, &did, &key).await.err();
let error = error.as_ref().map(ToString::to_string);
let level = [Level::INFO, Level::WARN][error.is_some() as usize];
let fields = Fields {
did: Some(&did),
key: Some(&key),
error,
..Fields::default()
};
telemetry::note(level, "code indexed", &fields);
}
/// Indexes `did`'s repository with `key` again, or forgets it if the
/// index no longer lists it or its default branch has no commit to index.
async fn reindex(app: &App, code: &Code, did: &str, key: &str) -> Result<()> {
let index = app.index.as_ref().context("no index")?;
index.reindexed(did, key);
let Some(listed) = index.listed(did, key)? else {
code.forget(did, key);
return Ok(());
};
let git = app.mirror(did, &listed.name).await.ok();
let git = git.context("its mirror could not be read")?;
if git.rev_parse("HEAD").await.is_err() {
code.forget(did, key);
return Ok(());
}
code.index(did, key, &git.dir).await
}
/// Zoekt over `dir`, as if its webserver had died: nothing answers.
#[cfg(test)]
pub(crate) fn dead(dir: &Path) -> Code {
Code {
dir: dir.to_path_buf(),
url: "http://127.0.0.1:1".into(),
http: reqwest::Client::new(),
_server: tokio::process::Command::new("true").spawn().unwrap(),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn languages_are_left_out_to_count_them() {
assert_eq!(unlanguaged("count lang:Rust"), "count");
assert_eq!(unlanguaged("lang:\"Vim Script\" a -lang:toml b"), "a b");
assert_eq!(unlanguaged("lang:\"Go\" x"), "x");
assert_eq!(unlanguaged("lang:Rust"), "lang:Rust");
// An operator joining a filter goes with it, before or after.
assert_eq!(unlanguaged("foo or lang:Rust"), "foo");
assert_eq!(unlanguaged("lang:Rust or foo"), "foo");
assert_eq!(unlanguaged("a or lang:Rust or b"), "a or b");
// A group the filter closed stays closed.
assert_eq!(unlanguaged("(foo or lang:Rust) bar"), "(foo ) bar");
assert_eq!(unlanguaged("(a or lang:\"Vim Script\") b"), "(a ) b");
assert_eq!(within("a", Some(("d", "k"))), "repo:^d\\/k$ (a)");
}
#[test]
fn prefixes_are_file_names_of_their_own() {
assert_eq!(prefix("did:plc:ab", "3k.x-y"), "did_3aplc_3aab_2f3k.x-y");
assert_ne!(prefix("did:plc:a", "b_c"), prefix("did:plc:a", "b:c"));
assert_eq!(escape("did:plc:a/k.1"), "did\\:plc\\:a\\/k\\.1");
}
#[test]
fn chunks_are_cut_into_lines() {
let content =
base64::engine::general_purpose::STANDARD.encode("fn a() {}\nlet abc = a();\n");
let at = |byte_offset, line_number| ZLocation {
byte_offset,
line_number,
};
let zc = ZChunk {
content: content.clone(),
// The chunk starts at byte 100, line 7, of its file.
content_start: at(100, 7),
// `a` in line 7, and from its `{}` to the end of the next line's `abc`.
ranges: vec![
ZRange {
start: at(103, 7),
end: at(104, 7),
},
ZRange {
start: at(107, 7),
end: at(117, 8),
},
],
file_name: false,
};
let c = chunk(zc, "src/a.rs");
assert_eq!(c.start_line, 7);
assert_eq!(c.content, "fn a() {}\nlet abc = a();");
let found = |line, byte_start, byte_end| Found {
line,
byte_start,
byte_end,
};
assert_eq!(c.ranges, [found(7, 3, 4), found(7, 7, 9), found(8, 0, 7)]);
// Painted as Rust reads, by its path: `fn` and `let` keywords.
assert!(c.highlights.unwrap().contains("keyword:0,2,"));
// A file no grammar here reads is painted not at all.
let plain = ZChunk {
content,
content_start: at(0, 1),
ranges: vec![],
file_name: false,
};
assert_eq!(chunk(plain, "notes.txt").highlights, None);
// A file's chunks, in the order of its lines.
let at_line = |line| ZChunk {
content: base64::engine::general_purpose::STANDARD.encode("x\n"),
content_start: at(0, line),
ranges: vec![],
file_name: false,
};
let found = hit(ZFile {
score: 0.0,
file_name: "a".into(),
repository: "did:plc:a/k".into(),
language: String::new(),
version: String::new(),
chunk_matches: Some(vec![at_line(30), at_line(2)]),
})
.unwrap();
let lines: Vec<u32> = found.chunks.iter().map(|c| c.start_line).collect();
assert_eq!(lines, [2, 30]);
assert!(
hit(ZFile {
score: 0.0,
file_name: "a".into(),
repository: "no-slash".into(),
language: String::new(),
version: String::new(),
chunk_matches: None,
})
.is_none()
);
}
}