use std::collections::BTreeMap;
use std::io::{Read as _, Seek as _, SeekFrom, Write as _};
use std::path::Path;
use anyhow::Context as _;
use bytes::Bytes;
use cid::Cid;
use ipld_core::ipld::Ipld;
use jacquard_common::types::tid::{Ticker, Tid};
use knot_record::chain::{self, CommitChunk, Emit};
use knot_record::git_refs::GIT_REFS_COLLECTION;
use knot_types::{RepoDid, UnixMicros};
use serde::{Deserialize, Serialize};
use crate::cas::atomic_write;
use crate::manifest::GenerationManifest;
use crate::plan::repo_key;
pub(crate) const DIR: &str = "firehose";
/// frame k carries seq k + 1, which makes a cursor a frame count.
pub(crate) const FRAMES: &str = "frames";
const STATE: &str = "state.json";
pub(crate) struct Framing<'a> {
ticker: Ticker,
did: &'a RepoDid,
#[cfg(test)]
at_micros: Option<u64>,
}
impl<'a> Framing<'a> {
pub(crate) fn new(did: &'a RepoDid) -> Self {
Self {
ticker: Ticker::new(),
did,
#[cfg(test)]
at_micros: None,
}
}
/// revs from a fixed clock, so two settles of the same records make the same commits
#[cfg(test)]
pub(crate) fn at(did: &'a RepoDid, micros: u64) -> Self {
Self {
at_micros: Some(micros),
..Self::new(did)
}
}
}
impl Emit for Framing<'_> {
fn rev(&mut self, prev: Option<&Tid>) -> Tid {
#[cfg(test)]
if let Some(micros) = &mut self.at_micros {
*micros += 1;
return Tid::from_time(*micros, 0);
}
self.ticker.next(prev.cloned())
}
fn frame(&self, chunk: &CommitChunk) -> Option<Bytes> {
match commit_frame(self.did, chunk) {
Ok(frame) => Some(Bytes::from(frame)),
Err(error) => {
tracing::warn!(repo = %self.did, "no firehose frame for a commit: {error:#}");
None
}
}
}
fn committed(&mut self, _frame: Option<Bytes>) {}
}
fn commit_frame(did: &RepoDid, chunk: &CommitChunk) -> anyhow::Result<Vec<u8>> {
let commit = chain::commit_cid(&chunk.commit)?;
let mut blocks = chunk.blocks.clone();
blocks.insert(commit, Bytes::from(chunk.commit.clone()));
let link = |cid: Option<Cid>| cid.map_or(Ipld::Null, Ipld::Link);
let ops = chunk.ops.iter().map(|op| {
let mut fields = BTreeMap::from([
("action".to_string(), Ipld::String(action(op).into())),
("path".to_string(), Ipld::String(op.key().as_str().into())),
("cid".to_string(), link(op.cid())),
]);
if let Some(prev) = op.prev() {
fields.insert("prev".to_string(), Ipld::Link(prev));
}
Ipld::Map(fields)
});
let mut body = BTreeMap::from([
("seq".to_string(), Ipld::Integer(0)),
("rebase".to_string(), Ipld::Bool(false)),
(
"tooBig".to_string(),
Ipld::Bool(chain::blocks_bytes(&chunk.blocks) > chain::MAX_COMMIT_BLOCKS_BYTES),
),
("repo".to_string(), Ipld::String(did.as_str().into())),
("commit".to_string(), Ipld::Link(commit)),
("rev".to_string(), Ipld::String(chunk.rev.to_string())),
(
"since".to_string(),
chunk
.prev_rev
.as_ref()
.map_or(Ipld::Null, |rev| Ipld::String(rev.to_string())),
),
(
"blocks".to_string(),
Ipld::Bytes(knot_record::car_bytes(commit, &blocks)),
),
("ops".to_string(), Ipld::List(ops.collect())),
("blobs".to_string(), Ipld::List(Vec::new())),
(
"time".to_string(),
Ipld::String(datetime(chunk.rev.timestamp())),
),
]);
if let Some(prev_data) = chunk.prev_data {
body.insert("prevData".to_string(), Ipld::Link(prev_data));
}
let header = Ipld::Map(BTreeMap::from([
("op".to_string(), Ipld::Integer(1)),
("t".to_string(), Ipld::String("#commit".into())),
]));
let mut frame = serde_ipld_dagcbor::to_vec(&header)?;
frame.extend(serde_ipld_dagcbor::to_vec(&Ipld::Map(body))?);
Ok(frame)
}
fn action(op: &chain::Op) -> &'static str {
match op {
chain::Op::Create { .. } => "create",
chain::Op::Update { .. } => "update",
chain::Op::Delete { .. } => "delete",
}
}
/// atproto wants millisecond datetimes ending in `Z`.
fn datetime(micros: u64) -> String {
let at = chrono::DateTime::from_timestamp_micros(micros as i64).unwrap_or_default();
at.to_rfc3339_opts(chrono::SecondsFormat::Millis, true)
}
#[derive(Debug, Default, Serialize, Deserialize, PartialEq)]
struct State {
len: u64,
next_seq: u64,
repos: BTreeMap<String, Appended>,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
struct Appended {
commit: String,
after: u64,
}
/// the first append writes no frames, it only records each repo's served commit.
pub(crate) fn append(output_root: &Path, heads: Vec<(String, String)>) -> anyhow::Result<usize> {
let dir = output_root.join(DIR);
std::fs::create_dir_all(&dir).with_context(|| format!("create {}", dir.display()))?;
let state_path = dir.join(STATE);
let Some(mut state) = std::fs::read(&state_path)
.ok()
.map(|bytes| serde_json::from_slice::<State>(&bytes))
.transpose()
.context("read the firehose state")?
else {
// `after` stops at the served commit. newer ones wait for a later activation.
let chain_root = output_root.join("atproto");
let repos = heads
.into_iter()
.map(|(did, commit)| {
let served = live_frames(&chain_root, &did, &commit, 0)
.with_context(|| format!("the firehose frames of {did}"))?;
let after = served.last().map_or(0, |(rev, _)| *rev);
Ok((did, Appended { commit, after }))
})
.collect::<anyhow::Result<_>>()?;
let state = State {
len: 0,
next_seq: 1,
repos,
};
atomic_write(&state_path, &serde_json::to_vec(&state)?)?;
return Ok(0);
};
let mut frames = std::fs::OpenOptions::new()
.create(true)
.read(true)
.append(true)
.open(dir.join(FRAMES))
.context("open the firehose frames")?;
recover(&mut frames, &mut state)?;
let chain_root = output_root.join("atproto");
let mut appended = 0;
for (did, commit) in heads {
if state
.repos
.get(&did)
.is_some_and(|known| known.commit == commit)
{
continue;
}
let after = state.repos.get(&did).map_or(0, |known| known.after);
let live = live_frames(&chain_root, &did, &commit, after)
.with_context(|| format!("the firehose frames of {did}"))?;
let mut newest = after;
for (rev, frame) in live {
let wire = ws_frame(&with_seq(&frame, state.next_seq)?);
frames
.write_all(&wire)
.context("append to the firehose frames")?;
state.len += wire.len() as u64;
state.next_seq += 1;
newest = rev;
appended += 1;
}
state.repos.insert(
did,
Appended {
commit,
after: newest,
},
);
}
frames.sync_data().context("sync the firehose frames")?;
atomic_write(&state_path, &serde_json::to_vec(&state)?)?;
Ok(appended)
}
pub(crate) fn heads(manifest: &GenerationManifest) -> Vec<(String, String)> {
let served = manifest
.profile
.get("atproto")
.and_then(|heads| heads.as_object());
manifest
.repos
.iter()
.filter_map(|entry| {
let commit = served?.get(&entry.repo_did)?.as_str()?;
Some((entry.repo_did.clone(), commit.to_string()))
})
.collect()
}
fn open_chain(chain_root: &Path, did: &str) -> anyhow::Result<Option<knot_git::Repo>> {
let path = chain_root.join(repo_key(did)).join("chain.git");
if !path.exists() {
return Ok(None);
}
Ok(Some(knot_git::Repo::open(&path)?))
}
fn live_frames(
chain_root: &Path,
did: &str,
commit: &str,
after: u64,
) -> anyhow::Result<Vec<(u64, Vec<u8>)>> {
let Some(repo) = open_chain(chain_root, did)? else {
return Ok(Vec::new());
};
let Some(tip) = chain::tip(&repo)? else {
return Ok(Vec::new());
};
let found = chain::frames_after(&repo, tip.git, UnixMicros::new(after), usize::MAX).found;
let mut live = Vec::new();
let mut reached = false;
for (rev, frame) in found {
reached = reached || commit_of(&frame)?.to_string() == commit;
if reached {
live.push((rev.timestamp(), frame.to_vec()));
}
}
live.reverse();
Ok(live)
}
fn split(frame: &[u8]) -> anyhow::Result<(Ipld, BTreeMap<String, Ipld>)> {
let mut read = frame;
let header: Ipld = serde_ipld_dagcbor::de::from_reader_once(&mut read)?;
match serde_ipld_dagcbor::from_slice(read)? {
Ipld::Map(body) => Ok((header, body)),
_ => anyhow::bail!("a frame body isn't a map"),
}
}
fn commit_of(frame: &[u8]) -> anyhow::Result<Cid> {
match split(frame)?.1.get("commit") {
Some(Ipld::Link(cid)) => Ok(*cid),
_ => anyhow::bail!("a frame names no commit"),
}
}
fn with_seq(frame: &[u8], seq: u64) -> anyhow::Result<Vec<u8>> {
let (header, mut body) = split(frame)?;
body.insert("seq".to_string(), Ipld::Integer(seq.into()));
let mut out = serde_ipld_dagcbor::to_vec(&header)?;
out.extend(serde_ipld_dagcbor::to_vec(&Ipld::Map(body))?);
Ok(out)
}
/// the frames a reader counts: whole ones from an offset up to the recorded length
#[derive(Debug, Default, PartialEq)]
pub(crate) struct Pushes {
/// the log is shorter than the offset asked for, so these start at its beginning
pub(crate) restarted: bool,
pub(crate) end: u64,
/// the utc day of each frame that moves a repo's refs, `yyyy-mm-dd`
pub(crate) days: Vec<String>,
}
pub(crate) fn pushes_after(output_root: &Path, offset: u64) -> anyhow::Result<Pushes> {
let dir = output_root.join(DIR);
let Some(state) = std::fs::read(dir.join(STATE))
.ok()
.map(|bytes| serde_json::from_slice::<State>(&bytes))
.transpose()
.context("read the firehose state")?
else {
return Ok(Pushes::default());
};
let restarted = offset > state.len;
let start = if restarted { 0 } else { offset };
let mut bytes = Vec::new();
if state.len > start {
let mut frames =
std::fs::File::open(dir.join(FRAMES)).context("open the firehose frames")?;
frames.seek(SeekFrom::Start(start))?;
frames
.take(state.len - start)
.read_to_end(&mut bytes)
.context("read the firehose frames")?;
}
let mut days = Vec::new();
let mut at = 0;
while let Some((header, payload)) = ws_header(&bytes[at..]) {
let end = at + header + payload as usize;
if end > bytes.len() {
break;
}
let (_, body) = split(&bytes[at + header..end])?;
let moves_refs = matches!(body.get("ops"), Some(Ipld::List(ops)) if ops.iter().any(|op| {
matches!(op, Ipld::Map(op) if matches!(op.get("path"),
Some(Ipld::String(path)) if path.split('/').next() == Some(GIT_REFS_COLLECTION)))
}));
if let (true, Some(Ipld::String(time))) = (moves_refs, body.get("time")) {
days.extend(time.get(..10).map(str::to_string));
}
at = end;
}
Ok(Pushes {
restarted,
end: start + at as u64,
days,
})
}
fn ws_frame(payload: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(payload.len() + 10);
out.push(0x82);
match payload.len() {
len if len < 126 => out.push(len as u8),
len if len <= 0xffff => {
out.push(126);
out.extend((len as u16).to_be_bytes());
}
len => {
out.push(127);
out.extend((len as u64).to_be_bytes());
}
}
out.extend_from_slice(payload);
out
}
fn ws_header(bytes: &[u8]) -> Option<(usize, u64)> {
match *bytes.get(1)? {
126 => Some((
4,
u16::from_be_bytes(bytes.get(2..4)?.try_into().ok()?).into(),
)),
127 => Some((10, u64::from_be_bytes(bytes.get(2..10)?.try_into().ok()?))),
len => Some((2, len.into())),
}
}
/// the state is written last, and a crash can leave whole frames past its length.
fn recover(frames: &mut std::fs::File, state: &mut State) -> anyhow::Result<()> {
let len = frames.metadata()?.len();
if len == state.len {
return Ok(());
}
anyhow::ensure!(
len > state.len,
"the firehose frames are shorter than the {} bytes recorded",
state.len
);
let mut tail = Vec::new();
frames.seek(SeekFrom::Start(state.len))?;
frames.read_to_end(&mut tail)?;
let mut at = 0;
while let Some((header, payload)) = ws_header(&tail[at..]) {
let end = at + header + payload as usize;
if end > tail.len() {
break;
}
let (_, body) = split(&tail[at + header..end])?;
let (Some(Ipld::String(repo)), Some(Ipld::Link(commit)), Some(Ipld::String(rev))) =
(body.get("repo"), body.get("commit"), body.get("rev"))
else {
anyhow::bail!("a frame past the recorded length isn't a commit");
};
let after = rev
.parse::<Tid>()
.map_err(|error| anyhow::anyhow!("{error}"))?
.timestamp();
state.repos.insert(
repo.clone(),
Appended {
commit: commit.to_string(),
after,
},
);
state.next_seq += 1;
at = end;
}
state.len += at as u64;
frames
.set_len(state.len)
.context("cut a torn frame off the firehose")?;
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::atproto::SettledRepo;
use crate::atproto::tests::{Fixture, fixture, git, settle_fixture};
fn advance(fixture: &Fixture, text: &str) {
std::fs::write(fixture.work.join("a"), text).unwrap();
git(&fixture.work, &["commit", "-q", "-am", text]);
let work = fixture.work.to_str().unwrap();
git(
&fixture.source.export_dir,
&["fetch", "-q", work, "+refs/heads/main:refs/heads/main"],
);
}
fn served(fixture: &Fixture, settled: &SettledRepo) -> Vec<(String, String)> {
vec![(fixture.source.did.to_string(), settled.commit.to_string())]
}
fn frames_path(root: &Path) -> std::path::PathBuf {
root.join(DIR).join(FRAMES)
}
fn read_frames(root: &Path) -> Vec<(Ipld, BTreeMap<String, Ipld>)> {
let bytes = std::fs::read(frames_path(root)).unwrap_or_default();
let mut frames = Vec::new();
let mut at = 0;
while at < bytes.len() {
assert_eq!(bytes[at], 0x82, "a final binary frame");
let (header, payload) = ws_header(&bytes[at..]).unwrap();
let end = at + header + payload as usize;
frames.push(split(&bytes[at + header..end]).unwrap());
at = end;
}
assert_eq!(at, bytes.len());
frames
}
fn tip_data(fixture: &Fixture) -> Cid {
let did = fixture.source.did.as_str();
let repo = open_chain(&fixture.chain_root, did).unwrap().unwrap();
chain::tip(&repo).unwrap().unwrap().data
}
fn two_commits() -> (Fixture, SettledRepo, SettledRepo, Cid) {
let fixture = fixture("sha1");
let root = fixture.chain_root.parent().unwrap().to_path_buf();
let first = settle_fixture(&fixture);
let first_data = tip_data(&fixture);
assert_eq!(append(&root, served(&fixture, &first)).unwrap(), 0);
advance(&fixture, "two");
let second = settle_fixture(&fixture);
(fixture, first, second, first_data)
}
#[test]
fn a_websocket_header_carries_each_length_encoding() {
for len in [0, 125, 126, 0xffff, 0x10000] {
let frame = ws_frame(&vec![7; len]);
let (header, payload) = ws_header(&frame).unwrap();
assert_eq!((header + len, payload), (frame.len(), len as u64));
assert_eq!(ws_header(&frame[..header - 1]), None);
}
}
#[test]
fn an_activation_appends_the_commits_its_generation_made_live() {
let (fixture, first, second, first_data) = two_commits();
let root = fixture.chain_root.parent().unwrap();
assert!(
read_frames(root).is_empty(),
"the first append starts at the served commit"
);
assert_eq!(
append(root, served(&fixture, &first)).unwrap(),
0,
"a commit no generation serves yet stays out"
);
assert_eq!(append(root, served(&fixture, &second)).unwrap(), 1);
assert_eq!(append(root, served(&fixture, &second)).unwrap(), 0);
let frames = read_frames(root);
assert_eq!(frames.len(), 1);
let (header, body) = &frames[0];
assert_eq!(
header,
&Ipld::Map(BTreeMap::from([
("op".to_string(), Ipld::Integer(1)),
("t".to_string(), Ipld::String("#commit".into())),
]))
);
assert_eq!(body["seq"], Ipld::Integer(1));
assert_eq!(body["repo"], Ipld::String(fixture.source.did.to_string()));
assert_eq!(body["commit"], Ipld::Link(second.commit));
assert_eq!(body["rev"], Ipld::String(second.rev.clone()));
assert_eq!(body["since"], Ipld::String(first.rev.clone()));
assert_eq!(body["prevData"], Ipld::Link(first_data));
let Ipld::List(ops) = &body["ops"] else {
panic!("ops")
};
let updates = ops.iter().filter(|op| {
matches!(op, Ipld::Map(op) if op["action"] == Ipld::String("update".into())
&& matches!(op.get("prev"), Some(Ipld::Link(_))))
});
assert_eq!(
updates.count(),
1,
"the refs record moved, naming its prev: {ops:?}"
);
let Ipld::Bytes(car) = &body["blocks"] else {
panic!("blocks")
};
let car = knot_record::car_blocks(car).unwrap();
assert_eq!(car.roots, vec![second.commit]);
assert!(car.blocks.contains_key(&second.commit));
}
#[test]
fn a_torn_frame_is_cut_off_and_a_whole_unrecorded_one_kept() {
let (fixture, _, second, _) = two_commits();
let root = fixture.chain_root.parent().unwrap();
let state = std::fs::read(root.join(DIR).join(STATE)).unwrap();
assert_eq!(append(root, served(&fixture, &second)).unwrap(), 1);
let whole = std::fs::read(frames_path(root)).unwrap();
atomic_write(&root.join(DIR).join(STATE), &state).unwrap();
let mut torn = whole.clone();
torn.extend_from_slice(&whole[..whole.len() / 2]);
std::fs::OpenOptions::new()
.append(true)
.open(frames_path(root))
.unwrap()
.write_all(&whole[..whole.len() / 2])
.unwrap();
assert_eq!(std::fs::read(frames_path(root)).unwrap(), torn);
assert_eq!(append(root, served(&fixture, &second)).unwrap(), 0);
assert_eq!(std::fs::read(frames_path(root)).unwrap(), whole);
advance(&fixture, "three");
let third = settle_fixture(&fixture);
assert_eq!(append(root, served(&fixture, &third)).unwrap(), 1);
let frames = read_frames(root);
assert_eq!(frames.len(), 2);
assert_eq!(frames[1].1["seq"], Ipld::Integer(2));
assert_eq!(frames[1].1["since"], Ipld::String(second.rev.clone()));
}
#[test]
fn the_first_append_starts_at_the_served_commit_not_the_chain_tip() {
let fixture = fixture("sha1");
let root = fixture.chain_root.parent().unwrap();
let first = settle_fixture(&fixture);
advance(&fixture, "two");
let second = settle_fixture(&fixture);
assert_eq!(append(root, served(&fixture, &first)).unwrap(), 0);
assert_eq!(append(root, served(&fixture, &second)).unwrap(), 1);
let frames = read_frames(root);
assert_eq!(frames[0].1["commit"], Ipld::Link(second.commit));
assert_eq!(frames[0].1["since"], Ipld::String(first.rev.clone()));
}
#[test]
fn pushes_are_counted_from_an_offset_by_the_day_their_refs_moved() {
let (fixture, _, second, _) = two_commits();
let root = fixture.chain_root.parent().unwrap();
assert_eq!(pushes_after(root, 0).unwrap(), Pushes::default());
assert_eq!(append(root, served(&fixture, &second)).unwrap(), 1);
let Ipld::String(time) = &read_frames(root)[0].1["time"] else {
panic!("time")
};
let len = std::fs::metadata(frames_path(root)).unwrap().len();
let all = pushes_after(root, 0).unwrap();
assert_eq!(
all,
Pushes {
restarted: false,
end: len,
days: vec![time[..10].to_string()],
}
);
assert_eq!(pushes_after(root, len).unwrap().days, Vec::<String>::new());
let past = pushes_after(root, len + 1).unwrap();
assert!(past.restarted);
assert_eq!(past.days.len(), 1);
}
}