Skip to main content

hyperfile_cleaner/
s3.rs

1//! S3-backed cleaner implementation.
2//!
3//! `S3Cleaner` enumerates the segments referenced by the latest
4//! published inode of a hyperfile and (optionally) deletes anything
5//! in a configured `[low, high]` segid range that isn't referenced.
6//!
7//! See the crate-level [`Cleaner`] trait for the public API shape.
8
9use std::io::{Result, Error, ErrorKind};
10use std::collections::BTreeSet;
11use aws_sdk_s3::Client;
12use log::{debug, error};
13use btree_ondisk::bmap::BMap;
14use btree_ondisk::BlockLoader;
15use btree_ondisk::NullNodeCache;
16use hyperfile::{BlockIndex, BlockPtr, SegmentId};
17use hyperfile::meta_loader::s3::S3BlockLoader;
18use hyperfile::ondisk::InodeRaw;
19use hyperfile::inode::{Inode, OnDiskState};
20use hyperfile::segment::{Segment, SegmentSum, SegmentReadWrite};
21use hyperfile::staging::{Staging, config::StagingConfig, s3::S3Staging};
22use hyperfile::config::HyperFileRuntimeConfig;
23use hyperfile::s3uri::S3Uri;
24use hyperfile::meta_format::BlockPtrFormat;
25use crate::Cleaner;
26
27/// S3-backed cleaner for a single hyperfile rooted at a given
28/// `s3://<bucket>/<prefix>` URI.
29///
30/// The cleaner does not retain or open the file's data. It only
31/// reads the published inode and walks the metadata B+tree to
32/// enumerate referenced segments, then deletes the unreferenced
33/// ones via `DeleteObjects`.
34///
35/// Concurrency: the cleaner is not aware of in-flight writers. If
36/// a writer publishes a new checkpoint while `collect` /
37/// `prune` is running, the result reflects whatever inode was
38/// loaded at the start of the call. New segments published after
39/// that point are outside the configured `[low, high]` range and
40/// are not affected.
41pub struct S3Cleaner {
42    client: Client,
43    bucket: String,
44    root_path: String,
45    staging: S3Staging,
46}
47
48impl S3Cleaner {
49    /// Construct a cleaner for `root_uri`, using the AWS SDK's
50    /// default credential / region resolution
51    /// (`aws_config::load_from_env`).
52    ///
53    /// Returns `InvalidInput` if `root_uri` is not a valid S3 URI.
54    pub async fn new(root_uri: &str) -> Result<Self> {
55        let cfg = aws_config::load_from_env().await;
56        let client = Client::new(&cfg);
57        Self::with_client(&client, root_uri).await
58    }
59
60    /// Construct a cleaner using a caller-supplied `Client`. Use
61    /// this when you need to control region / endpoint / credential
62    /// resolution explicitly (for example, in tests or in a tool
63    /// that already has a configured client).
64    ///
65    /// Returns `InvalidInput` if `root_uri` is not a valid S3 URI.
66    pub async fn with_client(client: &Client, root_uri: &str) -> Result<Self> {
67        let trimmed = root_uri.strip_suffix('/').unwrap_or(root_uri);
68        let s3uri = S3Uri::parse(trimmed).map_err(|e| {
69            Error::new(
70                ErrorKind::InvalidInput,
71                format!("invalid S3 URI {:?}: {}", root_uri, e),
72            )
73        })?;
74        let runtime_config = HyperFileRuntimeConfig::default();
75        let staging_config = StagingConfig::new_s3_staging(s3uri.bucket, "", s3uri.key, None);
76        let staging = S3Staging::from(client, staging_config, runtime_config).await?;
77        Ok(Self {
78            client: client.clone(),
79            bucket: s3uri.bucket.to_owned(),
80            root_path: s3uri.key.to_owned(),
81            staging,
82        })
83    }
84
85    async fn load_segment_sum(&self, segid: SegmentId) -> Result<SegmentSum> {
86        self.staging.open(segid).await
87    }
88
89    async fn load_raw_inode(&self) -> Result<(InodeRaw, Option<OnDiskState>)> {
90        let mut raw_inode = InodeRaw::default();
91        let inode_state = self.staging.load_inode(raw_inode.as_mut_u8_slice()).await?;
92        Ok((raw_inode, inode_state))
93    }
94
95    /// Walk the B+tree starting from the supplied `bmap` root and
96    /// collect the set of segment IDs that hold any of its
97    /// non-leaf-node-stored block pointers. In hyperfile's CoW
98    /// model this set is sufficient to identify all segments that
99    /// must be retained.
100    fn collect_segment_id<'a, L>(
101        bmap: BMap<'a, BlockIndex, BlockPtr, BlockPtr, L, NullNodeCache>,
102    ) -> BTreeSet<SegmentId>
103    where
104        L: BlockLoader<BlockPtr> + Clone,
105    {
106        let mut segid_set: BTreeSet<SegmentId> = BTreeSet::new();
107        for blkptr in bmap.nonleafnode_iter() {
108            if BlockPtrFormat::is_on_staging(&blkptr) {
109                segid_set.insert(BlockPtrFormat::decode_segid(&blkptr));
110            }
111        }
112        segid_set
113    }
114
115    /// Delete `delete_keys` from `bucket` in batches of 1000 (the
116    /// `DeleteObjects` per-call max).
117    ///
118    /// Failures are tracked at two levels and merged into the same
119    /// returned error:
120    ///
121    /// 1. **SDK-level (transport / auth / etc.)**: the entire batch
122    ///    is recorded as failed. The SDK error is logged at
123    ///    `error!` level.
124    /// 2. **Per-key (S3 partial success in `DeleteObjectsOutput`)**:
125    ///    `DeleteObjects` is a partial-success API; even when the
126    ///    batch returns `Ok`, individual keys may have failed (e.g.
127    ///    `AccessDenied`, `NoSuchKey`). The response's `errors`
128    ///    list is parsed and each failing key recorded.
129    ///
130    /// Returns `Err(Interrupted)` if any key (at either level)
131    /// failed. The error message includes the failed object count
132    /// and a sample failed key. Successfully-deleted keys stay
133    /// deleted (S3 has no transactional batch delete).
134    async fn do_delete_objects(
135        client: &Client,
136        bucket: &str,
137        delete_keys: Vec<String>,
138    ) -> Result<()> {
139        use aws_sdk_s3::types::ObjectIdentifier;
140        if delete_keys.is_empty() {
141            return Ok(());
142        }
143        let mut failed_keys: Vec<String> = Vec::new();
144        for keys in delete_keys.chunks(1000) {
145            let obj_ids = keys.iter()
146                .map(|k| {
147                    ObjectIdentifier::builder()
148                        .key(k)
149                        .build()
150                        .expect("ObjectIdentifier requires non-empty key")
151                })
152                .collect::<Vec<ObjectIdentifier>>();
153            // We use quiet=true so the response only carries the
154            // `errors` list (not the success list); this is what
155            // we want for partial-failure tracking.
156            let delete = aws_sdk_s3::types::Delete::builder()
157                .set_objects(Some(obj_ids))
158                .quiet(true)
159                .build()
160                .expect("Delete::builder requires non-empty object list");
161            match client.delete_objects()
162                .bucket(bucket)
163                .delete(delete)
164                .send()
165                .await
166            {
167                Ok(output) => {
168                    // Partial-success path: even though the batch
169                    // returned Ok at the SDK level, individual keys
170                    // may have failed. Pick those out of the
171                    // response.
172                    if let Some(errors) = output.errors {
173                        for e in errors {
174                            let key = e.key.unwrap_or_default();
175                            let code = e.code.as_deref().unwrap_or("?");
176                            let msg = e.message.as_deref().unwrap_or("");
177                            error!("DeleteObjects per-key failure: key={} code={} message={}",
178                                key, code, msg);
179                            if !key.is_empty() {
180                                failed_keys.push(key);
181                            }
182                        }
183                    }
184                },
185                Err(sdk_err) => {
186                    // SDK-level failure: treat the whole batch as
187                    // failed (we don't know which keys made it).
188                    error!("DeleteObjects batch failed (bucket={}, count={}): {}",
189                        bucket, keys.len(), sdk_err);
190                    failed_keys.extend(keys.iter().cloned());
191                }
192            }
193        }
194        if failed_keys.is_empty() {
195            Ok(())
196        } else {
197            Err(Error::new(
198                ErrorKind::Interrupted,
199                format!(
200                    "DeleteObjects failed for {} object(s); first failed key: {:?}",
201                    failed_keys.len(),
202                    failed_keys.first(),
203                ),
204            ))
205        }
206    }
207
208    /// Inner shared implementation used by both `collect` and
209    /// `prune`. Returns `(referenced, unreferenced)` segment id
210    /// vectors where:
211    ///  - `referenced` is the set of segids reachable via the
212    ///    latest published inode's metadata tree, intersected with
213    ///    `[low, high]`.
214    ///  - `unreferenced` is `[low, high]` minus the referenced set.
215    ///
216    /// `high == 0` is a sentinel meaning "use the latest checkpoint
217    /// (`inode.get_last_seq()`)".
218    async fn collect_unref_seg(
219        &self,
220        low: SegmentId,
221        high: SegmentId,
222    ) -> Result<(Vec<SegmentId>, Vec<SegmentId>)> {
223        let high_segid = if high == 0 {
224            let (raw_inode, inode_state) = self.load_raw_inode().await?;
225            let inode = Inode::from_raw(&raw_inode, inode_state);
226            inode.get_last_seq()
227        } else {
228            high
229        };
230
231        debug!("collect unreferenced segment low: {}, high: {}", low, high_segid);
232        if low >= high_segid {
233            let err_msg = format!("invalid segment range [{}, {}]", low, high_segid);
234            return Err(Error::new(ErrorKind::InvalidInput, err_msg));
235        }
236
237        // load segment summary for the high checkpoint
238        let ss = self.load_segment_sum(high_segid).await?;
239        let i_bmap = ss.hdr.s_inode.i_bmap;
240        let meta_block_size = 1 << ss.hdr.s_meta_blk_shift;
241
242        // Walk the B+tree on a blocking thread (btree-ondisk's
243        // sync API parks on the runtime; we don't want to park
244        // the caller's reactor thread).
245        let loader = S3BlockLoader::new(&self.client, &self.bucket, &self.root_path);
246        let join = tokio::task::spawn_blocking(move || -> Result<BTreeSet<SegmentId>> {
247            let bmap = BMap::<BlockIndex, BlockPtr, BlockPtr, S3BlockLoader, NullNodeCache>
248                ::read(&i_bmap, meta_block_size, loader, NullNodeCache)?;
249            Ok(Self::collect_segment_id(bmap))
250        });
251        let set = match join.await {
252            Ok(inner) => inner?,
253            Err(join_err) => {
254                return Err(Error::other(
255                    format!("bmap walk task failed to join: {}", join_err),
256                ));
257            }
258        };
259
260        // Build the unreferenced list within [low, high].
261        let mut non_ref = Vec::new();
262        for n in low..=high_segid {
263            if !set.contains(&n) {
264                non_ref.push(n);
265            }
266        }
267        let v_ref = set.into_iter().collect::<Vec<SegmentId>>();
268        Ok((v_ref, non_ref))
269    }
270}
271
272impl Cleaner for S3Cleaner {
273    async fn collect(
274        &self,
275        low: SegmentId,
276        high: SegmentId,
277    ) -> Result<(Vec<SegmentId>, Vec<SegmentId>)> {
278        self.collect_unref_seg(low, high).await
279    }
280
281    /// Prune unreferenced segments in `[low, high]`. Implemented
282    /// as `collect` followed by `DeleteObjects` of the
283    /// unreferenced list.
284    async fn prune(&self, low: SegmentId, high: SegmentId) -> Result<()> {
285        let (_, non_ref) = self.collect_unref_seg(low, high).await?;
286        let delete_keys = non_ref.iter()
287            .map(|segid| format!(
288                "{}/{}",
289                self.root_path,
290                Segment::segid_to_staging_file_id(*segid),
291            ))
292            .collect::<Vec<String>>();
293        debug!("unreferenced segments to delete: {:?}", delete_keys);
294        Self::do_delete_objects(&self.client, &self.bucket, delete_keys).await
295    }
296}
297
298#[cfg(test)]
299mod tests {
300    use super::*;
301
302    /// Builder calls in `do_delete_objects` panic only on empty
303    /// inputs that we guard against. Verify the empty-input fast
304    /// path.
305    #[tokio::test]
306    async fn delete_empty_keys_is_noop() {
307        let cfg = aws_config::load_from_env().await;
308        let client = Client::new(&cfg);
309        // Empty list -> no SDK call -> Ok regardless of credentials.
310        let r = S3Cleaner::do_delete_objects(&client, "any-bucket", Vec::new()).await;
311        assert!(r.is_ok(), "empty key list should not call S3");
312    }
313}