hyperfile_cleaner/
s3.rs

1use std::io::{Result, Error, ErrorKind};
2use std::collections::BTreeSet;
3use aws_sdk_s3::Client;
4use log::{debug, error};
5use btree_ondisk::bmap::BMap;
6use btree_ondisk::BlockLoader;
7use btree_ondisk::NullNodeCache;
8use hyperfile::{BlockIndex, BlockPtr, SegmentId};
9use hyperfile::meta_loader::s3::S3BlockLoader;
10use hyperfile::ondisk::InodeRaw;
11use hyperfile::inode::{Inode, OnDiskState};
12use hyperfile::segment::{Segment, SegmentSum, SegmentReadWrite};
13use hyperfile::staging::{Staging, config::StagingConfig, s3::S3Staging};
14use hyperfile::config::HyperFileRuntimeConfig;
15use hyperfile::s3uri::S3Uri;
16use hyperfile::meta_format::BlockPtrFormat;
17use crate::Cleaner;
18
19pub struct S3Cleaner {
20    client: Client,
21    bucket: String,
22    root_path: String,
23    staging: S3Staging,
24}
25
26impl S3Cleaner {
27    pub async fn new(root_uri: &str) -> Result<Self> {
28        let root_uri = if let Some(remove_tail_slash) = root_uri.strip_suffix('/') {
29            remove_tail_slash
30        } else {
31            root_uri
32        };
33        let s3uri = S3Uri::parse(root_uri).expect("failed to parse root path for s3");
34        let config = aws_config::load_from_env().await;
35        let client = aws_sdk_s3::Client::new(&config);
36        let runtime_config = HyperFileRuntimeConfig::default();
37        let staging_config = StagingConfig::new_s3_staging(s3uri.bucket, "", s3uri.key, None);
38        let staging = S3Staging::from(&client, staging_config, runtime_config).await?;
39        Ok(Self {
40            client: client,
41            bucket: s3uri.bucket.to_owned(),
42            root_path: s3uri.key.to_owned(),
43            staging: staging,
44        })
45    }
46
47    async fn load_segment_sum(&self, segid: SegmentId) -> Result<SegmentSum> {
48        self.staging.open(segid).await
49    }
50
51    async fn load_raw_inode(&self) -> Result<(InodeRaw, Option<OnDiskState>)> {
52        let mut raw_inode: InodeRaw = unsafe { std::mem::MaybeUninit::zeroed().assume_init() };
53        let inode_state = self.staging.load_inode(&mut raw_inode.as_mut_u8_slice()).await?;
54        Ok((raw_inode, inode_state))
55    }
56
57	fn collect_segment_id<'a, L>(bmap: BMap<'a, BlockIndex, BlockPtr, BlockPtr, L, NullNodeCache>) -> BTreeSet<u64>
58    	where L: BlockLoader<BlockPtr> + Clone,
59	{
60    	// collect segment id in u64
61    	let mut segid_set: BTreeSet<u64> = BTreeSet::new();
62
63    	let iter = bmap.nonleafnode_iter();
64    	for blkptr in iter {
65    	    if BlockPtrFormat::is_on_staging(&blkptr) {
66    	        let segid = BlockPtrFormat::decode_segid(&blkptr);
67    	        segid_set.insert(segid);
68    	    }
69    	}
70		segid_set
71	}
72
73    async fn do_delete_objects(client: &Client, bucket: &str, delete_keys: Vec<String>) -> Result<()> {
74        use aws_sdk_s3::types::ObjectIdentifier;
75        let mut err = false;
76        for keys in delete_keys.chunks(1000).into_iter() {
77            let obj_ids = keys.into_iter().map(|k| {
78                    // FIXME: move the value, avoid ref
79                    ObjectIdentifier::builder()
80                        .key(k.to_string())
81                        .build()
82                        .unwrap()
83                }).collect::<Vec<ObjectIdentifier>>();
84            let delete = aws_sdk_s3::types::Delete::builder()
85                    .set_objects(Some(obj_ids))
86                    .quiet(true)
87                    .build()
88                    .unwrap();
89            match client.delete_objects()
90                .bucket(bucket)
91                .delete(delete)
92                .send()
93                .await
94            {
95                Ok(_) => {},
96                Err(sdk_err) => {
97                    err = true;
98                    // TODO: check delete objects result with signle delete error
99                    error!("delete objects error: {}", sdk_err);
100                }
101            }
102        }
103        if err {
104            return Err(Error::new(ErrorKind::Interrupted, "at least one of delete objects op failed"));
105        }
106        Ok(())
107    }
108
109    // collect referenced and unreferenced segment between low and high
110    // return:
111    //   (referenced seg list, unreferenced seg list)
112    async fn collect_unref_seg(&self, low: SegmentId, high: SegmentId) -> Result<(Vec<SegmentId>, Vec<SegmentId>)> {
113        let high_segid = if high == 0 {
114            // load last inode from disk
115            let (raw_inode, inode_state) = self.load_raw_inode().await?;
116            let inode = Inode::from_raw(&raw_inode, inode_state);
117            inode.get_last_seq()
118        } else {
119            high
120        };
121
122        debug!("collect unreferenced segment low: {}, high: {}", low, high_segid);
123        if low >= high_segid {
124            let err_msg = format!("invalid segment range [{}, {}]", low, high_segid);
125            return Err(Error::new(ErrorKind::InvalidInput, err_msg));
126        }
127
128        // load segment sum
129        let ss = self.load_segment_sum(high_segid).await?;
130        let i_bmap = ss.hdr.s_inode.i_bmap;
131        let meta_block_size = 1 << ss.hdr.s_meta_blk_shift;
132
133        // build bmap
134        let loader = S3BlockLoader::new(&self.client, &self.bucket, &self.root_path);
135        let set = tokio::task::spawn_blocking(move || {
136            let bmap = BMap::<BlockIndex, BlockPtr, BlockPtr, S3BlockLoader, NullNodeCache>::read(&i_bmap, meta_block_size, loader, NullNodeCache);
137            Self::collect_segment_id(bmap)
138        }).await.expect("spawn blocking task failed");
139
140        let high = high_segid;
141        // build non-referenced list
142        let mut non_ref = Vec::new();
143        for n in low..=high {
144            if set.get(&n).is_none() {
145                non_ref.push(n);
146            }
147        }
148        let v_ref = set.into_iter().map(|segid| segid).collect::<Vec<SegmentId>>();
149        Ok((v_ref, non_ref))
150    }
151}
152
153impl Cleaner for S3Cleaner {
154    async fn collect(&self, low: SegmentId, high: SegmentId) -> Result<(Vec<SegmentId>, Vec<SegmentId>)> {
155        self.collect_unref_seg(low, high).await
156    }
157
158    // prune unreferenced segments between low and high
159    async fn prune(&self, low: SegmentId, high: SegmentId) -> Result<()> {
160        let (_, non_ref) = self.collect_unref_seg(low, high).await?;
161        let delete_keys = non_ref.iter().map(|segid| {
162                                format!("{}/{}", self.root_path, Segment::segid_to_staging_file_id(*segid))
163                            }).collect::<Vec<String>>();
164        debug!("unrefereneced segments to delete {:?}", delete_keys);
165        Self::do_delete_objects(&self.client, &self.bucket, delete_keys).await?;
166        Ok(())
167    }
168}
hyperfile_cleaner/s3.rs

hyperfile_cleaner/
s3.rs
