1use 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
27pub struct S3Cleaner {
42 client: Client,
43 bucket: String,
44 root_path: String,
45 staging: S3Staging,
46}
47
48impl S3Cleaner {
49 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 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 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 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 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 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 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 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 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 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 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 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 #[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 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}