use std::collections::BTreeMap;
use std::num::NonZeroUsize;
use haematite::branch::current_timestamp;
use haematite::db::ReadOnlyDatabase;
use haematite::tree::Hash;
use serde::{Deserialize, Serialize};
use super::branches::{self, ForkProvenance};
use super::mutations::PreparedWrite;
use super::release::{Answer, Rules};
use super::{
BranchRefRecord, Captured, Failure, SERVICE, Storage, branch_name, dead, history, identifier,
reads,
};
use crate::serve::protocol::{Change, Stage, Version};
pub const TAG: &[u8] = b"s\0";
const FORMAT: u32 = 1;
const RELEASE_ROWS: usize = 4_096;
pub const ROW_BYTES: usize = 65_536;
const UNFINISHED: &[u8] = b"line-release\0";
const UNFINISHED_PAGE: usize = 64;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct At {
pub branch: String,
pub version: Version,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Listed {
pub branch: String,
pub version: Version,
pub stage: Stage,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
pub struct Merged {
pub from: At,
pub stages: Vec<Listed>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct Row {
format: u32,
version: Version,
#[serde(default, skip_serializing_if = "Option::is_none")]
stage: Option<Stage>,
#[serde(default, skip_serializing_if = "Option::is_none")]
merged: Option<Merged>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pieces: Option<u32>,
}
#[derive(Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct Piece {
format: u32,
version: Version,
index: u32,
stages: Vec<Listed>,
}
pub(super) type UnfinishedRow = (Vec<u8>, Vec<u8>);
#[derive(Debug, Serialize, Deserialize)]
#[serde(deny_unknown_fields)]
struct Unfinished {
branch: String,
generation: u64,
}
#[derive(Debug, Default, PartialEq, Eq)]
pub struct Lineage {
pub stage: Option<Stage>,
pub merged: Option<Merged>,
}
#[must_use]
pub fn key(version: Version) -> Vec<u8> {
let mut key = Vec::with_capacity(TAG.len() + 16);
key.extend_from_slice(TAG);
key.extend_from_slice(&version.generation.to_be_bytes());
key.extend_from_slice(&version.sequence.to_be_bytes());
key
}
#[must_use]
pub fn piece_key(version: Version, index: u32) -> Vec<u8> {
let mut key = key(version);
key.extend_from_slice(&index.to_be_bytes());
key
}
fn key_parts(key: &[u8]) -> Result<(Version, Option<u32>), Failure> {
let corrupt = || Failure::new("stage_row_corrupt", "stage row key is not a version");
let rest = key.strip_prefix(TAG).ok_or_else(corrupt)?;
let (generation, rest) = rest.split_at_checked(8).ok_or_else(corrupt)?;
let (sequence, piece) = rest.split_at_checked(8).ok_or_else(corrupt)?;
let version = Version {
generation: u64::from_be_bytes(generation.try_into().map_err(|_| corrupt())?),
sequence: u64::from_be_bytes(sequence.try_into().map_err(|_| corrupt())?),
};
match piece.len() {
0 => Ok((version, None)),
4 => Ok((
version,
Some(u32::from_be_bytes(piece.try_into().map_err(|_| corrupt())?)),
)),
_ => Err(corrupt()),
}
}
fn named(value: &str, path: String) -> Result<(), Failure> {
identifier(value).map_err(|mut failure| {
failure.message = format!("{path}: {}", failure.message);
failure.path = Some(path);
failure
})
}
pub(super) fn validate(stage: &Stage, change: &Change) -> Result<(), Failure> {
if matches!(
change,
Change::Fork(_) | Change::Discard(_) | Change::Merge(_)
) {
return Err(Failure {
code: "invalid_operation".to_owned(),
message: "a stage is written by a data change: branch.fork, branch.discard and \
branch.merge take none"
.to_owned(),
path: Some("/stage".to_owned()),
cause: None,
});
}
named(&stage.line, "/stage/line".to_owned())?;
named(&stage.run, "/stage/run".to_owned())?;
named(&stage.stage, "/stage/stage".to_owned())?;
named(&stage.subject, "/stage/subject".to_owned())?;
for (index, input) in stage.inputs.iter().enumerate() {
named(&input.branch, format!("/stage/inputs/{index}/branch"))?;
}
Ok(())
}
fn encode(row: &Row) -> Result<Vec<u8>, Failure> {
serde_json::to_vec(row).map_err(Failure::storage)
}
pub(super) fn stage_write(next: Version, stage: &Stage) -> Result<PreparedWrite, Failure> {
let row = Row {
format: FORMAT,
version: next,
stage: Some(stage.clone()),
merged: None,
pieces: None,
};
Ok((key(next), Some(encode(&row)?)))
}
fn format_known(format: u32) -> Result<(), Failure> {
if format != FORMAT {
return Err(Failure::new(
"stage_row_format_unsupported",
format!("stage row format {format} is not {FORMAT}"),
));
}
Ok(())
}
fn decode(version: Version, bytes: &[u8]) -> Result<Row, Failure> {
let row: Row = serde_json::from_slice(bytes)
.map_err(|error| Failure::new("stage_row_corrupt", error.to_string()))?;
format_known(row.format)?;
let misplaced = row.pieces.is_some()
&& !row
.merged
.as_ref()
.is_some_and(|merged| merged.stages.is_empty());
if row.version != version || row.stage.is_some() == row.merged.is_some() || misplaced {
return Err(Failure::new(
"stage_row_corrupt",
"a stage row names its own version and holds a stage or a merge, not both, and \
only a merge whose list is in pieces names them",
));
}
Ok(row)
}
fn decode_piece(version: Version, index: u32, bytes: &[u8]) -> Result<Vec<Listed>, Failure> {
let piece: Piece = serde_json::from_slice(bytes)
.map_err(|error| Failure::new("stage_row_corrupt", error.to_string()))?;
format_known(piece.format)?;
if piece.version != version || piece.index != index {
return Err(Failure::new(
"stage_row_corrupt",
"a piece of a merged list names its own version and index",
));
}
Ok(piece.stages)
}
fn missing_piece(version: Version, index: u32) -> Failure {
Failure::new(
"stage_row_corrupt",
format!(
"the merged list at {}.{} has no piece {index}",
version.generation, version.sequence
),
)
}
fn merged_rows(next: Version, merged: Merged) -> Result<Vec<PreparedWrite>, Failure> {
let mut whole = Row {
format: FORMAT,
version: next,
stage: None,
merged: Some(merged),
pieces: None,
};
let bytes = encode(&whole)?;
if bytes.len() <= ROW_BYTES {
return Ok(vec![(key(next), Some(bytes))]);
}
let stages = whole
.merged
.as_mut()
.map(|merged| std::mem::take(&mut merged.stages))
.unwrap_or_default();
let envelope = encode_piece(next, u32::MAX, Vec::new())?.len();
let mut pieces: Vec<Vec<Listed>> = Vec::new();
let mut listed: Vec<Listed> = Vec::new();
let mut size = 0;
for stage in stages {
let one = serde_json::to_vec(&stage).map_err(Failure::storage)?.len();
let grown = size + one + usize::from(!listed.is_empty());
if !listed.is_empty() && envelope + grown > ROW_BYTES {
pieces.push(std::mem::take(&mut listed));
size = one;
} else {
size = grown;
}
listed.push(stage);
}
if !listed.is_empty() {
pieces.push(listed);
}
whole.pieces = Some(u32::try_from(pieces.len()).map_err(Failure::storage)?);
let mut writes = vec![(key(next), Some(encode(&whole)?))];
for (index, stages) in (0_u32..).zip(pieces) {
writes.push((
piece_key(next, index),
Some(encode_piece(next, index, stages)?),
));
}
Ok(writes)
}
fn encode_piece(version: Version, piece: u32, stages: Vec<Listed>) -> Result<Vec<u8>, Failure> {
serde_json::to_vec(&Piece {
format: FORMAT,
version,
index: piece,
stages,
})
.map_err(Failure::storage)
}
#[derive(Debug)]
enum Entry {
Listed(Vec<Listed>),
Pieces(u32),
}
#[derive(Debug, Default)]
pub(super) struct Found {
rows: Vec<(Version, Entry)>,
pieces: BTreeMap<(u64, u64, u32), Vec<Listed>>,
}
impl Found {
pub(super) fn add(&mut self, source: &str, key: &[u8], bytes: &[u8]) -> Result<(), Failure> {
let (version, piece) = key_parts(key)?;
if let Some(index) = piece {
let stages = decode_piece(version, index, bytes)?;
self.pieces
.insert((version.generation, version.sequence, index), stages);
return Ok(());
}
let row = decode(version, bytes)?;
let entry = match (row.stage, row.merged, row.pieces) {
(Some(stage), _, _) => Entry::Listed(vec![Listed {
branch: source.to_owned(),
version: row.version,
stage,
}]),
(None, Some(_merged), Some(pieces)) => Entry::Pieces(pieces),
(None, Some(merged), None) => Entry::Listed(merged.stages),
(None, None, _) => Entry::Listed(Vec::new()),
};
self.rows.push((row.version, entry));
Ok(())
}
const fn is_empty(&self) -> bool {
self.rows.is_empty()
}
fn write(
mut self,
next: Version,
from: &str,
version: Version,
) -> Result<Vec<PreparedWrite>, Failure> {
self.rows
.sort_unstable_by_key(|(at, _)| (at.generation, at.sequence));
let mut stages = Vec::new();
for (at, entry) in self.rows {
match entry {
Entry::Listed(listed) => stages.extend(listed),
Entry::Pieces(count) => {
for index in 0..count {
let piece = self
.pieces
.remove(&(at.generation, at.sequence, index))
.ok_or_else(|| missing_piece(at, index))?;
stages.extend(piece);
}
}
}
}
let merged = Merged {
from: At {
branch: from.to_owned(),
version,
},
stages,
};
merged_rows(next, merged)
}
}
pub(super) fn merged_write(
reader: &ReadOnlyDatabase,
source: &Captured,
found: Found,
next: Version,
) -> Result<Vec<PreparedWrite>, Failure> {
if found.is_empty() && !is_line(reader, source)? {
return Ok(Vec::new());
}
found.write(next, &source.branch, source.version)
}
pub fn read(reader: &ReadOnlyDatabase, at: &Captured) -> Result<Lineage, Failure> {
let key = key(at.version);
let Some(bytes) = reads::get_bytes(reader, at, &key)? else {
return Ok(Lineage::default());
};
let row = decode(at.version, &bytes)?;
let mut merged = row.merged;
if let (Some(merged), Some(count)) = (merged.as_mut(), row.pieces) {
for index in 0..count {
let piece = reads::get_bytes(reader, at, &piece_key(at.version, index))?
.ok_or_else(|| missing_piece(at.version, index))?;
merged
.stages
.extend(decode_piece(at.version, index, &piece)?);
}
}
Ok(Lineage {
stage: row.stage,
merged,
})
}
pub(super) fn is_line(reader: &ReadOnlyDatabase, at: &Captured) -> Result<bool, Failure> {
Ok(branches::origin(reader, at)?.is_some_and(|(_origin, line)| line))
}
fn unfinished_key(branch: &str, generation: u64) -> Vec<u8> {
let mut key = Vec::with_capacity(UNFINISHED.len() + 32 + 8);
key.extend_from_slice(UNFINISHED);
key.extend_from_slice(Hash::of(branch.as_bytes()).as_bytes());
key.extend_from_slice(&generation.to_be_bytes());
key
}
fn unfinished_end() -> Vec<u8> {
let mut end = UNFINISHED.to_vec();
if let Some(last) = end.last_mut() {
*last = last.saturating_add(1);
}
end
}
fn unfinished_page(
reader: &ReadOnlyDatabase,
service: &BranchRefRecord,
) -> Result<Vec<(Vec<u8>, Unfinished)>, Failure> {
let limit = NonZeroUsize::new(UNFINISHED_PAGE).unwrap_or(NonZeroUsize::MIN);
let end = unfinished_end();
let mut found = Vec::new();
for shard in &service.shards {
let read = reader
.range_page_at(shard.shard_id, shard.head, UNFINISHED, &end, limit)
.map_err(Failure::observer)?;
#[cfg(test)]
unfinished_reads::count(read.entries.len());
for (key, bytes) in read.entries {
let row: Unfinished = serde_json::from_slice(&bytes)
.map_err(|error| Failure::new("line_release_row_corrupt", error.to_string()))?;
if unfinished_key(&row.branch, row.generation) != key {
return Err(Failure::new(
"line_release_row_corrupt",
"an unfinished line release names another key",
));
}
found.push((key, row));
}
}
found.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
Ok(found)
}
impl Storage {
pub(super) fn unfinished_write(
&self,
branch: &str,
generation: u64,
origin: Option<&(ForkProvenance, bool)>,
) -> Result<Option<UnfinishedRow>, Failure> {
let Some((origin, line)) = origin else {
return Ok(None);
};
if !line && !self.origin_is_discarded_line(origin)? {
return Ok(None);
}
let row = Unfinished {
branch: branch.to_owned(),
generation,
};
let bytes = serde_json::to_vec(&row).map_err(Failure::storage)?;
Ok(Some((unfinished_key(branch, generation), bytes)))
}
fn origin_is_discarded_line(&self, origin: &ForkProvenance) -> Result<bool, Failure> {
Ok(self.refs.get(&branch_name(&origin.branch)).is_none()
&& matches!(
self.discarded_origin(&origin.branch, origin.version.generation)?,
Some((_, true))
))
}
pub(super) fn release_after_discard(
&mut self,
branch: &str,
generation: u64,
) -> Result<(), Failure> {
self.finish_release(branch, generation).map_err(|failure| {
Failure::new(
"line_release_unfinished",
format!(
"the discard of {branch} is durable; releasing its line versions stopped: \
{failure}"
),
)
})
}
fn finish_release(&mut self, branch: &str, generation: u64) -> Result<(), Failure> {
let unfinished = unfinished_key(branch, generation);
self.release_lines(branch, generation, &unfinished)?;
let service = self
.refs
.get(SERVICE)
.ok_or_else(|| Failure::new("reference_missing", "missing service namespace"))?;
let shard_id = self.reader.shard_for(&unfinished);
let root = service
.shards
.iter()
.find(|shard| shard.shard_id == shard_id)
.ok_or_else(|| Failure::new("incomplete_roots", "service row shard is missing"))?
.head;
let left = self
.reader
.get_at(shard_id, root, &unfinished)
.map_err(Failure::observer)?;
if left.is_some() {
self.commit_service(&[(unfinished, None)])?;
}
Ok(())
}
pub(super) fn finish_line_releases(&mut self) -> Result<(), Failure> {
let mut finished = false;
loop {
let service = self
.refs
.get(SERVICE)
.ok_or_else(|| Failure::new("reference_missing", "missing service namespace"))?;
let page = unfinished_page(&self.reader, service)?;
if page.is_empty() {
break;
}
for (_key, row) in page {
self.release_after_discard(&row.branch, row.generation)?;
finished = true;
}
}
if finished {
self.writer.request_reclaim().map_err(|error| {
Failure::new(
"line_release_unfinished",
format!("the line releases are finished; asking a reclaim pass: {error}"),
)
})?;
}
Ok(())
}
fn release_lines(
&mut self,
branch: &str,
generation: u64,
unfinished: &[u8],
) -> Result<(), Failure> {
let Some((origin, line)) = self.discarded_origin(branch, generation)? else {
return Ok(());
};
let parent = origin.version.generation;
if self.origin_is_discarded_line(&origin)? {
self.release_line(&origin.branch, parent, (!line).then_some(unfinished))?;
}
if line {
self.release_line(branch, generation, Some(unfinished))?;
}
Ok(())
}
fn discarded_origin(
&self,
branch: &str,
generation: u64,
) -> Result<Option<(ForkProvenance, bool)>, Failure> {
let service = self
.refs
.get(SERVICE)
.ok_or_else(|| Failure::new("reference_missing", "missing service namespace"))?;
let Some(marker) = dead::marker(&self.reader, service, branch, generation)? else {
return Ok(None);
};
let Some(last) = history::capture_kept(&self.reader, service, branch, marker.last)? else {
return Ok(None);
};
branches::origin(&self.reader, &last)
}
fn release_line(
&mut self,
branch: &str,
generation: u64,
unfinished: Option<&[u8]>,
) -> Result<(), Failure> {
if self.refs.get(&branch_name(branch)).is_some() {
return Ok(());
}
let max_rows = NonZeroUsize::new(RELEASE_ROWS).unwrap_or(NonZeroUsize::MIN);
loop {
let service = self
.refs
.get(SERVICE)
.ok_or_else(|| Failure::new("reference_missing", "missing service namespace"))?;
if dead::marker(&self.reader, service, branch, generation)?.is_none() {
return Ok(());
}
let rules = Rules {
keep: None,
max_rows,
now: current_timestamp(),
};
let answer: Answer = self.release_generation(branch, generation, &rules, unfinished)?;
if answer.released.is_none() {
return Ok(());
}
}
}
}
#[cfg(test)]
pub mod unfinished_reads {
use std::cell::Cell;
thread_local! {
static READS: Cell<(usize, usize)> = const { Cell::new((0, 0)) };
}
pub(super) fn count(rows: usize) {
READS.with(|reads| {
let (made, found) = reads.get();
reads.set((made + 1, found + rows));
});
}
pub fn take() -> (usize, usize) {
READS.with(|reads| reads.replace((0, 0)))
}
}
#[cfg(test)]
pub mod before_commit {
use std::cell::Cell;
use super::Failure;
thread_local! {
static ARMED: Cell<bool> = const { Cell::new(false) };
}
pub fn arm() {
ARMED.with(|armed| armed.set(true));
}
pub fn reached() -> Result<(), Failure> {
if ARMED.with(|armed| armed.replace(false)) {
return Err(Failure::new(
"injected_crash",
"stopped after the writes were staged and before the commit",
));
}
Ok(())
}
}
#[cfg(test)]
#[path = "lineage_pieces_tests.rs"]
mod pieces_tests;