Line data Source code
1 : use anyhow::Result;
2 : use clap::Parser;
3 : use semver::Version;
4 : use std::collections::HashMap;
5 :
6 : use crate::remote_access::RemoteAccess;
7 :
8 : /// Delete all but the newest release versions from R2
9 : #[derive(Parser, Debug)]
10 : pub struct VacuumArgs {
11 : /// List the objects to delete without deleting them
12 : #[arg(long, required = false, default_value = "false")]
13 : dry_run: bool,
14 : /// The number of newest prerelease versions to keep
15 : #[arg(long, required = false, default_value_t = 10)]
16 : keep_latest: usize,
17 : /// The number of newest stable versions to keep
18 : #[arg(long, required = false, default_value_t = 100)]
19 : keep_stable: usize,
20 : }
21 :
22 31 : fn retained_versions(mut versions: Vec<Version>, keep: usize) -> (Vec<Version>, Vec<Version>) {
23 31 : versions.sort_by(|a, b| b.cmp(a));
24 31 : let deleted = versions.split_off(keep.min(versions.len()));
25 31 : (versions, deleted)
26 31 : }
27 :
28 3 : pub async fn vacuum(args: VacuumArgs) -> Result<()> {
29 : // Is a dry-run?
30 3 : if args.dry_run {
31 1 : log::info!("DRY RUN, no persistent changes will be made");
32 2 : }
33 :
34 : // Check R2 credentials
35 3 : let remote_access = RemoteAccess::from_env()?;
36 2 : log::info!("r2 bucket: **** (from environment)");
37 2 : log::info!("r2 access key id: **** (from environment)");
38 2 : log::info!("r2 secret access key: **** (from environment)");
39 :
40 2 : vacuum_bucket(&remote_access.r2_client(), &args).await
41 3 : }
42 :
43 : /// The objects of the versions that are not among the newest ones.
44 : /// Stable and prerelease versions are retained separately, keys without a version are never expired.
45 14 : fn expired_objects(keys: &[String], keep_latest: usize, keep_stable: usize) -> Vec<String> {
46 14 : let mut stable_objects: HashMap<Version, Vec<String>> = HashMap::new();
47 14 : let mut latest_objects: HashMap<Version, Vec<String>> = HashMap::new();
48 4064 : for key in keys {
49 4064 : let relative = key.strip_prefix("releases/").unwrap_or(key);
50 4064 : let version_text = &relative[..relative.find('/').unwrap_or(relative.len())];
51 4064 : let Ok(version) = semver::Version::parse(version_text) else { continue };
52 4059 : let objects = if version.pre.is_empty() {
53 4041 : stable_objects.entry(version).or_default()
54 : } else {
55 18 : latest_objects.entry(version).or_default()
56 : };
57 4059 : objects.push(key.clone());
58 : }
59 :
60 14 : let (keep_latest, delete_latest) = retained_versions(latest_objects.keys().cloned().collect(), keep_latest);
61 14 : let (keep_stable, delete_stable) = retained_versions(stable_objects.keys().cloned().collect(), keep_stable);
62 :
63 14 : log::info!("keep prerelease versions: {:?}", keep_latest);
64 14 : log::info!("keep stable versions: {:?}", keep_stable);
65 14 : log::info!("delete stable versions: {:?}", delete_stable);
66 14 : log::info!("delete prerelease versions: {:?}", delete_latest);
67 :
68 14 : let latest = delete_latest.iter().flat_map(|v| latest_objects[v].iter());
69 14 : let stable = delete_stable.iter().flat_map(|v| stable_objects[v].iter());
70 14 : latest.chain(stable).cloned().collect()
71 14 : }
72 :
73 : /// List every object of the versioned release directories
74 10 : async fn list_release_objects(client: &aws_sdk_s3::Client) -> Result<Vec<String>> {
75 10 : let mut keys = Vec::new();
76 10 : let mut continuation_token: Option<String> = None;
77 : loop {
78 16 : let results = client
79 16 : .list_objects_v2()
80 16 : .bucket("dashql-get")
81 16 : .prefix("releases/")
82 16 : .set_continuation_token(continuation_token)
83 16 : .send()
84 16 : .await?;
85 4034 : keys.extend(results.contents().iter().filter_map(|entry| entry.key()).map(str::to_owned));
86 15 : continuation_token = results.next_continuation_token().map(str::to_owned);
87 15 : if continuation_token.is_none() { break; }
88 : }
89 9 : Ok(keys)
90 10 : }
91 :
92 10 : async fn vacuum_bucket(client: &aws_sdk_s3::Client, args: &VacuumArgs) -> Result<()> {
93 10 : let keys = list_release_objects(client).await?;
94 9 : let expired = expired_objects(&keys, args.keep_latest, args.keep_stable);
95 9 : log::info!("delete objects: {:?}", expired);
96 :
97 9 : if !args.dry_run {
98 8 : for chunk in expired.chunks(1000) {
99 8 : let objects = chunk
100 8 : .iter()
101 3516 : .map(|key| aws_sdk_s3::types::ObjectIdentifier::builder().key(key).build())
102 8 : .collect::<Result<Vec<_>, _>>()?;
103 8 : let objects = aws_sdk_s3::types::Delete::builder().set_objects(Some(objects)).build()?;
104 8 : let result = client
105 8 : .delete_objects()
106 8 : .bucket("dashql-get")
107 8 : .delete(objects)
108 8 : .send()
109 8 : .await?;
110 7 : if !result.errors().is_empty() {
111 1 : return Err(anyhow::anyhow!("R2 failed to delete objects: {:?}", result.errors()));
112 6 : }
113 6 : log::info!("deleted {} objects", chunk.len());
114 : }
115 2 : }
116 :
117 7 : Ok(())
118 10 : }
119 :
120 : #[cfg(test)]
121 : mod tests {
122 : use super::*;
123 : use crate::test_s3::{FakeS3, ENV_LOCK};
124 :
125 : const R2_ENV: [&str; 3] = ["DASHQL_GET_R2_ENDPOINT", "DASHQL_GET_R2_ACCESS_KEY_ID", "DASHQL_GET_R2_SECRET_ACCESS_KEY"];
126 :
127 7 : fn keys(values: &[&str]) -> Vec<String> {
128 17 : values.iter().map(|value| value.to_string()).collect()
129 7 : }
130 :
131 : /// The objects of the versions, one metadata file and one artifact per version
132 16 : fn objects_of(versions: &[&str]) -> Vec<String> {
133 16 : versions.iter()
134 34 : .flat_map(|version| [format!("releases/{version}/metadata.json"), format!("releases/{version}/macos/arm64/DashQL.dmg")])
135 16 : .collect()
136 16 : }
137 :
138 11 : fn args(keep_latest: usize, keep_stable: usize, dry_run: bool) -> VacuumArgs {
139 11 : VacuumArgs { dry_run, keep_latest, keep_stable }
140 11 : }
141 :
142 11 : fn sorted(mut values: Vec<String>) -> Vec<String> {
143 11 : values.sort();
144 11 : values
145 11 : }
146 :
147 : #[test]
148 1 : fn retention_keeps_newest_versions() {
149 1 : let versions = ["1.0.0-dev.1", "1.0.0-dev.3", "1.0.0-dev.2"]
150 3 : .into_iter().map(|value| Version::parse(value).unwrap()).collect();
151 1 : let (keep, delete) = retained_versions(versions, 2);
152 1 : assert_eq!(keep, vec![Version::parse("1.0.0-dev.3").unwrap(), Version::parse("1.0.0-dev.2").unwrap()]);
153 1 : assert_eq!(delete, vec![Version::parse("1.0.0-dev.1").unwrap()]);
154 1 : }
155 :
156 : #[test]
157 1 : fn retention_compares_versions_as_numbers_not_as_text() {
158 3 : let versions = ["1.9.0", "1.10.0", "1.2.0"].into_iter().map(|value| Version::parse(value).unwrap()).collect();
159 1 : let (keep, delete) = retained_versions(versions, 1);
160 1 : assert_eq!(keep, vec![Version::parse("1.10.0").unwrap()]);
161 1 : assert_eq!(delete.len(), 2);
162 :
163 1 : let expired = expired_objects(&keys(&["releases/1.0.0-dev.9/a", "releases/1.0.0-dev.10/a"]), 1, 1);
164 1 : assert_eq!(expired, keys(&["releases/1.0.0-dev.9/a"]));
165 1 : }
166 :
167 : #[test]
168 1 : fn retention_keeps_everything_with_fewer_versions_than_allowed() {
169 1 : let (keep, delete) = retained_versions(vec![Version::new(1, 0, 0)], 10);
170 1 : assert_eq!(keep.len(), 1);
171 1 : assert!(delete.is_empty());
172 1 : }
173 :
174 : #[test]
175 1 : fn expires_every_object_of_a_version() {
176 1 : let expired = expired_objects(&objects_of(&["1.0.0", "2.0.0"]), 10, 1);
177 1 : assert_eq!(sorted(expired), sorted(objects_of(&["1.0.0"])));
178 1 : }
179 :
180 : #[test]
181 1 : fn retains_stable_and_prerelease_versions_separately() {
182 : // The canaries are newer than every stable version, they must not push the stable versions out
183 1 : let objects = objects_of(&["1.0.0", "1.1.0", "1.1.1-dev.1", "1.1.1-dev.2", "1.1.1-dev.3"]);
184 1 : assert_eq!(sorted(expired_objects(&objects, 2, 2)), sorted(objects_of(&["1.1.1-dev.1"])));
185 1 : assert_eq!(sorted(expired_objects(&objects, 3, 1)), sorted(objects_of(&["1.0.0"])));
186 1 : }
187 :
188 : #[test]
189 1 : fn never_expires_keys_without_a_version() {
190 1 : let objects = keys(&["releases/notes.txt", "releases/latest/metadata.json", "releases/1.0/metadata.json", "stable.json"]);
191 1 : assert_eq!(expired_objects(&objects, 0, 0), Vec::<String>::new());
192 1 : }
193 :
194 : #[test]
195 1 : fn parses_the_arguments_of_the_publish_workflow() {
196 1 : let parsed = VacuumArgs::try_parse_from(["vacuum", "--keep-stable", "100", "--keep-latest", "10"]).unwrap();
197 1 : assert!(!parsed.dry_run);
198 1 : assert_eq!((parsed.keep_stable, parsed.keep_latest), (100, 10));
199 1 : }
200 :
201 : #[test]
202 1 : fn dry_run_is_a_flag_and_retention_has_defaults() {
203 1 : let parsed = VacuumArgs::try_parse_from(["vacuum", "--dry-run"]).unwrap();
204 1 : assert!(parsed.dry_run);
205 1 : assert_eq!((parsed.keep_stable, parsed.keep_latest), (100, 10));
206 1 : assert!(!VacuumArgs::try_parse_from(["vacuum"]).unwrap().dry_run);
207 1 : }
208 :
209 : #[tokio::test]
210 1 : async fn deletes_only_the_expired_objects_of_the_bucket() {
211 1 : let s3 = FakeS3::start().await;
212 1 : s3.put_objects(objects_of(&["1.0.0", "2.0.0", "2.0.1-dev.1", "2.0.1-dev.2"]));
213 1 : s3.put_objects(keys(&["stable.json", "latest.json", "channels/stable/macos/arm64/latest-mac.yml", "releases/notes.txt"]));
214 :
215 1 : vacuum_bucket(&s3.client(), &args(1, 1, false)).await.unwrap();
216 :
217 1 : let mut expected = objects_of(&["2.0.0", "2.0.1-dev.2"]);
218 1 : expected.extend(keys(&["stable.json", "latest.json", "channels/stable/macos/arm64/latest-mac.yml", "releases/notes.txt"]));
219 1 : assert_eq!(s3.object_keys(), sorted(expected));
220 1 : }
221 :
222 : #[tokio::test]
223 1 : async fn only_lists_the_release_directories() {
224 1 : let s3 = FakeS3::start().await;
225 1 : vacuum_bucket(&s3.client(), &args(1, 1, false)).await.unwrap();
226 1 : let requests = s3.requests();
227 1 : assert_eq!(requests.len(), 1);
228 1 : assert_eq!(requests[0].query.get("prefix").map(String::as_str), Some("releases/"));
229 1 : }
230 :
231 : #[tokio::test]
232 1 : async fn dry_run_lists_but_deletes_nothing() {
233 1 : let s3 = FakeS3::start().await;
234 1 : s3.put_objects(objects_of(&["1.0.0", "2.0.0"]));
235 1 : vacuum_bucket(&s3.client(), &args(0, 0, true)).await.unwrap();
236 1 : assert_eq!(s3.object_keys(), sorted(objects_of(&["1.0.0", "2.0.0"])));
237 1 : assert!(s3.requests().iter().all(|request| request.method == "GET"));
238 1 : }
239 :
240 : #[tokio::test]
241 1 : async fn sees_the_newest_versions_on_later_pages() {
242 : // A version on the last page decides what expires, stopping early would delete the newest kept version
243 1 : let s3 = FakeS3::start().await;
244 1 : s3.set_page_size(3);
245 1 : s3.put_objects(objects_of(&["1.0.0", "1.0.1", "1.0.2", "1.0.3", "1.0.4"]));
246 1 : vacuum_bucket(&s3.client(), &args(1, 1, false)).await.unwrap();
247 1 : assert_eq!(s3.object_keys(), sorted(objects_of(&["1.0.4"])));
248 5 : assert!(s3.requests().iter().filter(|request| request.method == "GET").count() > 1);
249 1 : }
250 :
251 : #[tokio::test]
252 1 : async fn deletes_in_batches_of_a_thousand_objects() {
253 1 : let s3 = FakeS3::start().await;
254 1 : s3.set_page_size(1000);
255 2500 : let mut objects: Vec<String> = (0..2500).map(|i| format!("releases/1.0.0/file{i}")).collect();
256 1 : objects.extend(keys(&["releases/2.0.0/metadata.json"]));
257 1 : s3.put_objects(objects);
258 :
259 1 : vacuum_bucket(&s3.client(), &args(0, 1, false)).await.unwrap();
260 :
261 1 : assert_eq!(s3.object_keys(), keys(&["releases/2.0.0/metadata.json"]));
262 1 : let batches: Vec<usize> = s3.requests().iter()
263 6 : .filter(|request| request.method == "POST")
264 3 : .map(|request| String::from_utf8_lossy(&request.body).matches("<Key>").count())
265 1 : .collect();
266 1 : assert_eq!(batches, vec![1000, 1000, 500]);
267 1 : }
268 :
269 : #[tokio::test]
270 1 : async fn deletes_nothing_when_the_listing_fails() {
271 1 : let s3 = FakeS3::start().await;
272 1 : s3.put_objects(objects_of(&["1.0.0", "2.0.0"]));
273 1 : s3.deny_listing();
274 1 : assert!(vacuum_bucket(&s3.client(), &args(0, 0, false)).await.is_err());
275 1 : assert_eq!(s3.object_keys().len(), 4);
276 1 : }
277 :
278 : #[tokio::test]
279 1 : async fn reports_objects_that_could_not_be_deleted() {
280 1 : let s3 = FakeS3::start().await;
281 1 : s3.put_objects(objects_of(&["1.0.0", "2.0.0"]));
282 1 : s3.keep_object("releases/1.0.0/metadata.json");
283 1 : let error = vacuum_bucket(&s3.client(), &args(0, 1, false)).await.unwrap_err();
284 1 : assert!(error.to_string().contains("failed to delete"), "{error}");
285 1 : }
286 :
287 : #[tokio::test]
288 1 : async fn stops_at_the_first_failed_batch() {
289 1 : let s3 = FakeS3::start().await;
290 1 : s3.set_page_size(1000);
291 1500 : s3.put_objects((0..1500).map(|i| format!("releases/1.0.0/file{i:04}")).collect());
292 1 : s3.deny_deletes();
293 1 : assert!(vacuum_bucket(&s3.client(), &args(0, 0, false)).await.is_err());
294 3 : assert_eq!(s3.requests().iter().filter(|request| request.method == "POST").count(), 1);
295 1 : assert_eq!(s3.object_keys().len(), 1500);
296 1 : }
297 :
298 : #[tokio::test]
299 1 : async fn the_command_reads_the_bucket_from_the_environment() {
300 1 : let _guard = ENV_LOCK.lock().await;
301 1 : let s3 = FakeS3::start().await;
302 1 : s3.put_objects(objects_of(&["1.0.0", "2.0.0"]));
303 1 : std::env::set_var(R2_ENV[0], s3.endpoint());
304 1 : std::env::set_var(R2_ENV[1], "key");
305 1 : std::env::set_var(R2_ENV[2], "secret");
306 1 : vacuum(args(0, 1, false)).await.unwrap();
307 1 : assert_eq!(s3.object_keys(), sorted(objects_of(&["2.0.0"])));
308 1 : vacuum(args(0, 0, true)).await.unwrap();
309 1 : assert_eq!(s3.object_keys(), sorted(objects_of(&["2.0.0"])));
310 1 : }
311 :
312 : #[tokio::test]
313 1 : async fn the_command_needs_credentials() {
314 1 : let _guard = ENV_LOCK.lock().await;
315 3 : for name in R2_ENV {
316 3 : std::env::remove_var(name);
317 3 : }
318 1 : let error = vacuum(args(0, 0, false)).await.unwrap_err();
319 1 : assert!(error.to_string().contains("DASHQL_GET_R2_ENDPOINT"), "{error}");
320 1 : }
321 : }
|
