Skip to main content

max / makenotwork

7.8 KB · 212 lines History Blame Raw
1 //! Tracking for presigned uploads that have not yet been confirmed.
2
3 use sqlx::PgPool;
4
5 use crate::db::UserId;
6 use crate::error::Result;
7
8 /// Record that a presigned upload URL was issued.
9 pub(crate) async fn record_pending_upload(
10 pool: &PgPool,
11 user_id: UserId,
12 s3_key: &str,
13 bucket: &str,
14 ) -> Result<()> {
15 // Re-presigning the same key (idempotent retry, multi-part flow) must refresh
16 // `created_at`, otherwise the stale-pending reaper can delete a freshly-pending
17 // object that's actively being uploaded right now.
18 //
19 // Pin the refresh to the original owner: if a different user collides on the
20 // same key, do NOT refresh, let the reaper age the original row out on its
21 // own schedule. Otherwise a re-presign loop by another principal could keep
22 // an orphan object alive indefinitely.
23 sqlx::query!(
24 "INSERT INTO pending_uploads (user_id, s3_key, bucket) VALUES ($1, $2, $3)
25 ON CONFLICT (s3_key, bucket) DO UPDATE SET created_at = NOW()
26 WHERE pending_uploads.user_id = EXCLUDED.user_id",
27 user_id as UserId,
28 s3_key,
29 bucket,
30 )
31 .execute(pool)
32 .await?;
33 Ok(())
34 }
35
36 /// Record a pending multipart upload, persisting the size the caller declared at
37 /// `start`. `parts` reads it back (see [`declared_size`]) to bind the part
38 /// geometry to the tier-checked size instead of trusting a later request body.
39 pub(crate) async fn record_pending_multipart_upload(
40 pool: &PgPool,
41 user_id: UserId,
42 s3_key: &str,
43 bucket: &str,
44 declared_size_bytes: i64,
45 ) -> Result<()> {
46 // Same owner-pinned created_at refresh as record_pending_upload; also refresh
47 // the declared size so a re-issued `start` for the same key stays consistent.
48 sqlx::query!(
49 "INSERT INTO pending_uploads (user_id, s3_key, bucket, declared_size_bytes)
50 VALUES ($1, $2, $3, $4)
51 ON CONFLICT (s3_key, bucket) DO UPDATE
52 SET created_at = NOW(), declared_size_bytes = EXCLUDED.declared_size_bytes
53 WHERE pending_uploads.user_id = EXCLUDED.user_id",
54 user_id as UserId,
55 s3_key,
56 bucket,
57 declared_size_bytes,
58 )
59 .execute(pool)
60 .await?;
61 Ok(())
62 }
63
64 /// The size the caller declared when it opened this pending upload, or `None`
65 /// if the row does not exist or carries no declared size (a single-PUT presign).
66 /// Scoped to `(user_id, s3_key, bucket)`, the same ownership proof the rest of
67 /// this module uses.
68 pub(crate) async fn declared_size(
69 pool: &PgPool,
70 user_id: UserId,
71 s3_key: &str,
72 bucket: &str,
73 ) -> Result<Option<i64>> {
74 let row = sqlx::query!(
75 "SELECT declared_size_bytes FROM pending_uploads
76 WHERE s3_key = $1 AND user_id = $2 AND bucket = $3",
77 s3_key,
78 user_id as UserId,
79 bucket,
80 )
81 .fetch_optional(pool)
82 .await?;
83 Ok(row.and_then(|r| r.declared_size_bytes))
84 }
85
86 /// Refresh a pending upload's `created_at` so an actively-progressing transfer
87 /// (e.g. a multipart session still requesting part URLs) is not aged out by the
88 /// stale-pending reaper mid-flight.
89 pub(crate) async fn touch_pending_upload(
90 pool: &PgPool,
91 user_id: UserId,
92 s3_key: &str,
93 bucket: &str,
94 ) -> Result<()> {
95 sqlx::query!(
96 "UPDATE pending_uploads SET created_at = NOW()
97 WHERE s3_key = $1 AND user_id = $2 AND bucket = $3",
98 s3_key,
99 user_id as UserId,
100 bucket,
101 )
102 .execute(pool)
103 .await?;
104 Ok(())
105 }
106
107 /// Remove the pending upload record after a successful confirm. Scoped to
108 /// `user_id` so a future caller that accepts a partially user-supplied key
109 /// can't delete another user's pending row, today's per-handler prefix
110 /// validation makes cross-user collision unreachable, but the function
111 /// signature shouldn't be broader than the invariant it protects.
112 ///
113 /// Also scoped to `bucket`: the unique key is `(s3_key, bucket)` (migration
114 /// 137), so a key present in both the main and synckit buckets has two distinct
115 /// rows. Matching on `s3_key` alone would let one bucket's confirm clear the
116 /// other bucket's pending record (Run #18 Storage B7).
117 pub(crate) async fn remove_pending_upload<'e>(
118 executor: impl sqlx::PgExecutor<'e>,
119 user_id: UserId,
120 s3_key: &str,
121 bucket: &str,
122 ) -> Result<()> {
123 sqlx::query!(
124 "DELETE FROM pending_uploads WHERE s3_key = $1 AND user_id = $2 AND bucket = $3",
125 s3_key,
126 user_id as UserId,
127 bucket,
128 )
129 .execute(executor)
130 .await?;
131 Ok(())
132 }
133
134 /// Whether `s3_key` (in `bucket`) is a pending upload this user presigned.
135 ///
136 /// The confirm handlers used to prove ownership by checking the client-supplied
137 /// key started with the user's `{user_id}/{item_id}/...` prefix. With scan-then-
138 /// promote the presigned key is an owner-less `staging/{uuid}.{ext}`, so that
139 /// structural check no longer binds the key to a user. Every presign records the
140 /// key here against its owner (`record_pending_upload`); confirm now proves
141 /// ownership by looking it up, a caller cannot confirm a staging key it did not
142 /// presign (and cannot guess another user's random staging uuid).
143 pub(crate) async fn is_owned(
144 pool: &PgPool,
145 user_id: UserId,
146 s3_key: &str,
147 bucket: &str,
148 ) -> Result<bool> {
149 let owned = sqlx::query_scalar::<_, bool>(
150 "SELECT EXISTS (SELECT 1 FROM pending_uploads WHERE s3_key = $1 AND user_id = $2 AND bucket = $3)",
151 )
152 .bind(s3_key)
153 .bind(user_id)
154 .bind(bucket)
155 .fetch_one(pool)
156 .await?;
157 Ok(owned)
158 }
159
160 /// Per-tick cap on the orphan-upload reaper. The reaper runs every scheduler
161 /// tick under the tick-wide advisory lock and deletes serially (one S3 round-
162 /// trip per row), so an unbounded result set lets a backlog wedge the tick
163 /// (PERF-S1, Run #23). Oldest-first + a bound drains across ticks instead.
164 pub(crate) const STALE_UPLOAD_BATCH: i64 = 200;
165
166 /// Fetch up to [`STALE_UPLOAD_BATCH`] presigned uploads older than `max_age`
167 /// that were never confirmed, oldest first.
168 pub(crate) async fn get_stale_pending_uploads(
169 pool: &PgPool,
170 max_age: chrono::Duration,
171 ) -> Result<Vec<(String, String)>> {
172 let cutoff = chrono::Utc::now() - max_age;
173 // runtime-checked: binds a chrono `DateTime<Utc>` (`$1`). With sqlx's `time`
174 // and `chrono` features unified (the session store pulls `time`), the macro
175 // infers the TIMESTAMPTZ bind as `time::OffsetDateTime`, which a chrono value
176 // won't satisfy, and a bind parameter's type can't be overridden in the macro
177 // (only output columns can). Mirrors mt-db's one chrono-binding write path.
178 let rows: Vec<(String, String)> = sqlx::query_as(
179 "SELECT s3_key, bucket FROM pending_uploads WHERE created_at < $1 \
180 ORDER BY created_at LIMIT $2",
181 )
182 .bind(cutoff)
183 .bind(STALE_UPLOAD_BATCH)
184 .fetch_all(pool)
185 .await?;
186 Ok(rows)
187 }
188
189 /// Bulk-delete pending upload records by `(s3_key, bucket)` pair. Bucket-scoped
190 /// to match the `(s3_key, bucket)` uniqueness (migration 137): the stale reaper
191 /// processes per-bucket, so it must clear only the row for the bucket it acted
192 /// on, not every bucket that happens to share the key (Run #18 Storage B7).
193 pub(crate) async fn delete_pending_uploads(pool: &PgPool, keys: &[(String, String)]) -> Result<()> {
194 if keys.is_empty() {
195 return Ok(());
196 }
197 let s3_keys: Vec<&str> = keys.iter().map(|(k, _)| k.as_str()).collect();
198 let buckets: Vec<&str> = keys.iter().map(|(_, b)| b.as_str()).collect();
199 // Match each (key, bucket) pair positionally via UNNEST so a key in two
200 // buckets only clears the rows whose bucket was actually reaped.
201 sqlx::query!(
202 "DELETE FROM pending_uploads pu
203 USING UNNEST($1::text[], $2::text[]) AS t(s3_key, bucket)
204 WHERE pu.s3_key = t.s3_key AND pu.bucket = t.bucket",
205 &s3_keys as &[&str],
206 &buckets as &[&str],
207 )
208 .execute(pool)
209 .await?;
210 Ok(())
211 }
212