use std::{collections::BTreeSet, future::Future, io, path::PathBuf};
use bao_tree::{
io::fsm::{BaoContentItem, Outboard},
BaoTree, ChunkRanges,
};
use bytes::Bytes;
use futures_lite::{Stream, StreamExt};
use genawaiter::rc::{Co, Gen};
use iroh_base::rpc::RpcError;
use iroh_io::AsyncSliceReader;
use serde::{Deserialize, Serialize};
use tokio::io::AsyncRead;
use tokio_util::task::LocalPoolHandle;
use crate::{
hashseq::parse_hash_seq,
protocol::RangeSpec,
util::{
progress::{BoxedProgressSender, IdGenerator, ProgressSender},
Tag,
},
BlobFormat, Hash, HashAndFormat, TempTag, IROH_BLOCK_SIZE,
};
pub use bao_tree;
pub use range_collections;
pub type DbIter<T> = Box<dyn Iterator<Item = io::Result<T>> + Send + Sync + 'static>;
pub type ExportProgressCb = Box<dyn Fn(u64) -> io::Result<()> + Send + Sync + 'static>;
#[derive(Debug, Clone, Eq, PartialEq)]
pub enum EntryStatus {
Complete,
Partial,
NotFound,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, Eq, PartialEq)]
pub enum BaoBlobSize {
Unverified(u64),
Verified(u64),
}
impl BaoBlobSize {
pub fn new(size: u64, verified: bool) -> Self {
if verified {
BaoBlobSize::Verified(size)
} else {
BaoBlobSize::Unverified(size)
}
}
pub fn value(&self) -> u64 {
match self {
BaoBlobSize::Unverified(size) => *size,
BaoBlobSize::Verified(size) => *size,
}
}
}
pub trait MapEntry: std::fmt::Debug + Clone + Send + Sync + 'static {
fn hash(&self) -> Hash;
fn size(&self) -> BaoBlobSize;
fn is_complete(&self) -> bool;
fn outboard(&self) -> impl Future<Output = io::Result<impl Outboard>> + Send;
fn data_reader(&self) -> impl Future<Output = io::Result<impl AsyncSliceReader>> + Send;
}
pub trait Map: Clone + Send + Sync + 'static {
type Entry: MapEntry;
fn get(&self, hash: &Hash) -> impl Future<Output = io::Result<Option<Self::Entry>>> + Send;
}
pub trait MapEntryMut: MapEntry {
fn batch_writer(&self) -> impl Future<Output = io::Result<impl BaoBatchWriter>> + Send;
}
pub trait BaoBatchWriter {
fn write_batch(
&mut self,
size: u64,
batch: Vec<BaoContentItem>,
) -> impl Future<Output = io::Result<()>>;
fn sync(&mut self) -> impl Future<Output = io::Result<()>>;
}
impl<W: BaoBatchWriter> BaoBatchWriter for &mut W {
async fn write_batch(&mut self, size: u64, batch: Vec<BaoContentItem>) -> io::Result<()> {
(**self).write_batch(size, batch).await
}
async fn sync(&mut self) -> io::Result<()> {
(**self).sync().await
}
}
#[derive(Debug)]
pub(crate) struct FallibleProgressBatchWriter<W, F>(W, F);
impl<W: BaoBatchWriter, F: Fn(u64, usize) -> io::Result<()> + 'static>
FallibleProgressBatchWriter<W, F>
{
pub fn new(inner: W, on_write: F) -> Self {
Self(inner, on_write)
}
}
impl<W: BaoBatchWriter, F: Fn(u64, usize) -> io::Result<()> + 'static> BaoBatchWriter
for FallibleProgressBatchWriter<W, F>
{
async fn write_batch(&mut self, size: u64, batch: Vec<BaoContentItem>) -> io::Result<()> {
let chunk = batch
.iter()
.filter_map(|item| {
if let BaoContentItem::Leaf(leaf) = item {
Some((leaf.offset, leaf.data.len()))
} else {
None
}
})
.next();
self.0.write_batch(size, batch).await?;
if let Some((offset, len)) = chunk {
(self.1)(offset, len)?;
}
Ok(())
}
async fn sync(&mut self) -> io::Result<()> {
self.0.sync().await
}
}
pub trait MapMut: Map {
type EntryMut: MapEntryMut;
fn get_mut(
&self,
hash: &Hash,
) -> impl Future<Output = io::Result<Option<Self::EntryMut>>> + Send;
fn get_or_create(
&self,
hash: Hash,
size: u64,
) -> impl Future<Output = io::Result<Self::EntryMut>> + Send;
fn entry_status(&self, hash: &Hash) -> impl Future<Output = io::Result<EntryStatus>> + Send;
fn entry_status_sync(&self, hash: &Hash) -> io::Result<EntryStatus>;
fn insert_complete(&self, entry: Self::EntryMut)
-> impl Future<Output = io::Result<()>> + Send;
}
pub trait ReadableStore: Map {
fn blobs(&self) -> impl Future<Output = io::Result<DbIter<Hash>>> + Send;
fn tags(&self) -> impl Future<Output = io::Result<DbIter<(Tag, HashAndFormat)>>> + Send;
fn temp_tags(&self) -> Box<dyn Iterator<Item = HashAndFormat> + Send + Sync + 'static>;
fn consistency_check(
&self,
repair: bool,
tx: BoxedProgressSender<ConsistencyCheckProgress>,
) -> impl Future<Output = io::Result<()>> + Send;
fn partial_blobs(&self) -> impl Future<Output = io::Result<DbIter<Hash>>> + Send;
fn export(
&self,
hash: Hash,
target: PathBuf,
mode: ExportMode,
progress: ExportProgressCb,
) -> impl Future<Output = io::Result<()>> + Send;
}
pub trait Store: ReadableStore + MapMut {
fn import_file(
&self,
data: PathBuf,
mode: ImportMode,
format: BlobFormat,
progress: impl ProgressSender<Msg = ImportProgress> + IdGenerator,
) -> impl Future<Output = io::Result<(TempTag, u64)>> + Send;
fn import_bytes(
&self,
bytes: Bytes,
format: BlobFormat,
) -> impl Future<Output = io::Result<TempTag>> + Send;
fn import_stream(
&self,
data: impl Stream<Item = io::Result<Bytes>> + Send + Unpin + 'static,
format: BlobFormat,
progress: impl ProgressSender<Msg = ImportProgress> + IdGenerator,
) -> impl Future<Output = io::Result<(TempTag, u64)>> + Send;
fn import_reader(
&self,
data: impl AsyncRead + Send + Unpin + 'static,
format: BlobFormat,
progress: impl ProgressSender<Msg = ImportProgress> + IdGenerator,
) -> impl Future<Output = io::Result<(TempTag, u64)>> + Send {
let stream = tokio_util::io::ReaderStream::new(data);
self.import_stream(stream, format, progress)
}
fn set_tag(
&self,
name: Tag,
hash: Option<HashAndFormat>,
) -> impl Future<Output = io::Result<()>> + Send;
fn create_tag(&self, hash: HashAndFormat) -> impl Future<Output = io::Result<Tag>> + Send;
fn temp_tag(&self, value: HashAndFormat) -> TempTag;
fn gc_start(&self) -> impl Future<Output = io::Result<()>> + Send;
fn gc_mark(&self, live: &mut BTreeSet<Hash>) -> impl Stream<Item = GcMarkEvent> + Unpin {
Gen::new(|co| async move {
if let Err(e) = gc_mark_task(self, live, &co).await {
co.yield_(GcMarkEvent::Error(e)).await;
}
})
}
fn gc_sweep(&self, live: &BTreeSet<Hash>) -> impl Stream<Item = GcSweepEvent> + Unpin {
Gen::new(|co| async move {
if let Err(e) = gc_sweep_task(self, live, &co).await {
co.yield_(GcSweepEvent::Error(e)).await;
}
})
}
fn delete(&self, hashes: Vec<Hash>) -> impl Future<Output = io::Result<()>> + Send;
fn shutdown(&self) -> impl Future<Output = ()> + Send;
fn validate(
&self,
repair: bool,
tx: BoxedProgressSender<ValidateProgress>,
) -> impl Future<Output = io::Result<()>> + Send {
validate_impl(self, repair, tx)
}
}
async fn validate_impl(
store: &impl Store,
repair: bool,
tx: BoxedProgressSender<ValidateProgress>,
) -> io::Result<()> {
use futures_buffered::BufferedStreamExt;
let validate_parallelism: usize = num_cpus::get();
let lp = LocalPoolHandle::new(validate_parallelism);
let complete = store.blobs().await?.collect::<io::Result<Vec<_>>>()?;
let partial = store
.partial_blobs()
.await?
.collect::<io::Result<Vec<_>>>()?;
tx.send(ValidateProgress::Starting {
total: complete.len() as u64,
})
.await?;
let complete_result = futures_lite::stream::iter(complete)
.map(|hash| {
let store = store.clone();
let tx = tx.clone();
lp.spawn_pinned(move || async move {
let entry = store
.get(&hash)
.await?
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "entry not found"))?;
let size = entry.size().value();
let outboard = entry.outboard().await?;
let data = entry.data_reader().await?;
let chunk_ranges = ChunkRanges::all();
let mut ranges = bao_tree::io::fsm::valid_ranges(outboard, data, &chunk_ranges);
let id = tx.new_id();
tx.send(ValidateProgress::Entry {
id,
hash,
path: None,
size,
})
.await?;
let mut actual_chunk_ranges = ChunkRanges::empty();
while let Some(item) = ranges.next().await {
let item = item?;
let offset = item.start.to_bytes();
actual_chunk_ranges |= ChunkRanges::from(item);
tx.try_send(ValidateProgress::EntryProgress { id, offset })?;
}
let expected_chunk_range =
ChunkRanges::from(..BaoTree::new(size, IROH_BLOCK_SIZE).chunks());
let incomplete = actual_chunk_ranges == expected_chunk_range;
let error = if incomplete {
None
} else {
Some(format!(
"expected chunk ranges {:?}, got chunk ranges {:?}",
expected_chunk_range, actual_chunk_ranges
))
};
tx.send(ValidateProgress::EntryDone { id, error }).await?;
drop(ranges);
drop(entry);
io::Result::Ok((hash, incomplete))
})
})
.buffered_unordered(validate_parallelism)
.collect::<Vec<_>>()
.await;
let partial_result = futures_lite::stream::iter(partial)
.map(|hash| {
let store = store.clone();
let tx = tx.clone();
lp.spawn_pinned(move || async move {
let entry = store
.get(&hash)
.await?
.ok_or_else(|| io::Error::new(io::ErrorKind::NotFound, "entry not found"))?;
let size = entry.size().value();
let outboard = entry.outboard().await?;
let data = entry.data_reader().await?;
let chunk_ranges = ChunkRanges::all();
let mut ranges = bao_tree::io::fsm::valid_ranges(outboard, data, &chunk_ranges);
let id = tx.new_id();
tx.send(ValidateProgress::PartialEntry {
id,
hash,
path: None,
size,
})
.await?;
let mut actual_chunk_ranges = ChunkRanges::empty();
while let Some(item) = ranges.next().await {
let item = item?;
let offset = item.start.to_bytes();
actual_chunk_ranges |= ChunkRanges::from(item);
tx.try_send(ValidateProgress::PartialEntryProgress { id, offset })?;
}
tx.send(ValidateProgress::PartialEntryDone {
id,
ranges: RangeSpec::new(&actual_chunk_ranges),
})
.await?;
drop(ranges);
drop(entry);
io::Result::Ok(())
})
})
.buffered_unordered(validate_parallelism)
.collect::<Vec<_>>()
.await;
let mut to_downgrade = Vec::new();
for item in complete_result {
let (hash, incomplete) = item??;
if incomplete {
to_downgrade.push(hash);
}
}
for item in partial_result {
item??;
}
if repair {
return Err(io::Error::new(
io::ErrorKind::Other,
"repair not implemented",
));
}
Ok(())
}
async fn gc_mark_task<'a>(
store: &'a impl Store,
live: &'a mut BTreeSet<Hash>,
co: &Co<GcMarkEvent>,
) -> anyhow::Result<()> {
macro_rules! debug {
($($arg:tt)*) => {
co.yield_(GcMarkEvent::CustomDebug(format!($($arg)*))).await;
};
}
macro_rules! warn {
($($arg:tt)*) => {
co.yield_(GcMarkEvent::CustomWarning(format!($($arg)*), None)).await;
};
}
let mut roots = BTreeSet::new();
debug!("traversing tags");
for item in store.tags().await? {
let (name, haf) = item?;
debug!("adding root {:?} {:?}", name, haf);
roots.insert(haf);
}
debug!("traversing temp roots");
for haf in store.temp_tags() {
debug!("adding temp pin {:?}", haf);
roots.insert(haf);
}
for HashAndFormat { hash, format } in roots {
if live.insert(hash) && !format.is_raw() {
let Some(entry) = store.get(&hash).await? else {
warn!("gc: {} not found", hash);
continue;
};
if !entry.is_complete() {
warn!("gc: {} is partial", hash);
continue;
}
let Ok(reader) = entry.data_reader().await else {
warn!("gc: {} creating data reader failed", hash);
continue;
};
let Ok((mut stream, count)) = parse_hash_seq(reader).await else {
warn!("gc: {} parse failed", hash);
continue;
};
debug!("parsed collection {} {:?}", hash, count);
loop {
let item = match stream.next().await {
Ok(Some(item)) => item,
Ok(None) => break,
Err(_err) => {
warn!("gc: {} parse failed", hash);
break;
}
};
live.insert(item);
}
}
}
debug!("gc mark done. found {} live blobs", live.len());
Ok(())
}
async fn gc_sweep_task<'a>(
store: &'a impl Store,
live: &BTreeSet<Hash>,
co: &Co<GcSweepEvent>,
) -> anyhow::Result<()> {
let blobs = store.blobs().await?.chain(store.partial_blobs().await?);
let mut count = 0;
let mut batch = Vec::new();
for hash in blobs {
let hash = hash?;
if !live.contains(&hash) {
batch.push(hash);
count += 1;
}
if batch.len() >= 100 {
store.delete(batch.clone()).await?;
batch.clear();
}
}
if !batch.is_empty() {
store.delete(batch).await?;
}
co.yield_(GcSweepEvent::CustomDebug(format!(
"deleted {} blobs",
count
)))
.await;
Ok(())
}
#[derive(Debug)]
pub enum GcMarkEvent {
CustomDebug(String),
CustomWarning(String, Option<anyhow::Error>),
Error(anyhow::Error),
}
#[derive(Debug)]
pub enum GcSweepEvent {
CustomDebug(String),
CustomWarning(String, Option<anyhow::Error>),
Error(anyhow::Error),
}
#[allow(missing_docs)]
#[derive(Debug)]
pub enum ImportProgress {
Found { id: u64, name: String },
CopyProgress { id: u64, offset: u64 },
Size { id: u64, size: u64 },
OutboardProgress { id: u64, offset: u64 },
OutboardDone { id: u64, hash: Hash },
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ImportMode {
#[default]
Copy,
TryReference,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize, Serialize)]
pub enum ExportMode {
#[default]
Copy,
TryReference,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub enum ExportFormat {
#[default]
Blob,
Collection,
}
#[allow(missing_docs)]
#[derive(Debug)]
pub enum ExportProgress {
Start {
id: u64,
hash: Hash,
path: PathBuf,
stable: bool,
},
Progress { id: u64, offset: u64 },
Done { id: u64 },
}
#[derive(
Debug, Clone, Copy, derive_more::Display, Serialize, Deserialize, PartialOrd, Ord, PartialEq, Eq,
)]
pub enum ReportLevel {
Trace,
Info,
Warn,
Error,
}
#[derive(Debug, Serialize, Deserialize)]
pub enum ConsistencyCheckProgress {
Start,
Update {
message: String,
entry: Option<Hash>,
level: ReportLevel,
},
Done,
Abort(RpcError),
}
#[derive(Debug, Serialize, Deserialize)]
pub enum ValidateProgress {
Starting {
total: u64,
},
Entry {
id: u64,
hash: Hash,
path: Option<String>,
size: u64,
},
EntryProgress {
id: u64,
offset: u64,
},
EntryDone {
id: u64,
error: Option<String>,
},
PartialEntry {
id: u64,
hash: Hash,
path: Option<String>,
size: u64,
},
PartialEntryProgress {
id: u64,
offset: u64,
},
PartialEntryDone {
id: u64,
ranges: RangeSpec,
},
AllDone,
Abort(RpcError),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Event {
GcStarted,
GcCompleted,
}