Skip to main content

max / makenotwork

10.2 KB · 285 lines History Blame Raw
1 //! Durable queue for S3 object deletions that must survive server crashes.
2
3 use crate::error::Result;
4 use sqlx::PgPool;
5 use uuid::Uuid;
6
7 /// A pending S3 deletion record.
8 #[derive(Debug, sqlx::FromRow)]
9 pub struct PendingS3Deletion {
10 pub id: Uuid,
11 pub s3_key: String,
12 pub bucket: String,
13 pub source: String,
14 pub attempts: i32,
15 }
16
17 /// Enqueue S3 keys for deletion. Each key is (s3_key, bucket).
18 #[tracing::instrument(skip_all)]
19 pub async fn enqueue_deletions<'e>(
20 executor: impl sqlx::PgExecutor<'e>,
21 keys: &[(String, String)],
22 source: &str,
23 ) -> Result<()> {
24 if keys.is_empty() {
25 return Ok(());
26 }
27 let s3_keys: Vec<&str> = keys.iter().map(|(k, _)| k.as_str()).collect();
28 let buckets: Vec<&str> = keys.iter().map(|(_, b)| b.as_str()).collect();
29 // Takes any `PgExecutor` so callers can enqueue inside a transaction,
30 // letting a row delete, its storage refund, and the S3-delete enqueue commit
31 // atomically (see `versions::delete_version`).
32 sqlx::query!(
33 "INSERT INTO pending_s3_deletions (s3_key, bucket, source) SELECT * FROM unnest($1::text[], $2::text[], $3::text[]) ON CONFLICT DO NOTHING",
34 &s3_keys as &[&str],
35 &buckets as &[&str],
36 &vec![source; keys.len()] as &[&str],
37 )
38 .execute(executor)
39 .await?;
40 Ok(())
41 }
42
43 /// Remove completed deletions by ID.
44 #[tracing::instrument(skip_all)]
45 pub async fn remove_completed(pool: &PgPool, ids: &[Uuid]) -> Result<()> {
46 if ids.is_empty() {
47 return Ok(());
48 }
49 sqlx::query!("DELETE FROM pending_s3_deletions WHERE id = ANY($1)", ids)
50 .execute(pool)
51 .await?;
52 Ok(())
53 }
54
55 /// Move permanently-failing deletions off the hot queue into the operator-visible
56 /// dead-letter table (see migration 129). Atomic per the single-statement CTE:
57 /// the rows are deleted from `pending_s3_deletions` and inserted into
58 /// `pending_s3_deletions_dead_letter` in one operation, preserving their
59 /// `created_at`/`attempts`/`last_attempted_at`. Returns the number moved.
60 ///
61 /// Without this the worker only logged and DELETEd such rows, so a malformed
62 /// key / gone bucket / ACL gap left an S3 object orphaned with no durable
63 /// record, the exact leak the dead-letter table was added to prevent.
64 #[tracing::instrument(skip_all)]
65 pub async fn move_to_dead_letter(pool: &PgPool, ids: &[Uuid]) -> Result<u64> {
66 if ids.is_empty() {
67 return Ok(0);
68 }
69 let result = sqlx::query!(
70 r#"
71 WITH moved AS (
72 DELETE FROM pending_s3_deletions
73 WHERE id = ANY($1)
74 RETURNING id, s3_key, bucket, source, created_at, attempts, last_attempted_at
75 )
76 INSERT INTO pending_s3_deletions_dead_letter
77 (id, s3_key, bucket, source, created_at, attempts, last_attempted_at)
78 SELECT id, s3_key, bucket, source, created_at, attempts, last_attempted_at
79 FROM moved
80 ON CONFLICT (id) DO NOTHING
81 "#,
82 ids,
83 )
84 .execute(pool)
85 .await?;
86 Ok(result.rows_affected())
87 }
88
89 /// One place a live S3 object key can still be referenced. [`is_s3_key_live`] is
90 /// generated from this registry, so a new s3_key-bearing table is added in
91 /// exactly one place AND declares its own bucket, it can't be checked against
92 /// the wrong bucket (the OTA-key-checked-in-the-main-branch bug this replaced)
93 /// or silently forgotten by one branch of a hand-written query. Every registered
94 /// column stores a bare s3_key, so the check is always `column = $1`, covers
95 /// gained `cover_s3_key` in migration 152, retiring the old URL-suffix match.
96 struct S3KeyRef {
97 bucket: &'static str,
98 table: &'static str,
99 column: &'static str,
100 }
101
102 /// The registry. To add a table that stores an S3 key, add a row here.
103 const S3_KEY_REFS: &[S3KeyRef] = &[
104 // ── main bucket: creator content ──
105 S3KeyRef {
106 bucket: "main",
107 table: "media_files",
108 column: "s3_key",
109 },
110 S3KeyRef {
111 bucket: "main",
112 table: "versions",
113 column: "s3_key",
114 },
115 S3KeyRef {
116 bucket: "main",
117 table: "items",
118 column: "audio_s3_key",
119 },
120 S3KeyRef {
121 bucket: "main",
122 table: "items",
123 column: "video_s3_key",
124 },
125 // ── public bucket: CDN-served, unsigned image content ──
126 // These four surfaces render straight from `cdn.makenot.work/{key}`, so after
127 // promote their content object lives in the public bucket, NOT main (see
128 // `ScanTargetKind::content_served_from_public_bucket`). Their staging object
129 // is still in main, deleted under the `main` tag by the promote, only the
130 // promoted content object is reaped against `public`. Covers store the bare
131 // key directly (migration 152 added projects.cover_s3_key; items.cover_s3_key
132 // predates it).
133 S3KeyRef {
134 bucket: "public",
135 table: "items",
136 column: "cover_s3_key",
137 },
138 S3KeyRef {
139 bucket: "public",
140 table: "projects",
141 column: "cover_s3_key",
142 },
143 S3KeyRef {
144 bucket: "public",
145 table: "item_images",
146 column: "s3_key",
147 },
148 S3KeyRef {
149 bucket: "public",
150 table: "project_images",
151 column: "s3_key",
152 },
153 // content_insertions is served PRESIGNED from the private bucket, so it stays main.
154 S3KeyRef {
155 bucket: "main",
156 table: "content_insertions",
157 column: "storage_key",
158 },
159 // ── synckit bucket: SyncKit blobs + OTA artifacts (both deterministic keys) ──
160 S3KeyRef {
161 bucket: "synckit",
162 table: "sync_blobs",
163 column: "s3_key",
164 },
165 S3KeyRef {
166 bucket: "synckit",
167 table: "ota_artifacts",
168 column: "s3_key",
169 },
170 ];
171
172 /// Returns true if any live row in `bucket` still references `s3_key`. Used by
173 /// the deletion worker to detect the delete-then-reupload race: a freshly
174 /// uploaded object reusing a queued key must not be torpedoed by the worker
175 /// draining the queue. The tables checked are [`S3_KEY_REFS`] filtered to
176 /// `bucket`, add a new key-bearing table there.
177 ///
178 /// Note: this deliberately does NOT consult `pending_uploads`. An in-flight
179 /// upload whose S3 PUT finished but whose durable row hasn't committed is
180 /// invisible here. That window is safe in practice because the durable row is
181 /// written early for every key class (OTA writes `ota_artifacts` at presign
182 /// time; same-name item/version/media replaces are caught by their confirm
183 /// handlers' idempotency guards before any delete is queued), and a queued
184 /// delete only fires after a >=10-minute staleness check. Adding
185 /// `pending_uploads` would broaden the live-set to ephemeral rows the reaper is
186 /// meant to clean, so it is intentionally excluded.
187 #[tracing::instrument(skip_all)]
188 pub async fn is_s3_key_live(pool: &PgPool, bucket: &str, s3_key: &str) -> Result<bool> {
189 let refs: Vec<&S3KeyRef> = S3_KEY_REFS.iter().filter(|r| r.bucket == bucket).collect();
190 if refs.is_empty() {
191 // Only "main"/"synckit" exist, so an unregistered bucket is unexpected.
192 // Refuse to declare the key dead, don't let the worker delete an object
193 // we have no way to verify.
194 tracing::warn!(
195 bucket,
196 "is_s3_key_live: no registered tables for bucket; treating key as live"
197 );
198 return Ok(true);
199 }
200
201 // Table/column names are compile-time constants from S3_KEY_REFS (never user
202 // input); the key value is always the bound parameter $1. Every class stores
203 // a bare key, so liveness is a flat OR of equality checks.
204 let clauses: Vec<String> = refs
205 .iter()
206 .map(|r| format!("EXISTS(SELECT 1 FROM {} WHERE {} = $1)", r.table, r.column))
207 .collect();
208 let sql = format!("SELECT {}", clauses.join(" OR "));
209
210 Ok(sqlx::query_scalar::<_, bool>(&sql)
211 .bind(s3_key)
212 .fetch_one(pool)
213 .await?)
214 }
215
216 /// Fetch stale pending deletions (older than min_age, up to limit).
217 /// Atomically increments attempt count.
218 #[tracing::instrument(skip_all)]
219 pub async fn get_stale_pending(
220 pool: &PgPool,
221 min_age: chrono::Duration,
222 limit: i64,
223 ) -> Result<Vec<PendingS3Deletion>> {
224 let cutoff = chrono::Utc::now() - min_age;
225 let rows = sqlx::query_as::<_, PendingS3Deletion>(
226 r"
227 UPDATE pending_s3_deletions
228 SET attempts = attempts + 1, last_attempted_at = NOW()
229 WHERE id IN (
230 SELECT id FROM pending_s3_deletions
231 WHERE created_at < $1
232 ORDER BY created_at
233 LIMIT $2
234 FOR UPDATE SKIP LOCKED
235 )
236 RETURNING id, s3_key, bucket, source, attempts
237 ",
238 )
239 .bind(cutoff)
240 .bind(limit)
241 .fetch_all(pool)
242 .await?;
243 Ok(rows)
244 }
245
246 #[cfg(test)]
247 mod tests {
248 use super::*;
249 use std::collections::BTreeSet;
250
251 /// Pin the set of S3 key-bearing columns the deletion-worker liveness check
252 /// consults. A new key-bearing column added anywhere must be registered in
253 /// `S3_KEY_REFS` too, or a freshly re-uploaded object reusing a queued key
254 /// could be torpedoed by the worker. This forces any change to be deliberate
255 /// (ultra-fuzz Run 10 Sto N-1).
256 #[test]
257 fn s3_key_ref_registry_covers_every_known_class() {
258 let actual: BTreeSet<(&str, &str, &str)> = S3_KEY_REFS
259 .iter()
260 .map(|r| (r.bucket, r.table, r.column))
261 .collect();
262 let expected: BTreeSet<(&str, &str, &str)> = [
263 ("main", "media_files", "s3_key"),
264 ("main", "versions", "s3_key"),
265 ("main", "items", "audio_s3_key"),
266 ("main", "items", "video_s3_key"),
267 ("main", "content_insertions", "storage_key"),
268 // CDN-served image content lives in the public bucket post-promote.
269 ("public", "items", "cover_s3_key"),
270 ("public", "projects", "cover_s3_key"),
271 ("public", "item_images", "s3_key"),
272 ("public", "project_images", "s3_key"),
273 ("synckit", "sync_blobs", "s3_key"),
274 ("synckit", "ota_artifacts", "s3_key"),
275 ]
276 .into_iter()
277 .collect();
278 assert_eq!(
279 actual, expected,
280 "S3_KEY_REFS changed: if you added or removed a key-bearing column, update this \
281 test and confirm the deletion-worker liveness check still covers every key class"
282 );
283 }
284 }
285