use std::collections::HashSet;
use std::fs::{File, OpenOptions};
use std::io::{Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};
use std::sync::Mutex;
use anyhow::Context;
use gunnar_core::txlog::{Crc32c, Entry};
use gunnar_grpc::{ControlError, ControlResult};
use super::Applied;
use crate::stores::ServerStores;
pub const PARTIAL_DIR: &str = "txlog-restore-partial";
const DATA: &str = "entry.blob";
const RECORD: &str = "entry.rec";
const RECORD_MAGIC: &str = "gunnar-restore-partial v1";
pub const COMMIT_EVERY: u64 = 8 << 20;
pub const IN_MEMORY_UP_TO: u64 = 4 << 20;
pub const TWIN_TOO_OLD: &str = "restore-twin-too-old";
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Staged {
pub seq: u64,
pub kind: u16,
pub blob_len: u64,
pub meta: Vec<u8>,
pub durable: u64,
}
impl Staged {
fn to_text(&self) -> String {
format!(
"{RECORD_MAGIC}\nseq {}\nkind {}\nblob_len {}\ndurable {}\nmeta {}\n",
self.seq,
self.kind,
self.blob_len,
self.durable,
hex::encode(&self.meta)
)
}
fn from_text(t: &str) -> Option<Staged> {
let mut lines = t.lines();
if lines.next()? != RECORD_MAGIC {
return None;
}
let mut field = |name: &str| lines.next()?.strip_prefix(name)?.strip_prefix(' ').map(str::to_owned);
let seq = field("seq")?.parse().ok()?;
let kind = field("kind")?.parse().ok()?;
let blob_len = field("blob_len")?.parse().ok()?;
let durable: u64 = field("durable")?.parse().ok()?;
let meta = hex::decode(field("meta")?).ok()?;
(durable <= blob_len && seq > 0).then_some(Staged { seq, kind, blob_len, meta, durable })
}
}
pub fn staged(root: &Path) -> Option<Staged> {
std::fs::read_to_string(root.join(PARTIAL_DIR).join(RECORD)).ok().as_deref().and_then(Staged::from_text)
}
pub fn partial(root: &Path) -> Option<(u64, u64)> {
staged(root).map(|s| (s.seq, s.durable))
}
pub fn forget(root: &Path) -> anyhow::Result<()> {
let dir = root.join(PARTIAL_DIR);
match std::fs::remove_dir_all(&dir) {
Ok(()) => {}
Err(e) if e.kind() == std::io::ErrorKind::NotFound => return Ok(()),
Err(e) => return Err(e).with_context(|| format!("removing {}", dir.display())),
}
let r = File::open(root).with_context(|| format!("opening {} to flush it", root.display()))?;
r.sync_all().with_context(|| format!("flushing the removal of {PARTIAL_DIR} from {}", root.display()))?;
Ok(())
}
fn write_record(root: &Path, rec: &Staged) -> anyhow::Result<()> {
let dir = root.join(PARTIAL_DIR);
if !dir.is_dir() {
std::fs::create_dir_all(&dir).with_context(|| format!("creating {}", dir.display()))?;
crate::durable::sync_dir(root).with_context(|| format!("flushing {PARTIAL_DIR} into {}", root.display()))?;
}
crate::durable::write_atomic_unframed(&dir, RECORD, rec.to_text().as_bytes())
}
fn sanitize(root: &Path, next: u64) -> anyhow::Result<Option<(Staged, File)>> {
if !root.join(PARTIAL_DIR).exists() {
return Ok(None);
}
let Some(rec) = staged(root) else {
forget(root)?;
return Ok(None);
};
if rec.seq != next {
tracing::info!(
target: gunnar_core::observe::TARGET_TWIN,
staged = rec.seq, next, "a staged restore entry that is not the next one is dropped (applied already, or another history)"
);
forget(root)?;
return Ok(None);
}
let path = root.join(PARTIAL_DIR).join(DATA);
let file = match OpenOptions::new().read(true).write(true).open(&path) {
Ok(f) => f,
Err(_) => {
forget(root)?;
return Ok(None);
}
};
let len = file.metadata().with_context(|| format!("reading {}", path.display()))?.len();
if len < rec.durable {
tracing::warn!(
target: gunnar_core::observe::TARGET_TWIN,
seq = rec.seq, durable = rec.durable, on_disk = len,
"the staged restore entry is SHORTER than its record says was made durable; it is not trusted and starts again at byte 0"
);
drop(file);
forget(root)?;
return Ok(None);
}
let mut file = file;
if len > rec.durable {
file.set_len(rec.durable).with_context(|| format!("cutting {} back to its durable {} bytes", path.display(), rec.durable))?;
file.sync_all().with_context(|| format!("flushing {}", path.display()))?;
}
file.seek(SeekFrom::Start(rec.durable)).with_context(|| format!("seeking {}", path.display()))?;
Ok(Some((rec, file)))
}
static OPEN: Mutex<Option<HashSet<PathBuf>>> = Mutex::new(None);
struct Held(PathBuf);
impl Held {
fn take(root: &Path) -> Option<Held> {
let mut open = OPEN.lock().unwrap_or_else(|p| p.into_inner());
open.get_or_insert_with(HashSet::new).insert(root.to_path_buf()).then(|| Held(root.to_path_buf()))
}
}
impl Drop for Held {
fn drop(&mut self) {
let mut open = OPEN.lock().unwrap_or_else(|p| p.into_inner());
if let Some(set) = open.as_mut() {
set.remove(&self.0);
}
}
}
enum Body {
Skip,
Memory(Vec<u8>),
Spool { file: File, rec: Staged, written: u64 },
}
struct Open {
seq: u64,
kind: u16,
meta: Vec<u8>,
blob_len: u64,
pos: u64,
explicit: bool,
held_before: u64,
body: Body,
}
pub struct Intake<'a> {
stores: &'a ServerStores,
_held: Held,
next: u64,
kept: Option<(Staged, File)>,
open: Option<Open>,
applied: u64,
lost: u64,
snap: Option<super::snapshot::RestoreSnapshot>,
snap_touched: bool,
}
impl<'a> Intake<'a> {
pub fn open(stores: &'a ServerStores, owner: &str) -> anyhow::Result<Option<Intake<'a>>> {
let Some(held) = Held::take(stores.root()) else { return Ok(None) };
let next = super::restore_begin_by(stores, owner)?;
let kept = sanitize(stores.root(), next)?;
if let Some((rec, _)) = &kept {
tracing::info!(
target: gunnar_core::observe::TARGET_TWIN,
seq = rec.seq, held = rec.durable, blob_len = rec.blob_len,
"a restore stream resumes an entry this box holds part of: {} of {} bytes of seq {}", rec.durable, rec.blob_len, rec.seq
);
}
let mut intake = Intake { stores, _held: held, next, kept, open: None, applied: 0, lost: 0, snap: None, snap_touched: false };
if super::is_restoring(stores.root()) {
intake.snap = super::snapshot::RestoreSnapshot::resume(stores.root())?;
}
if let Some(s) = &intake.snap {
let (items, bytes) = s.position();
tracing::info!(
target: gunnar_core::observe::TARGET_TWIN,
id = s.id(), items, bytes, whole = s.whole().is_some(),
"a restore stream resumes the twin's snapshot this box holds part of: {items} item(s) and {bytes} byte(s) of the next"
);
if s.whole().is_some() {
intake.install_snapshot().map_err(|e| anyhow::anyhow!("{e}"))?;
}
}
Ok(Some(intake))
}
fn install_snapshot(&mut self) -> ControlResult<()> {
let Some(s) = self.snap.take() else { return Ok(()) };
let id = s.id().to_owned();
let staged = s.finish().map_err(|err| ControlError::Internal(format!("the twin's snapshot {id}: {err:#}")))?;
let at = staged.at;
super::snapshot::install(self.stores, &staged).map_err(|err| ControlError::Internal(format!("installing the twin's snapshot {id}: {err:#}")))?;
drop(staged);
super::snapshot::forget_restore_snapshot_or_warn(self.stores.root());
tracing::warn!(target: gunnar_core::observe::TARGET_TWIN, id = %id, at, "the twin's snapshot, staged across streams, is installed: this box holds everything up to seq {at}; its log follows from {}", at + 1);
self.next = at + 1;
Ok(())
}
fn take_snapshot_piece(&mut self, chunk: &gunnar_proto::LogChunk) -> ControlResult<()> {
let root = self.stores.root();
let id = chunk.snapshot_id.trim();
let (item, offset) = (chunk.snapshot_item.unwrap_or(0), chunk.offset.unwrap_or(0));
self.snap_touched = true;
if self.snap.as_ref().is_none_or(|s| s.id() != id) {
if (item, offset) != (0, 0) {
let held = self.snap.as_ref().map(|s| format!("part of snapshot {}", s.id())).unwrap_or_else(|| "none of it".into());
return Err(ControlError::FailedPrecondition(format!(
"a piece at item {item} byte {offset} of the twin's snapshot {id} arrived, and this box holds {held}; resend snapshot {id} from its first item"
)));
}
if let Some(old) = self.snap.take() {
tracing::warn!(target: gunnar_core::observe::TARGET_TWIN, old = old.id(), new = id, "the twin's snapshot is another one now; what this box staged of the old one is dropped and the new one is taken from its first item");
}
self.snap = Some(super::snapshot::RestoreSnapshot::start(root, id).map_err(|err| ControlError::Internal(format!("staging the twin's snapshot {id}: {err:#}")))?);
}
let s = self.snap.as_mut().expect("a snapshot is staged");
let (done, held) = s.position();
if item < done {
return Ok(());
}
if item > done || offset > held {
return Err(ControlError::FailedPrecondition(format!(
"a piece at item {item} byte {offset} of the twin's snapshot {id} arrived, and this box holds {done} item(s) and {held} byte(s) of the next; resend from item {done} byte {held}"
)));
}
let skip = (held - offset).min(chunk.blob.len() as u64) as usize;
if skip == chunk.blob.len() && !chunk.last {
return Ok(());
}
let kind = u16::try_from(chunk.kind).map_err(|_| ControlError::InvalidArgument(format!("snapshot piece of kind {}", chunk.kind)))?;
let piece = super::snapshot::Piece { kind, meta: &chunk.meta, blob_len: chunk.blob_len, blob: &chunk.blob[skip..], last: chunk.last };
if let Err(err) = s.feed(piece, chunk.entry_crc) {
if let Some(bad) = err.downcast_ref::<super::snapshot::ItemCrcMismatch>() {
let said = bad.to_string();
self.snap = None;
super::snapshot::forget_restore_snapshot_or_warn(root);
return Err(ControlError::FailedPrecondition(format!(
"{said}: nothing of the twin's snapshot {id} is installed — this box drops what it staged of it and takes it again from its first item"
)));
}
return Err(ControlError::InvalidArgument(format!("the twin's snapshot {id}, item {item}: {err:#}")));
}
if s.whole().is_some() {
self.install_snapshot()?;
}
Ok(())
}
pub fn next(&self) -> u64 {
self.next
}
pub fn take(&mut self, chunk: &gunnar_proto::LogChunk) -> ControlResult<()> {
if chunk.seq == 0 {
if let Some(o) = &self.open {
return Err(ControlError::InvalidArgument(format!("a snapshot piece arrived with seq {} unfinished", o.seq)));
}
if !chunk.snapshot_id.trim().is_empty() {
return self.take_snapshot_piece(chunk);
}
let kind = u16::try_from(chunk.kind).map_err(|_| ControlError::InvalidArgument(format!("snapshot piece of kind {}", chunk.kind)))?;
let piece = super::snapshot::Piece { kind, meta: &chunk.meta, blob_len: chunk.blob_len, blob: &chunk.blob, last: chunk.last };
if let Some(at) = super::snapshot::restore_piece(self.stores, piece).map_err(|err| ControlError::Internal(format!("installing the snapshot: {err:#}")))? {
self.next = at + 1;
}
return Ok(());
}
let kind = u16::try_from(chunk.kind).map_err(|_| ControlError::InvalidArgument(format!("seq {}: kind {} is not a frame kind", chunk.seq, chunk.kind)))?;
match &self.open {
Some(o) if o.seq == chunk.seq => {
if o.kind != kind || o.blob_len != chunk.blob_len {
return Err(ControlError::InvalidArgument(format!(
"seq {}: a later piece says kind {kind} of {} bytes, the first said kind {} of {}",
o.seq, chunk.blob_len, o.kind, o.blob_len
)));
}
if let Some(at) = chunk.offset {
if at != o.pos {
return Err(ControlError::InvalidArgument(format!("seq {}: a piece at byte {at} arrived where byte {} was due", o.seq, o.pos)));
}
}
}
Some(o) => return Err(ControlError::InvalidArgument(format!("seq {} arrived with seq {} unfinished", chunk.seq, o.seq))),
None => self.open_entry(chunk, kind)?,
}
let stores = self.stores;
let root = stores.root();
let o = self.open.as_mut().expect("an entry is open");
let len = chunk.blob.len() as u64;
if o.pos + len > o.blob_len {
return Err(ControlError::InvalidArgument(format!("seq {}: more than the {} blob bytes it declared arrived", o.seq, o.blob_len)));
}
match &mut o.body {
Body::Skip => {}
Body::Memory(v) => v.extend_from_slice(&chunk.blob),
Body::Spool { file, rec, written } => {
if o.pos + len > *written {
let from = (*written - o.pos) as usize;
file.write_all(&chunk.blob[from..]).map_err(|e| internal(o.seq, "writing the staged entry", e))?;
*written = o.pos + len;
if *written - rec.durable >= COMMIT_EVERY {
commit(root, file, rec, *written).map_err(|e| internal(o.seq, "committing the staged entry", e))?;
}
}
}
}
o.pos += len;
if chunk.last {
self.finish_entry(chunk.entry_crc)?;
}
Ok(())
}
fn open_entry(&mut self, chunk: &gunnar_proto::LogChunk, kind: u16) -> ControlResult<()> {
let seq = chunk.seq;
let start = chunk.offset.unwrap_or(0);
let explicit = chunk.offset.is_some();
let mut open = Open { seq, kind, meta: chunk.meta.clone(), blob_len: chunk.blob_len, pos: start, explicit, held_before: 0, body: Body::Skip };
if start > chunk.blob_len {
return Err(ControlError::InvalidArgument(format!("seq {seq}: a piece at byte {start} of a {}-byte blob", chunk.blob_len)));
}
if seq < self.next {
self.open = Some(open);
return Ok(());
}
if seq > self.next {
return Err(ControlError::FailedPrecondition(format!("seq {seq} arrived where {} was due; resend from {}", self.next, self.next)));
}
let stores = self.stores;
let root = stores.root();
let held = match self.kept.take() {
Some((rec, file)) if rec.seq == seq && rec.kind == kind && rec.blob_len == chunk.blob_len && (chunk.meta.is_empty() || chunk.meta == rec.meta) => Some((rec, file)),
Some((rec, _)) => {
tracing::warn!(
target: gunnar_core::observe::TARGET_TWIN,
seq, staged_kind = rec.kind, staged_len = rec.blob_len, kind, blob_len = chunk.blob_len,
"the staged restore entry is not the entry now arriving at its number; it is dropped and the entry starts at byte 0"
);
forget(root).map_err(|e| internal(seq, "dropping the staged entry", e))?;
None
}
None => None,
};
let held_bytes = held.as_ref().map(|(rec, _)| rec.durable).unwrap_or(0);
if start > held_bytes {
self.kept = held;
return Err(ControlError::FailedPrecondition(format!(
"seq {seq}: this stream starts it at byte {start}, and this box holds {held_bytes} byte(s) of it; resend seq {seq} from byte {held_bytes}"
)));
}
open.held_before = held_bytes;
open.body = match held {
Some((rec, file)) => {
if open.meta.is_empty() {
open.meta.clone_from(&rec.meta);
}
Body::Spool { file, rec, written: held_bytes }
}
None if start == 0 && (chunk.blob_len <= IN_MEMORY_UP_TO || (chunk.last && chunk.blob.len() as u64 == chunk.blob_len)) => {
Body::Memory(Vec::with_capacity(gunnar_proto::reserve_for(chunk.blob_len)))
}
None => {
let rec = Staged { seq, kind, blob_len: chunk.blob_len, meta: open.meta.clone(), durable: 0 };
let file = start_spool(root, &rec).map_err(|e| internal(seq, "staging the entry", e))?;
Body::Spool { file, rec, written: 0 }
}
};
self.open = Some(open);
Ok(())
}
fn finish_entry(&mut self, entry_crc: Option<u32>) -> ControlResult<()> {
let stores = self.stores;
let root = stores.root();
let o = self.open.take().expect("an entry is open");
if o.pos != o.blob_len {
return Err(ControlError::InvalidArgument(format!("seq {}: {} of {} blob bytes arrived", o.seq, o.pos, o.blob_len)));
}
let (blob, staged) = match o.body {
Body::Skip => return Ok(()),
Body::Memory(v) => (v, false),
Body::Spool { mut file, mut rec, written } => {
if written != o.blob_len {
return Err(ControlError::InvalidArgument(format!("seq {}: {written} of {} blob bytes are staged", o.seq, o.blob_len)));
}
commit(root, &mut file, &mut rec, written).map_err(|e| internal(o.seq, "committing the staged entry", e))?;
let mut blob = Vec::with_capacity(gunnar_proto::reserve_for(o.blob_len));
file.seek(SeekFrom::Start(0)).and_then(|_| file.read_to_end(&mut blob)).map_err(|e| internal(o.seq, "reading the staged entry back", e))?;
if blob.len() as u64 != o.blob_len {
forget(root).map_err(|e| internal(o.seq, "dropping the staged entry", e))?;
return Err(ControlError::FailedPrecondition(format!(
"seq {}: {} of {} staged bytes read back; they are dropped and seq {} starts again from byte 0",
o.seq,
blob.len(),
o.blob_len,
o.seq
)));
}
(blob, true)
}
};
if let Some(want) = entry_crc {
let mut crc = Crc32c::new();
crc.update(&o.meta);
crc.update(&blob);
let got = crc.finish();
if got != want {
if staged {
forget(root).map_err(|e| internal(o.seq, "dropping the staged entry", e))?;
}
return Err(ControlError::FailedPrecondition(format!(
"seq {}: the {} bytes this box put together do not match the twin's crc ({got:08x}, the twin's {want:08x}); \
nothing was applied, the staged bytes are dropped and seq {} starts again from byte 0",
o.seq, o.blob_len, o.seq
)));
}
}
let e = Entry { seq: o.seq, kind: o.kind, meta: o.meta, blob };
match super::restore_apply(self.stores, &e) {
Ok(Applied::Lost(_)) => self.lost += 1,
Ok(_) => {}
Err(err) => return Err(ControlError::Internal(format!("restoring seq {}: {err:#}", e.seq))),
}
if staged {
tracing::info!(target: gunnar_core::observe::TARGET_TWIN, seq = e.seq, blob_len = e.blob.len(), "a restore entry staged across streams is whole, checked and applied");
}
if staged || root.join(PARTIAL_DIR).exists() {
forget(root).map_err(|err| internal(e.seq, "removing the staged entry after its apply", err))?;
}
self.applied += 1;
self.next += 1;
Ok(())
}
pub fn end_of_stream(&mut self) -> ControlResult<()> {
let stores = self.stores;
let root = stores.root();
if self.snap_touched {
if let Some(s) = self.snap.as_mut() {
s.commit().map_err(|err| ControlError::Internal(format!("committing the twin's snapshot {}: {err:#}", s.id())))?;
let (items, bytes) = s.position();
return Err(ControlError::FailedPrecondition(format!(
"the stream ended inside the twin's snapshot {}: this box holds {items} item(s) of it and {bytes} byte(s) of the next, \
durably; the next stream resumes there (Bootstrap.Status restore_snapshot_items/bytes, Mirror.SnapshotFetch \
from_item/from_offset) — nothing of it is sent twice",
s.id()
)));
}
}
let Some(mut o) = self.open.take() else { return Ok(()) };
match &mut o.body {
Body::Skip => Err(ControlError::InvalidArgument(format!(
"the stream ended inside seq {}, which this box holds already; resend from the box's head + 1 ({})",
o.seq, self.next
))),
Body::Memory(_) => Err(ControlError::InvalidArgument(format!(
"the stream ended inside seq {} ({} of {} bytes, small enough to be gathered in memory, and dropped); \
resend from the box's head + 1 ({})",
o.seq, o.pos, o.blob_len, self.next
))),
Body::Spool { file, rec, written } => {
commit(root, file, rec, *written).map_err(|e| internal(o.seq, "committing the staged entry", e))?;
if !o.explicit && o.held_before > 0 && *written <= o.held_before {
return Err(ControlError::FailedPrecondition(format!(
"REFUSED [{TWIN_TOO_OLD}] the twin is too old for a resumed restore; upgrade the twin. Seq {seq} is \
{len} bytes and needs more than one call: this box holds {held} of them, durably, but this stream \
sent the entry from its first byte (no `offset` on its chunks — a twin, or a restore lane, built \
before resumed restores) and ended before reaching byte {held}, so skipping what is held gains \
nothing and every resend ends the same way. With the twin upgraded the next stream resumes seq \
{seq} at byte {held} (Mirror.LogFetch from_offset); nothing here is lost meanwhile",
seq = o.seq,
len = o.blob_len,
held = o.held_before,
)));
}
Err(ControlError::FailedPrecondition(format!(
"the stream ended inside seq {seq}: this box holds {w} of its {len} bytes, durably; the next stream \
resumes seq {seq} at byte {w} (Bootstrap.Status restore_partial_bytes, Mirror.LogFetch from_offset) — \
nothing of it is sent twice",
seq = o.seq,
w = written,
len = o.blob_len,
)))
}
}
}
pub fn counts(&self) -> (u64, u64) {
(self.applied, self.lost)
}
}
impl Drop for Intake<'_> {
fn drop(&mut self) {
let root = self.stores.root().to_path_buf();
if let Some(Open { seq, body: Body::Spool { file, rec, written }, .. }) = self.open.as_mut() {
if *written > rec.durable {
if let Err(err) = commit(&root, file, rec, *written) {
tracing::warn!(target: gunnar_core::observe::TARGET_TWIN, seq = *seq, %err, "the staged restore entry could not be committed as its stream ended; the next stream resumes from its last commit");
}
}
}
}
}
impl gunnar_grpc::RestoreIntake for Intake<'_> {
fn take(&mut self, chunk: &gunnar_proto::LogChunk) -> ControlResult<()> {
Intake::take(self, chunk)
}
fn end_of_stream(&mut self) -> ControlResult<()> {
Intake::end_of_stream(self)
}
fn counts(&self) -> (u64, u64) {
Intake::counts(self)
}
}
fn commit(root: &Path, file: &mut File, rec: &mut Staged, written: u64) -> anyhow::Result<()> {
file.sync_data().context("fdatasync of the staged entry")?;
if rec.durable != written {
rec.durable = written;
write_record(root, rec)?;
}
Ok(())
}
fn start_spool(root: &Path, rec: &Staged) -> anyhow::Result<File> {
forget(root)?;
write_record(root, rec)?;
let path = root.join(PARTIAL_DIR).join(DATA);
let file = OpenOptions::new().read(true).write(true).create(true).truncate(true).open(&path).with_context(|| format!("creating {}", path.display()))?;
crate::durable::sync_dir(&root.join(PARTIAL_DIR)).with_context(|| format!("flushing {} into place", path.display()))?;
Ok(file)
}
fn internal(seq: u64, what: &str, e: impl Into<anyhow::Error>) -> ControlError {
ControlError::Internal(format!("seq {seq}: {what}: {:#}", e.into()))
}