|
23 | 23 | MultipartUploadState, |
24 | 24 | PartMetadata, |
25 | 25 | delete_upload_state, |
| 26 | + load_multipart_metadata, |
26 | 27 | persist_upload_state, |
27 | 28 | plaintext_attr_cache, |
28 | 29 | save_multipart_metadata, |
@@ -151,107 +152,158 @@ async def handle_complete_multipart_upload( |
151 | 152 | bucket, key = self._parse_path(request.url.path) |
152 | 153 | async with self._client(creds) as client: |
153 | 154 | upload_id, _ = self._extract_multipart_params(request) |
154 | | - |
155 | | - state = await self.multipart_manager.complete_upload(bucket, key, upload_id) |
156 | | - if not state: |
157 | | - state = await self._recover_upload_state( |
158 | | - client, bucket, key, upload_id, context="for complete" |
| 155 | + async with self.complete_upload_lock.hold(bucket, key, upload_id): |
| 156 | + return await self._handle_complete_multipart_upload_locked( |
| 157 | + request, creds, client, bucket, key, upload_id |
159 | 158 | ) |
160 | 159 |
|
161 | | - if state.deferred_copy_tail: |
162 | | - logger.info( |
163 | | - "COMPLETE_MULTIPART_DEFERRED_TAIL_PENDING", |
164 | | - bucket=bucket, |
165 | | - key=key, |
166 | | - upload_id=upload_id[:20] + "...", |
167 | | - tail_bytes=len(state.deferred_copy_tail), |
168 | | - ) |
169 | | - state = await self._flush_deferred_copy_tail_for_complete( |
170 | | - client, bucket, key, upload_id, state |
171 | | - ) |
172 | | - |
173 | | - # Parse client's part list |
174 | | - body = await request.body() |
175 | | - client_parts = self._parse_client_parts(body) |
| 160 | + async def _handle_complete_multipart_upload_locked( |
| 161 | + self, |
| 162 | + request: Request, |
| 163 | + creds: S3Credentials, |
| 164 | + client: S3Client, |
| 165 | + bucket: str, |
| 166 | + key: str, |
| 167 | + upload_id: str, |
| 168 | + ) -> Response: |
| 169 | + idempotent = await self._try_idempotent_complete_response(client, bucket, key, upload_id) |
| 170 | + if idempotent is not None: |
| 171 | + return idempotent |
176 | 172 |
|
177 | | - # Build S3 parts list |
178 | | - s3_parts, completed_parts, total_plaintext = self._build_s3_parts( |
179 | | - client_parts, state, bucket, key, upload_id |
| 173 | + state = await self.multipart_manager.complete_upload(bucket, key, upload_id) |
| 174 | + if not state: |
| 175 | + state = await self._recover_upload_state( |
| 176 | + client, bucket, key, upload_id, context="for complete" |
180 | 177 | ) |
181 | 178 |
|
| 179 | + if state.deferred_copy_tail: |
182 | 180 | logger.info( |
183 | | - "COMPLETE_MULTIPART", |
| 181 | + "COMPLETE_MULTIPART_DEFERRED_TAIL_PENDING", |
184 | 182 | bucket=bucket, |
185 | 183 | key=key, |
186 | 184 | upload_id=upload_id[:20] + "...", |
187 | | - client_parts=len(completed_parts), |
188 | | - s3_parts=len(s3_parts), |
189 | | - total_mb=f"{total_plaintext / 1024 / 1024:.2f}MB", |
| 185 | + tail_bytes=len(state.deferred_copy_tail), |
| 186 | + ) |
| 187 | + state = await self._flush_deferred_copy_tail_for_complete( |
| 188 | + client, bucket, key, upload_id, state |
190 | 189 | ) |
191 | 190 |
|
192 | | - # Complete in S3 |
193 | | - try: |
194 | | - complete_resp = await self._complete_multipart_upload_with_retry( |
195 | | - client, bucket, key, upload_id, s3_parts, completed_parts |
196 | | - ) |
197 | | - except ClientError as e: |
198 | | - await self._handle_complete_error( |
199 | | - e, client, bucket, key, upload_id, s3_parts, completed_parts, total_plaintext |
200 | | - ) |
201 | | - else: |
202 | | - plaintext_attr_cache.put( |
203 | | - bucket, |
204 | | - key, |
205 | | - str(complete_resp.get("ETag", "")).strip('"'), |
206 | | - total_plaintext, |
207 | | - synthetic_multipart_etag(total_plaintext), |
208 | | - ) |
| 191 | + # Parse client's part list |
| 192 | + body = await request.body() |
| 193 | + client_parts = self._parse_client_parts(body) |
209 | 194 |
|
210 | | - # Save metadata first, then delete state. |
211 | | - # Order matters: if metadata save fails, state is preserved |
212 | | - # so the upload can be retried. Deleting state first would |
213 | | - # lose the DEK, making the object permanently undecryptable. |
214 | | - # Prefer the kid recorded when the upload was created; if the state |
215 | | - # predates it (e.g. older recovered state), fall back to the |
216 | | - # completing credential's key. |
217 | | - if state.kid: |
218 | | - kid, kek = state.kid, self.keyring.key_by_id(state.kid) |
219 | | - else: |
220 | | - kid, kek = self.keyring.key_for(creds.access_key) |
221 | | - wrapped_dek = crypto.wrap_key(state.dek, kek) |
222 | | - await save_multipart_metadata( |
223 | | - client, |
| 195 | + # Build S3 parts list |
| 196 | + s3_parts, completed_parts, total_plaintext = self._build_s3_parts( |
| 197 | + client_parts, state, bucket, key, upload_id |
| 198 | + ) |
| 199 | + |
| 200 | + logger.info( |
| 201 | + "COMPLETE_MULTIPART", |
| 202 | + bucket=bucket, |
| 203 | + key=key, |
| 204 | + upload_id=upload_id[:20] + "...", |
| 205 | + client_parts=len(completed_parts), |
| 206 | + s3_parts=len(s3_parts), |
| 207 | + total_mb=f"{total_plaintext / 1024 / 1024:.2f}MB", |
| 208 | + ) |
| 209 | + |
| 210 | + # Complete in S3 |
| 211 | + try: |
| 212 | + complete_resp = await self._complete_multipart_upload_with_retry( |
| 213 | + client, bucket, key, upload_id, s3_parts, completed_parts |
| 214 | + ) |
| 215 | + except ClientError as e: |
| 216 | + await self._handle_complete_error( |
| 217 | + e, client, bucket, key, upload_id, s3_parts, completed_parts, total_plaintext |
| 218 | + ) |
| 219 | + else: |
| 220 | + plaintext_attr_cache.put( |
224 | 221 | bucket, |
225 | 222 | key, |
226 | | - MultipartMetadata( |
227 | | - version=2, |
228 | | - part_count=len(completed_parts), |
229 | | - total_plaintext_size=total_plaintext, |
230 | | - parts=completed_parts, |
231 | | - wrapped_dek=wrapped_dek, |
232 | | - kid=kid, |
233 | | - ), |
| 223 | + str(complete_resp.get("ETag", "")).strip('"'), |
| 224 | + total_plaintext, |
| 225 | + synthetic_multipart_etag(total_plaintext), |
234 | 226 | ) |
235 | | - await delete_upload_state(client, bucket, key, upload_id) |
236 | 227 |
|
237 | | - logger.info( |
238 | | - "COMPLETE_MULTIPART_SUCCESS", |
239 | | - bucket=bucket, |
240 | | - key=key, |
241 | | - upload_id=upload_id[:20] + "...", |
242 | | - total_parts=len(completed_parts), |
243 | | - total_mb=f"{total_plaintext / 1024 / 1024:.2f}MB", |
244 | | - ) |
| 228 | + # Save metadata first, then delete state. |
| 229 | + # Order matters: if metadata save fails, state is preserved |
| 230 | + # so the upload can be retried. Deleting state first would |
| 231 | + # lose the DEK, making the object permanently undecryptable. |
| 232 | + # Prefer the kid recorded when the upload was created; if the state |
| 233 | + # predates it (e.g. older recovered state), fall back to the |
| 234 | + # completing credential's key. |
| 235 | + if state.kid: |
| 236 | + kid, kek = state.kid, self.keyring.key_by_id(state.kid) |
| 237 | + else: |
| 238 | + kid, kek = self.keyring.key_for(creds.access_key) |
| 239 | + wrapped_dek = crypto.wrap_key(state.dek, kek) |
| 240 | + await save_multipart_metadata( |
| 241 | + client, |
| 242 | + bucket, |
| 243 | + key, |
| 244 | + MultipartMetadata( |
| 245 | + version=2, |
| 246 | + part_count=len(completed_parts), |
| 247 | + total_plaintext_size=total_plaintext, |
| 248 | + parts=completed_parts, |
| 249 | + wrapped_dek=wrapped_dek, |
| 250 | + kid=kid, |
| 251 | + ), |
| 252 | + ) |
| 253 | + await delete_upload_state(client, bucket, key, upload_id) |
245 | 254 |
|
246 | | - location = f"{self.settings.s3_endpoint}/{bucket}/{key}" |
247 | | - etag = hashlib.md5( |
248 | | - str(state.total_plaintext_size).encode(), usedforsecurity=False |
249 | | - ).hexdigest() |
| 255 | + logger.info( |
| 256 | + "COMPLETE_MULTIPART_SUCCESS", |
| 257 | + bucket=bucket, |
| 258 | + key=key, |
| 259 | + upload_id=upload_id[:20] + "...", |
| 260 | + total_parts=len(completed_parts), |
| 261 | + total_mb=f"{total_plaintext / 1024 / 1024:.2f}MB", |
| 262 | + ) |
250 | 263 |
|
251 | | - return Response( |
252 | | - content=xml_responses.complete_multipart(location, bucket, key, etag), |
253 | | - media_type="application/xml", |
254 | | - ) |
| 264 | + location = f"{self.settings.s3_endpoint}/{bucket}/{key}" |
| 265 | + etag = hashlib.md5( |
| 266 | + str(state.total_plaintext_size).encode(), usedforsecurity=False |
| 267 | + ).hexdigest() |
| 268 | + |
| 269 | + return Response( |
| 270 | + content=xml_responses.complete_multipart(location, bucket, key, etag), |
| 271 | + media_type="application/xml", |
| 272 | + ) |
| 273 | + |
| 274 | + async def _try_idempotent_complete_response( |
| 275 | + self, client: S3Client, bucket: str, key: str, upload_id: str |
| 276 | + ) -> Response | None: |
| 277 | + """Return success if a peer pod already finished this upload.""" |
| 278 | + meta = await load_multipart_metadata(client, bucket, key) |
| 279 | + if meta is None: |
| 280 | + return None |
| 281 | + |
| 282 | + try: |
| 283 | + head = await client.head_object(bucket, key) |
| 284 | + except ClientError: |
| 285 | + return None |
| 286 | + |
| 287 | + expected_ciphertext_size = sum(p.ciphertext_size for p in meta.parts) |
| 288 | + if head.get("ContentLength") != expected_ciphertext_size: |
| 289 | + return None |
| 290 | + |
| 291 | + logger.info( |
| 292 | + "COMPLETE_MULTIPART_IDEMPOTENT", |
| 293 | + bucket=bucket, |
| 294 | + key=key, |
| 295 | + upload_id=upload_id[:20] + "...", |
| 296 | + total_mb=f"{meta.total_plaintext_size / 1024 / 1024:.2f}MB", |
| 297 | + ) |
| 298 | + |
| 299 | + location = f"{self.settings.s3_endpoint}/{bucket}/{key}" |
| 300 | + etag = hashlib.md5( |
| 301 | + str(meta.total_plaintext_size).encode(), usedforsecurity=False |
| 302 | + ).hexdigest() |
| 303 | + return Response( |
| 304 | + content=xml_responses.complete_multipart(location, bucket, key, etag), |
| 305 | + media_type="application/xml", |
| 306 | + ) |
255 | 307 |
|
256 | 308 | def _parse_client_parts(self, body: bytes) -> list[dict]: |
257 | 309 | client_parts = [] |
|
0 commit comments