use std::io::{Result, Error, ErrorKind};
use std::collections::BTreeSet;
use std::sync::Arc;
use aws_sdk_s3::Client;
use log::{debug, error};
use btree_ondisk::bmap::BMap;
use btree_ondisk::BlockLoader;
use btree_ondisk::NullNodeCache;
use hyperfile::{BlockIndex, BlockPtr, Cno, SegmentId, s3_code_kind, s3_error_kind};
use hyperfile::file::ReadTiming;
use hyperfile::meta_loader::s3::S3BlockLoader;
use hyperfile::ondisk::InodeRaw;
use hyperfile::inode::{Inode, OnDiskState};
use hyperfile::segment::{Segment, SegmentSum, SegmentReadWrite};
use hyperfile::staging::{Staging, config::StagingConfig, s3::S3Staging};
use hyperfile::config::{HyperFileMetaConfig, HyperFileRuntimeConfig};
use hyperfile::s3uri::S3Uri;
use hyperfile::meta_format::BlockPtrFormat;
use crate::Cleaner;
pub struct S3Cleaner {
client: Client,
bucket: String,
root_path: String,
staging: S3Staging,
}
impl S3Cleaner {
pub async fn new(root_uri: &str) -> Result<Self> {
let cfg = aws_config::load_from_env().await;
let client = Client::new(&cfg);
Self::with_client(&client, root_uri).await
}
pub async fn with_client(client: &Client, root_uri: &str) -> Result<Self> {
let trimmed = trim_root_uri(root_uri);
let s3uri = S3Uri::parse(trimmed).map_err(|e| {
Error::new(
ErrorKind::InvalidInput,
format!("invalid S3 URI {:?}: {}", root_uri, e),
)
})?;
let runtime_config = HyperFileRuntimeConfig::default();
let staging_config = StagingConfig::new_s3_staging(s3uri.bucket, "", s3uri.key, None);
let staging = S3Staging::from(client, staging_config, runtime_config).await?;
Ok(Self {
client: client.clone(),
bucket: s3uri.bucket.to_owned(),
root_path: s3uri.key.to_owned(),
staging,
})
}
async fn load_segment_sum(&self, segid: SegmentId) -> Result<SegmentSum> {
self.staging.open(segid).await
}
pub(crate) fn root_uri(&self) -> String {
format!("s3://{}/{}", self.bucket, self.root_path)
}
pub(crate) fn client(&self) -> &Client {
&self.client
}
async fn load_raw_inode(&self) -> Result<(InodeRaw, Option<OnDiskState>)> {
let mut raw_inode = InodeRaw::default();
let inode_state = self.staging.load_inode(raw_inode.as_mut_u8_slice()).await?;
Ok((raw_inode, inode_state))
}
fn walk_live_ptrs<L>(
bmap: BMap<'_, BlockIndex, BlockPtr, BlockPtr, L, NullNodeCache>,
) -> Result<LivePtrs>
where
L: BlockLoader<BlockPtr> + Clone,
{
let handle = tokio::runtime::Handle::current();
let on_staging = |v: Vec<BlockPtr>| -> BTreeSet<BlockPtr> {
v.into_iter().filter(BlockPtrFormat::is_on_staging).collect()
};
let all = on_staging(handle.block_on(bmap.collect_ptrs())?);
let meta = on_staging(handle.block_on(bmap.collect_meta_ptrs())?);
let mut windows: Vec<(BlockIndex, usize)> = Vec::new();
let mut pos: BlockIndex = 0;
loop {
let key = match handle.block_on(bmap.seek_key(&pos)) {
Ok(key) => key,
Err(e) if e.kind() == ErrorKind::NotFound => break,
Err(e) => return Err(e),
};
match windows.last_mut() {
Some((start, n)) if *start + *n as BlockIndex == key
&& *n + DATA_WINDOW_BLOCKS <= 4 * DATA_WINDOW_BLOCKS =>
{
*n += DATA_WINDOW_BLOCKS;
},
_ => windows.push((key, DATA_WINDOW_BLOCKS)),
}
match key.checked_add(DATA_WINDOW_BLOCKS as BlockIndex) {
Some(next) => pos = next,
None => break,
}
}
Ok(LivePtrs { all, meta, windows })
}
async fn list_segment_objects(&self) -> Result<ObjectListing> {
let prefix = format!("{}/", self.root_path);
let mut listing = ObjectListing::default();
let mut pages = self.client.list_objects_v2()
.bucket(&self.bucket)
.prefix(&prefix)
.into_paginator()
.send();
while let Some(page) = pages.next().await {
let page = page.map_err(|e| Error::new(
s3_error_kind(&e),
format!("listing segment objects under s3://{}/{} failed: {}",
self.bucket, prefix, e),
))?;
for obj in page.contents() {
let Some(key) = obj.key() else { continue };
if key == self.staging.inode_file {
continue;
}
let Some(name) = key.strip_prefix(&prefix) else { continue };
match parse_segment_object_name(name) {
Some(segid) => {
listing.objects.push((segid, obj.size().unwrap_or(0).max(0) as u64));
},
None => {
listing.foreign += 1;
if listing.foreign_sample.is_none() {
listing.foreign_sample = Some(key.to_string());
}
},
}
}
}
listing.objects.sort();
Ok(listing)
}
async fn do_delete_objects(
client: &Client,
bucket: &str,
delete_keys: Vec<String>,
) -> Result<()> {
use aws_sdk_s3::types::ObjectIdentifier;
if delete_keys.is_empty() {
return Ok(());
}
let mut failed_keys: Vec<String> = Vec::new();
let mut first_kind: Option<ErrorKind> = None;
for keys in delete_keys.chunks(1000) {
let obj_ids = keys.iter()
.map(|k| {
ObjectIdentifier::builder()
.key(k)
.build()
.expect("ObjectIdentifier requires non-empty key")
})
.collect::<Vec<ObjectIdentifier>>();
let delete = aws_sdk_s3::types::Delete::builder()
.set_objects(Some(obj_ids))
.quiet(true)
.build()
.expect("Delete::builder requires non-empty object list");
match client.delete_objects()
.bucket(bucket)
.delete(delete)
.send()
.await
{
Ok(output) => {
if let Some(errors) = output.errors {
for e in errors {
let key = e.key.unwrap_or_default();
let code = e.code.as_deref().unwrap_or("?");
let msg = e.message.as_deref().unwrap_or("");
error!("DeleteObjects per-key failure: key={} code={} message={}",
key, code, msg);
first_kind.get_or_insert(s3_code_kind(e.code.as_deref()));
if !key.is_empty() {
failed_keys.push(key);
}
}
}
},
Err(sdk_err) => {
error!("DeleteObjects batch failed (bucket={}, count={}): {}",
bucket, keys.len(), sdk_err);
first_kind.get_or_insert(s3_error_kind(&sdk_err));
failed_keys.extend(keys.iter().cloned());
}
}
}
if failed_keys.is_empty() {
Ok(())
} else {
Err(Error::new(
first_kind.unwrap_or(ErrorKind::Other),
format!(
"DeleteObjects failed for {} object(s); first failed key: {:?}",
failed_keys.len(),
failed_keys.first(),
),
))
}
}
pub(crate) async fn view(&self, low: Cno, high: Cno) -> Result<ContainerView> {
let (raw_inode, inode_state) = self.load_raw_inode().await?;
refuse_unsegmented(&raw_inode)?;
let fmt = container_ptr_format(&raw_inode);
let latest_compacted = raw_inode.is_compacted();
let high_cno = if high == 0 {
let inode = Inode::from_raw(&raw_inode, inode_state);
inode.get_last_cno()
} else {
high
};
debug!("view of container low: {}, high: {}, format: {:?}", low, high_cno, fmt);
validate_segid_range(low, high_cno)?;
let high_summary = summary_obj(high_cno, &fmt);
let ss = self.load_segment_sum(high_summary).await?;
let ss_inode = ss.hdr.s_inode;
let meta_block_size = 1usize << ss.hdr.s_meta_blk_shift;
let data_block_size = 1usize << ss.hdr.s_data_blk_shift;
let loader = S3BlockLoader::new(
&self.client, &self.bucket, &self.root_path,
Arc::new(ReadTiming::default()),
);
let join = tokio::task::spawn_blocking(move || -> Result<LivePtrs> {
let bmap = BMap::<BlockIndex, BlockPtr, BlockPtr, S3BlockLoader, NullNodeCache>
::read(&ss_inode.i_bmap, meta_block_size, loader, NullNodeCache)?;
Self::walk_live_ptrs(bmap)
});
let live = match join.await {
Ok(inner) => inner?,
Err(join_err) => {
return Err(Error::other(
format!("bmap walk task failed to join: {}", join_err),
));
}
};
let listing = self.list_segment_objects().await?;
Ok(ContainerView {
low,
high_cno,
latest_compacted,
fmt,
high_summary,
data_block_size,
meta_block_size,
live,
listing,
})
}
async fn collect_unref_seg(
&self,
low: Cno,
high: Cno,
) -> Result<(Vec<SegmentId>, Vec<SegmentId>)> {
let view = self.view(low, high).await?;
Ok(partition_objects(
&view.object_ids(),
&view.referenced_objects(),
view.low,
view.high_cno,
view.high_summary,
))
}
}
pub(crate) struct LivePtrs {
pub(crate) all: BTreeSet<BlockPtr>,
pub(crate) meta: BTreeSet<BlockPtr>,
pub(crate) windows: Vec<(BlockIndex, usize)>,
}
pub(crate) const DATA_WINDOW_BLOCKS: usize = 64 * 1024;
#[derive(Default)]
pub(crate) struct ObjectListing {
pub(crate) objects: Vec<(SegmentId, u64)>,
pub(crate) foreign: usize,
pub(crate) foreign_sample: Option<String>,
}
pub(crate) struct ContainerView {
pub(crate) low: Cno,
pub(crate) high_cno: Cno,
pub(crate) latest_compacted: bool,
pub(crate) fmt: BlockPtrFormat,
pub(crate) high_summary: SegmentId,
pub(crate) data_block_size: usize,
pub(crate) meta_block_size: usize,
pub(crate) live: LivePtrs,
pub(crate) listing: ObjectListing,
}
impl ContainerView {
pub(crate) fn object_ids(&self) -> Vec<SegmentId> {
self.listing.objects.iter().map(|(segid, _)| *segid).collect()
}
pub(crate) fn referenced_objects(&self) -> BTreeSet<SegmentId> {
self.live.all.iter()
.map(|p| BlockPtrFormat::decode(p, &self.fmt).0)
.collect()
}
pub(crate) fn object_of(&self, blkptr: &BlockPtr) -> SegmentId {
BlockPtrFormat::decode(blkptr, &self.fmt).0
}
}
fn trim_root_uri(root_uri: &str) -> &str {
match root_uri.strip_suffix('/') {
Some(r) if !r.ends_with(":/") => r,
_ => root_uri,
}
}
const S_IFMT: u32 = 0o170000;
const S_IFCHR: u32 = 0o020000;
const S_IFBLK: u32 = 0o060000;
fn refuse_unsegmented(raw: &InodeRaw) -> Result<()> {
let fmt = raw.i_mode & S_IFMT;
let what = if raw.is_inline() {
Some("an inline file, whose data lives in the inode")
} else if fmt == S_IFCHR || fmt == S_IFBLK {
Some("a device node, which stores rdev where a checkpoint number would be")
} else {
None
};
match what {
Some(what) => Err(Error::new(
ErrorKind::InvalidInput,
format!("this container is {}: it has no segments to collect, prune or compact", what),
)),
None => Ok(()),
}
}
fn container_ptr_format(raw_inode: &InodeRaw) -> BlockPtrFormat {
match HyperFileMetaConfig::try_from_u32(raw_inode.i_meta_config) {
Ok(cfg) => cfg.block_ptr_format,
Err(e) => {
debug!("inode meta config {:#x} is not one this build can read ({}); \
reading pointers as {:?}", raw_inode.i_meta_config, e,
BlockPtrFormat::MicroGroup);
BlockPtrFormat::MicroGroup
},
}
}
fn summary_obj(cno: Cno, fmt: &BlockPtrFormat) -> SegmentId {
let segid = SegmentId::new_from_cno(cno);
if fmt.is_parted() {
segid.at_part(Segment::SUMMARY_PART)
} else {
segid
}
}
fn parse_segment_object_name(name: &str) -> Option<SegmentId> {
if name.contains('/') {
return None;
}
let (seq_str, part_str) = match name.split_once('.') {
Some((seq, part)) => (seq, Some(part)),
None => (name, None),
};
if seq_str.len() != 10 || !seq_str.bytes().all(|b| b.is_ascii_digit()) {
return None;
}
let seq_id = seq_str.parse::<u32>().ok()?;
match part_str {
None => Some(SegmentId::new(seq_id)),
Some(p) => {
if p.is_empty() || !p.bytes().all(|b| b.is_ascii_digit()) {
return None;
}
Some(SegmentId::with_part(seq_id, p.parse::<u16>().ok()?))
},
}
}
fn partition_objects(
existing: &[SegmentId],
referenced: &BTreeSet<SegmentId>,
low: Cno,
high: Cno,
high_summary: SegmentId,
) -> (Vec<SegmentId>, Vec<SegmentId>) {
let live_cnos: BTreeSet<Cno> = referenced.iter().map(|s| s.as_cno()).collect();
let mut unref = Vec::new();
for obj in existing {
let cno = obj.as_cno();
if cno < low || cno > high {
continue;
}
if referenced.contains(obj) {
continue;
}
if *obj == high_summary {
continue;
}
if !obj.is_partial() && live_cnos.contains(&cno) {
continue;
}
unref.push(*obj);
}
let v_ref = referenced.iter().copied().collect::<Vec<SegmentId>>();
(v_ref, unref)
}
fn validate_segid_range(low: Cno, high_cno: Cno) -> Result<()> {
if low > high_cno {
return Err(Error::new(
ErrorKind::InvalidInput,
format!("invalid segment range [{}, {}]", low, high_cno),
));
}
if high_cno > u32::MAX as Cno {
return Err(Error::new(
ErrorKind::InvalidInput,
format!("checkpoint {} is past what a block pointer can name ({})",
high_cno, u32::MAX),
));
}
Ok(())
}
impl Cleaner for S3Cleaner {
async fn collect(
&self,
low: Cno,
high: Cno,
) -> Result<(Vec<SegmentId>, Vec<SegmentId>)> {
self.collect_unref_seg(low, high).await
}
async fn prune(&self, low: Cno, high: Cno) -> Result<()> {
let (_, non_ref) = self.collect_unref_seg(low, high).await?;
let delete_keys = non_ref.iter()
.map(|segid| format!(
"{}/{}",
self.root_path,
Segment::segid_to_staging_file_id(*segid),
))
.collect::<Vec<String>>();
debug!("unreferenced segments to delete: {:?}", delete_keys);
Self::do_delete_objects(&self.client, &self.bucket, delete_keys).await
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn delete_empty_keys_is_noop() {
let cfg = aws_config::load_from_env().await;
let client = Client::new(&cfg);
let r = S3Cleaner::do_delete_objects(&client, "any-bucket", Vec::new()).await;
assert!(r.is_ok(), "empty key list should not call S3");
}
#[test]
fn segid_range_accepts_a_single_checkpoint() {
assert!(validate_segid_range(1, 1).is_ok());
}
#[test]
fn segid_range_accepts_a_normal_span() {
assert!(validate_segid_range(1, 5).is_ok());
}
#[test]
fn segid_range_rejects_an_empty_span() {
for (low, high) in [(5u64, 3u64), (2, 1), (1, 0)] {
let e = validate_segid_range(low, high)
.expect_err("low > high must be rejected");
assert_eq!(e.kind(), ErrorKind::InvalidInput);
assert!(
e.to_string().contains(&format!("[{}, {}]", low, high)),
"error should name the range, got {:?}", e.to_string(),
);
}
}
#[test]
fn a_trailing_slash_is_trimmed_but_not_the_schemes() {
assert_eq!(trim_root_uri("s3://b/k/"), "s3://b/k");
assert_eq!(trim_root_uri("s3://b/k"), "s3://b/k");
assert_eq!(trim_root_uri("s3://"), "s3://");
assert_eq!(trim_root_uri("S3://"), "S3://");
assert_eq!(trim_root_uri(""), "");
}
#[test]
fn an_inode_without_segments_is_refused() {
let mut inline = InodeRaw { i_mode: 0o100644, ..Default::default() };
inline.set_inline(b"hello");
let e = refuse_unsegmented(&inline).expect_err("inline must be refused");
assert_eq!(e.kind(), ErrorKind::InvalidInput);
assert!(e.to_string().contains("inline"), "{}", e);
for mode in [0o020644u32, 0o060644] {
let dev = InodeRaw { i_mode: mode, ..Default::default() };
let e = refuse_unsegmented(&dev).expect_err("a device node must be refused");
assert!(e.to_string().contains("device"), "{}", e);
}
for mode in [0o100644u32, 0o040755, 0o120777] {
let ok = InodeRaw { i_mode: mode, ..Default::default() };
assert!(refuse_unsegmented(&ok).is_ok(), "mode {:o} has segments", mode);
}
}
#[test]
fn segid_range_rejects_a_checkpoint_past_a_pointer() {
let e = validate_segid_range(1, u32::MAX as Cno + 1)
.expect_err("a cno past u32::MAX must be rejected");
assert_eq!(e.kind(), ErrorKind::InvalidInput);
assert!(validate_segid_range(1, u32::MAX as Cno).is_ok());
}
#[test]
fn object_names_round_trip_through_the_parser() {
for segid in [
SegmentId::new(1),
SegmentId::new(4294967295),
SegmentId::with_part(7, 0),
SegmentId::with_part(7, 16383),
] {
let name = Segment::segid_to_staging_file_id(segid);
assert_eq!(parse_segment_object_name(&name), Some(segid),
"name {:?} did not parse back", name);
}
}
#[test]
fn a_bare_name_is_not_part_zero() {
assert_eq!(parse_segment_object_name("0000000007"), Some(SegmentId::new(7)));
assert_eq!(parse_segment_object_name("0000000007.0"),
Some(SegmentId::with_part(7, 0)));
}
#[test]
fn keys_that_are_not_segment_objects_are_skipped() {
for name in [
"inode", "wal/0000000001", "0000000001/0_0_4096",
"000000001", "00000000012", "9999999999", "0000000001.", "0000000001.x",
"0000000001.65536", "0000000001.0.1",
"",
] {
assert_eq!(parse_segment_object_name(name), None,
"{:?} should not be read as a segment object", name);
}
}
fn partition(
existing: &[SegmentId],
referenced: &[SegmentId],
low: Cno,
high: Cno,
high_summary: SegmentId,
) -> Vec<SegmentId> {
let set: BTreeSet<SegmentId> = referenced.iter().copied().collect();
let (_, unref) = partition_objects(existing, &set, low, high, high_summary);
unref
}
#[test]
fn an_unreferenced_checkpoint_is_reclaimable() {
let existing = [SegmentId::new(1), SegmentId::new(2), SegmentId::new(3)];
let unref = partition(&existing, &[SegmentId::new(3)], 1, 3, SegmentId::new(3));
assert_eq!(unref, vec![SegmentId::new(1), SegmentId::new(2)]);
}
#[test]
fn objects_outside_the_range_are_left_alone() {
let existing = [
SegmentId::new(1), SegmentId::new(2), SegmentId::new(3), SegmentId::new(4),
];
let unref = partition(&existing, &[], 2, 3, SegmentId::new(9));
assert_eq!(unref, vec![SegmentId::new(2), SegmentId::new(3)]);
}
#[test]
fn the_summary_of_high_is_kept_even_when_nothing_points_into_it() {
let existing = [SegmentId::new(1), SegmentId::new(2)];
let unref = partition(&existing, &[], 1, 2, SegmentId::new(2));
assert_eq!(unref, vec![SegmentId::new(1)]);
let existing = [SegmentId::with_part(2, 0), SegmentId::with_part(2, 1)];
let unref = partition(&existing, &[], 1, 2, SegmentId::with_part(2, 0));
assert_eq!(unref, vec![SegmentId::with_part(2, 1)]);
}
#[test]
fn an_abandoned_partial_is_reclaimable_while_its_checkpoint_lives() {
let existing = [
SegmentId::with_part(5, 0), SegmentId::with_part(5, 1), SegmentId::with_part(5, 2), ];
let referenced = [SegmentId::with_part(5, 0), SegmentId::with_part(5, 1)];
let unref = partition(&existing, &referenced, 1, 9, SegmentId::with_part(9, 0));
assert_eq!(unref, vec![SegmentId::with_part(5, 2)]);
}
#[test]
fn the_summary_part_of_a_live_checkpoint_is_kept() {
let existing = [SegmentId::with_part(4, 0), SegmentId::with_part(4, 3)];
let unref = partition(&existing, &[SegmentId::with_part(4, 3)], 1, 9,
SegmentId::with_part(9, 0));
assert!(unref.is_empty(), "part 0 describes the live part 3, got {:?}", unref);
}
#[test]
fn a_wholly_dead_parted_checkpoint_goes_including_its_summary() {
let existing = [
SegmentId::with_part(4, 0), SegmentId::with_part(4, 1),
SegmentId::with_part(6, 0),
];
let unref = partition(&existing, &[SegmentId::with_part(6, 0)], 1, 9,
SegmentId::with_part(6, 0));
assert_eq!(unref, vec![SegmentId::with_part(4, 0), SegmentId::with_part(4, 1)]);
}
}