Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/worker_communication/cleanup_queue/memory.py: 21%
321 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-10-07 02:04 +0000
1from __future__ import annotations
3import asyncio
4from collections.abc import Iterable, Mapping
5from copy import deepcopy
6from dataclasses import dataclass
7from datetime import timedelta
8from secrets import token_urlsafe
9from typing import Any, TypeVar
10from uuid import UUID
12from prefect.client.schemas.worker_channel import CleanupKind
13from prefect.server.worker_communication.cleanup_queue import (
14 CleanupQueueDeadLetter,
15 CleanupQueueLeaseExpiryResult,
16 CleanupQueueMessage,
17 CleanupQueueOperation,
18 CleanupQueueOperationResult,
19 CleanupQueueReservation,
20 CleanupQueueWakeup,
21 record_cleanup_queue_dead_letter,
22 record_cleanup_queue_lease_expiry_result,
23 record_cleanup_queue_operation,
24)
25from prefect.server.worker_communication.cleanup_queue import (
26 WorkerCleanupQueue as _WorkerCleanupQueue,
27)
28from prefect.settings.context import get_current_settings
29from prefect.types import DateTime
30from prefect.types._datetime import now
32_T = TypeVar("_T", bound=CleanupQueueMessage | CleanupQueueDeadLetter)
35@dataclass(frozen=True)
36class _Reservation:
37 token: str
38 lease_expires_at: DateTime
39 reserved_at: DateTime
42@dataclass(frozen=True)
43class _AckedMessage:
44 message: CleanupQueueMessage
45 completed_at: DateTime
48@dataclass(frozen=True)
49class _QueuePolicy:
50 lease_duration: timedelta
51 max_delivery_attempts: int
52 completed_idempotency_retention: timedelta | None
55def _copy_model(model: _T) -> _T:
56 return model.model_copy(deep=True)
59class WorkerCleanupQueue(_WorkerCleanupQueue):
60 """
61 Singleton in-memory cleanup queue storage for a single server process.
62 """
64 _DEFAULT_EXPIRE_LEASE_LIMIT = 100
65 _instance: "WorkerCleanupQueue | None" = None
66 _initialized = False
68 def __new__(cls) -> "WorkerCleanupQueue":
69 if cls._instance is None: 69 ↛ 71line 69 didn't jump to line 71 because the condition on line 69 was always true
70 cls._instance = super().__new__(cls)
71 return cls._instance
73 def __init__(self) -> None:
74 if self.__class__._initialized: 74 ↛ 75line 74 didn't jump to line 75 because the condition on line 74 was never true
75 return
77 self._messages: dict[UUID, CleanupQueueMessage] = {}
78 self._reservations: dict[UUID, _Reservation] = {}
79 self._dead_letters: dict[UUID, CleanupQueueDeadLetter] = {}
80 self._acked_messages: dict[UUID, _AckedMessage] = {}
81 self._idempotency_keys: dict[tuple[UUID, str], UUID] = {}
82 self._wakeup_sequences: dict[UUID, int] = {}
83 self._lock = asyncio.Lock()
84 self._condition = asyncio.Condition()
85 self.__class__._initialized = True
87 def clear(self) -> None:
88 self._messages.clear()
89 self._reservations.clear()
90 self._dead_letters.clear()
91 self._acked_messages.clear()
92 self._idempotency_keys.clear()
93 self._wakeup_sequences.clear()
94 self._lock = asyncio.Lock()
95 self._condition = asyncio.Condition()
97 async def enqueue(
98 self,
99 *,
100 message_id: UUID,
101 idempotency_key: str,
102 work_pool_id: UUID,
103 kind: CleanupKind,
104 target: Mapping[str, Any],
105 data: Mapping[str, Any] | None = None,
106 work_queue_id: UUID | None = None,
107 ) -> CleanupQueueMessage:
108 if not idempotency_key:
109 raise ValueError("idempotency_key must be non-empty")
111 should_wake = False
112 policy = self._policy()
113 async with self._lock:
114 current_time = now("UTC")
115 self._prune_completed_idempotency_locked(
116 current_time=current_time, policy=policy
117 )
118 existing = self._read_existing_message_locked(
119 message_id=message_id,
120 idempotency_key=idempotency_key,
121 work_pool_id=work_pool_id,
122 )
123 if existing is not None:
124 should_wake = (
125 existing.message_id in self._messages
126 and existing.message_id not in self._reservations
127 )
128 result = _copy_model(existing)
129 record_cleanup_queue_operation(
130 "enqueue",
131 status="duplicate",
132 work_pool_id=work_pool_id,
133 message_id=message_id,
134 cleanup_kind=str(kind),
135 )
136 else:
137 message = CleanupQueueMessage(
138 message_id=message_id,
139 idempotency_key=idempotency_key,
140 work_pool_id=work_pool_id,
141 work_queue_id=work_queue_id,
142 kind=kind,
143 target=deepcopy(dict(target)),
144 data=deepcopy(dict(data or {})),
145 created_at=current_time,
146 updated_at=current_time,
147 )
148 self._messages[message_id] = message
149 self._idempotency_keys[(work_pool_id, idempotency_key)] = message_id
150 result = _copy_model(message)
151 should_wake = True
152 record_cleanup_queue_operation(
153 "enqueue",
154 status="accepted",
155 work_pool_id=work_pool_id,
156 message_id=message_id,
157 cleanup_kind=str(kind),
158 )
160 if should_wake:
161 await self.wake_dispatchers(work_pool_id)
162 return result
164 async def reserve(
165 self,
166 *,
167 work_pool_id: UUID,
168 cleanup_kinds: Iterable[CleanupKind] | None = None,
169 preferred_work_queue_ids: Iterable[UUID] | None = None,
170 allow_fallback_to_any_queue: bool = True,
171 ) -> CleanupQueueReservation | None:
172 policy = self._policy()
173 cleanup_kind_filter = set(cleanup_kinds) if cleanup_kinds is not None else None
174 queue_preference_passes = self._queue_preference_passes(
175 preferred_work_queue_ids=preferred_work_queue_ids,
176 allow_fallback_to_any_queue=allow_fallback_to_any_queue,
177 )
179 async with self._lock:
180 current_time = now("UTC")
181 expiry_result = self._expire_due_leases_locked(
182 current_time=current_time,
183 policy=policy,
184 limit=self._DEFAULT_EXPIRE_LEASE_LIMIT,
185 work_pool_id=work_pool_id,
186 )
187 wake_work_pool_ids = {
188 message.work_pool_id for message in expiry_result.redelivered
189 }
191 # Queue affinity is advisory: reserve preferred queue work first, then
192 # fall back to any eligible work in the pool when fallback is enabled.
193 result: CleanupQueueReservation | None = None
194 for queue_filter in queue_preference_passes:
195 for message in tuple(self._messages.values()):
196 if message.work_pool_id != work_pool_id:
197 continue
198 if message.message_id in self._reservations:
199 continue
200 if (
201 cleanup_kind_filter is not None
202 and message.kind not in cleanup_kind_filter
203 ):
204 continue
205 if (
206 queue_filter is not None
207 and message.work_queue_id not in queue_filter
208 ):
209 continue
210 if message.delivery_count >= policy.max_delivery_attempts:
211 self._move_to_dead_letter_locked(
212 message_id=message.message_id,
213 reason="max_delivery_attempts_reached",
214 current_time=current_time,
215 )
216 continue
218 lease_expires_at = current_time + policy.lease_duration
219 updated_message = message.model_copy(
220 update={
221 "delivery_count": message.delivery_count + 1,
222 "updated_at": current_time,
223 }
224 )
225 self._messages[message.message_id] = updated_message
226 reservation = _Reservation(
227 token=token_urlsafe(32),
228 lease_expires_at=lease_expires_at,
229 reserved_at=current_time,
230 )
231 self._reservations[message.message_id] = reservation
232 result = CleanupQueueReservation(
233 **updated_message.model_dump(),
234 reservation_token=reservation.token,
235 lease_expires_at=lease_expires_at,
236 )
237 break
238 if result is not None:
239 break
241 for wake_work_pool_id in wake_work_pool_ids:
242 await self.wake_dispatchers(wake_work_pool_id)
243 record_cleanup_queue_lease_expiry_result(expiry_result)
245 if result is not None:
246 record_cleanup_queue_operation(
247 "reserve",
248 status="accepted",
249 work_pool_id=work_pool_id,
250 message_id=result.message_id,
251 cleanup_kind=str(result.kind),
252 )
254 return result
256 async def ack(
257 self,
258 *,
259 work_pool_id: UUID,
260 message_id: UUID,
261 reservation_token: str,
262 ) -> CleanupQueueOperationResult:
263 policy = self._policy()
264 result: CleanupQueueOperationResult
265 wake_work_pool_id: UUID | None = None
266 async with self._lock:
267 current_time = now("UTC")
268 operation_result = self._validate_current_reservation_locked(
269 operation="ack",
270 work_pool_id=work_pool_id,
271 message_id=message_id,
272 reservation_token=reservation_token,
273 current_time=current_time,
274 policy=policy,
275 )
276 if operation_result is not None:
277 result = operation_result
278 wake_work_pool_id = self._expired_result_wakeup_work_pool_id_locked(
279 message_id=message_id,
280 result=result,
281 )
282 else:
283 message = self._messages.pop(message_id)
284 self._reservations.pop(message_id, None)
285 self._acked_messages[message_id] = _AckedMessage(
286 message=message.model_copy(update={"updated_at": current_time}),
287 completed_at=current_time,
288 )
289 result = CleanupQueueOperationResult(
290 message_id=message_id,
291 operation="ack",
292 status="accepted",
293 )
295 record_cleanup_queue_operation(
296 "ack",
297 status=result.status,
298 work_pool_id=work_pool_id,
299 message_id=message_id,
300 )
301 if wake_work_pool_id is not None:
302 await self.wake_dispatchers(wake_work_pool_id)
303 return result
305 async def release(
306 self,
307 *,
308 work_pool_id: UUID,
309 message_id: UUID,
310 reservation_token: str,
311 reason: str,
312 ) -> CleanupQueueOperationResult:
313 if not reason:
314 raise ValueError("release reason must be non-empty")
316 policy = self._policy()
317 result: CleanupQueueOperationResult
318 wake_work_pool_id: UUID | None = None
319 async with self._lock:
320 current_time = now("UTC")
321 operation_result = self._validate_current_reservation_locked(
322 operation="release",
323 work_pool_id=work_pool_id,
324 message_id=message_id,
325 reservation_token=reservation_token,
326 current_time=current_time,
327 policy=policy,
328 )
329 if operation_result is not None:
330 result = operation_result
331 wake_work_pool_id = self._expired_result_wakeup_work_pool_id_locked(
332 message_id=message_id,
333 result=result,
334 )
335 else:
336 message = self._messages[message_id]
337 if message.delivery_count >= policy.max_delivery_attempts:
338 dead_letter = self._move_to_dead_letter_locked(
339 message_id=message_id,
340 reason="max_delivery_attempts_reached",
341 current_time=current_time,
342 release_reason=reason,
343 )
344 result = CleanupQueueOperationResult(
345 message_id=message_id,
346 operation="release",
347 status="dead_lettered",
348 reason="max_delivery_attempts_reached",
349 dead_letter=dead_letter,
350 )
351 else:
352 self._reservations.pop(message_id, None)
353 self._messages[message_id] = message.model_copy(
354 update={"updated_at": current_time}
355 )
356 wake_work_pool_id = message.work_pool_id
357 result = CleanupQueueOperationResult(
358 message_id=message_id,
359 operation="release",
360 status="accepted",
361 )
363 record_cleanup_queue_operation(
364 "release",
365 status=result.status,
366 work_pool_id=work_pool_id,
367 message_id=message_id,
368 )
369 if wake_work_pool_id is not None:
370 await self.wake_dispatchers(wake_work_pool_id)
371 return result
373 async def renew(
374 self,
375 *,
376 work_pool_id: UUID,
377 message_id: UUID,
378 reservation_token: str,
379 ) -> CleanupQueueOperationResult:
380 policy = self._policy()
381 result: CleanupQueueOperationResult
382 wake_work_pool_id: UUID | None = None
383 async with self._lock:
384 current_time = now("UTC")
385 operation_result = self._validate_current_reservation_locked(
386 operation="renew",
387 work_pool_id=work_pool_id,
388 message_id=message_id,
389 reservation_token=reservation_token,
390 current_time=current_time,
391 policy=policy,
392 )
393 if operation_result is not None:
394 result = operation_result
395 wake_work_pool_id = self._expired_result_wakeup_work_pool_id_locked(
396 message_id=message_id,
397 result=result,
398 )
399 else:
400 lease_expires_at = current_time + policy.lease_duration
401 reservation = self._reservations[message_id]
402 self._reservations[message_id] = _Reservation(
403 token=reservation.token,
404 lease_expires_at=lease_expires_at,
405 reserved_at=reservation.reserved_at,
406 )
407 self._messages[message_id] = self._messages[message_id].model_copy(
408 update={"updated_at": current_time}
409 )
410 result = CleanupQueueOperationResult(
411 message_id=message_id,
412 operation="renew",
413 status="accepted",
414 lease_expires_at=lease_expires_at,
415 )
417 record_cleanup_queue_operation(
418 "renew",
419 status=result.status,
420 work_pool_id=work_pool_id,
421 message_id=message_id,
422 )
423 if wake_work_pool_id is not None:
424 await self.wake_dispatchers(wake_work_pool_id)
425 return result
427 async def expire_leases(
428 self,
429 *,
430 limit: int = _DEFAULT_EXPIRE_LEASE_LIMIT,
431 work_pool_id: UUID | None = None,
432 ) -> CleanupQueueLeaseExpiryResult:
433 policy = self._policy()
434 async with self._lock:
435 result = self._expire_due_leases_locked(
436 current_time=now("UTC"),
437 policy=policy,
438 limit=limit,
439 work_pool_id=work_pool_id,
440 )
442 for message in result.redelivered: 442 ↛ 443line 442 didn't jump to line 443 because the loop on line 442 never started
443 await self.wake_dispatchers(message.work_pool_id)
444 record_cleanup_queue_lease_expiry_result(result)
446 return result
448 async def read_message(
449 self, *, work_pool_id: UUID, message_id: UUID
450 ) -> CleanupQueueMessage | None:
451 async with self._lock:
452 message = self._messages.get(message_id)
453 if message is not None and message.work_pool_id != work_pool_id:
454 return None
455 return _copy_model(message) if message is not None else None
457 async def read_dead_letter(
458 self, *, work_pool_id: UUID, message_id: UUID
459 ) -> CleanupQueueDeadLetter | None:
460 async with self._lock:
461 dead_letter = self._dead_letters.get(message_id)
462 if (
463 dead_letter is not None
464 and dead_letter.message.work_pool_id != work_pool_id
465 ):
466 return None
467 return _copy_model(dead_letter) if dead_letter is not None else None
469 async def wake_dispatchers(self, work_pool_id: UUID) -> CleanupQueueWakeup:
470 async with self._condition:
471 sequence = self._wakeup_sequences.get(work_pool_id, 0) + 1
472 self._wakeup_sequences[work_pool_id] = sequence
473 wakeup = CleanupQueueWakeup(
474 work_pool_id=work_pool_id,
475 sequence=sequence,
476 )
477 self._condition.notify_all()
478 return wakeup
480 async def read_wakeup_sequence(self, work_pool_id: UUID) -> int:
481 async with self._condition:
482 return self._wakeup_sequences.get(work_pool_id, 0)
484 async def wait_for_wakeup(
485 self,
486 work_pool_id: UUID,
487 *,
488 after: int = 0,
489 timeout: float | None = None,
490 ) -> CleanupQueueWakeup | None:
491 async with self._condition:
492 try:
493 if timeout is None:
494 await self._condition.wait_for(
495 lambda: self._wakeup_sequences.get(work_pool_id, 0) > after
496 )
497 else:
498 await asyncio.wait_for(
499 self._condition.wait_for(
500 lambda: self._wakeup_sequences.get(work_pool_id, 0) > after
501 ),
502 timeout=timeout,
503 )
504 except (TimeoutError, asyncio.TimeoutError):
505 return None
507 return CleanupQueueWakeup(
508 work_pool_id=work_pool_id,
509 sequence=self._wakeup_sequences[work_pool_id],
510 )
512 def _read_existing_message_locked(
513 self,
514 *,
515 message_id: UUID,
516 idempotency_key: str,
517 work_pool_id: UUID,
518 ) -> CleanupQueueMessage | None:
519 existing = self._messages.get(message_id)
520 if existing is None and (dead_letter := self._dead_letters.get(message_id)):
521 existing = dead_letter.message
522 if existing is None and (acked_message := self._acked_messages.get(message_id)):
523 existing = acked_message.message
525 if existing is not None:
526 if (
527 existing.idempotency_key != idempotency_key
528 or existing.work_pool_id != work_pool_id
529 ):
530 raise ValueError(
531 "message_id is already associated with a different cleanup message"
532 )
533 return existing
535 existing_message_id = self._idempotency_keys.get(
536 (work_pool_id, idempotency_key)
537 )
538 if existing_message_id is None:
539 return None
541 if existing := self._messages.get(existing_message_id):
542 return existing
543 if dead_letter := self._dead_letters.get(existing_message_id):
544 return dead_letter.message
545 if acked_message := self._acked_messages.get(existing_message_id):
546 return acked_message.message
548 self._idempotency_keys.pop((work_pool_id, idempotency_key), None)
549 return None
551 def _validate_current_reservation_locked(
552 self,
553 *,
554 operation: CleanupQueueOperation,
555 work_pool_id: UUID,
556 message_id: UUID,
557 reservation_token: str,
558 current_time: DateTime,
559 policy: _QueuePolicy,
560 ) -> CleanupQueueOperationResult | None:
561 message = self._messages.get(message_id)
562 if message is None:
563 return CleanupQueueOperationResult(
564 message_id=message_id,
565 operation=operation,
566 status="not_found",
567 reason="message_not_found",
568 )
570 if message.work_pool_id != work_pool_id:
571 return CleanupQueueOperationResult(
572 message_id=message_id,
573 operation=operation,
574 status="unauthorized",
575 reason="work_pool_mismatch",
576 )
578 reservation = self._reservations.get(message_id)
579 if reservation is None:
580 return CleanupQueueOperationResult(
581 message_id=message_id,
582 operation=operation,
583 status="not_current",
584 reason="no_active_reservation",
585 )
587 if reservation.token != reservation_token:
588 return CleanupQueueOperationResult(
589 message_id=message_id,
590 operation=operation,
591 status="invalid_token",
592 reason="reservation_token_mismatch",
593 )
595 if reservation.lease_expires_at <= current_time:
596 dead_letter = self._expire_message_lease_locked(
597 message_id=message_id,
598 current_time=current_time,
599 policy=policy,
600 )
601 if dead_letter is not None:
602 return CleanupQueueOperationResult(
603 message_id=message_id,
604 operation=operation,
605 status="dead_lettered",
606 reason="max_delivery_attempts_reached",
607 dead_letter=dead_letter,
608 )
609 return CleanupQueueOperationResult(
610 message_id=message_id,
611 operation=operation,
612 status="expired",
613 reason="lease_expired",
614 )
616 return None
618 def _expire_due_leases_locked(
619 self,
620 *,
621 current_time: DateTime,
622 policy: _QueuePolicy,
623 limit: int,
624 work_pool_id: UUID | None,
625 ) -> CleanupQueueLeaseExpiryResult:
626 if limit < 1: 626 ↛ 627line 626 didn't jump to line 627 because the condition on line 626 was never true
627 raise ValueError("limit must be at least 1")
629 redelivered: list[CleanupQueueMessage] = []
630 dead_lettered: list[CleanupQueueDeadLetter] = []
631 expired_count = 0
633 for message_id, reservation in tuple(self._reservations.items()): 633 ↛ 634line 633 didn't jump to line 634 because the loop on line 633 never started
634 if reservation.lease_expires_at > current_time:
635 continue
636 message = self._messages.get(message_id)
637 if work_pool_id is not None and (
638 message is None or message.work_pool_id != work_pool_id
639 ):
640 continue
642 dead_letter = self._expire_message_lease_locked(
643 message_id=message_id,
644 current_time=current_time,
645 policy=policy,
646 )
647 expired_count += 1
648 if dead_letter is not None:
649 dead_lettered.append(dead_letter)
650 elif message := self._messages.get(message_id):
651 redelivered.append(_copy_model(message))
653 if expired_count >= limit:
654 break
656 return CleanupQueueLeaseExpiryResult(
657 redelivered=redelivered,
658 dead_lettered=dead_lettered,
659 )
661 def _expire_message_lease_locked(
662 self,
663 *,
664 message_id: UUID,
665 current_time: DateTime,
666 policy: _QueuePolicy,
667 ) -> CleanupQueueDeadLetter | None:
668 message = self._messages.get(message_id)
669 if message is None:
670 self._reservations.pop(message_id, None)
671 return None
673 if message.delivery_count >= policy.max_delivery_attempts:
674 return self._move_to_dead_letter_locked(
675 message_id=message_id,
676 reason="max_delivery_attempts_reached",
677 current_time=current_time,
678 )
680 self._reservations.pop(message_id, None)
681 self._messages[message_id] = message.model_copy(
682 update={"updated_at": current_time}
683 )
684 return None
686 def _move_to_dead_letter_locked(
687 self,
688 *,
689 message_id: UUID,
690 reason: str,
691 current_time: DateTime,
692 release_reason: str | None = None,
693 ) -> CleanupQueueDeadLetter:
694 message = self._messages.pop(message_id)
695 reservation = self._reservations.pop(message_id, None)
696 dead_letter = CleanupQueueDeadLetter(
697 message=message,
698 reason=reason,
699 final_delivery_count=message.delivery_count,
700 moved_at=current_time,
701 reservation_token=reservation.token if reservation is not None else None,
702 lease_expires_at=(
703 reservation.lease_expires_at if reservation is not None else None
704 ),
705 release_reason=release_reason,
706 )
707 self._dead_letters[message_id] = dead_letter
708 record_cleanup_queue_dead_letter(dead_letter, source="memory_cleanup_queue")
709 return _copy_model(dead_letter)
711 def _expired_result_wakeup_work_pool_id_locked(
712 self, *, message_id: UUID, result: CleanupQueueOperationResult
713 ) -> UUID | None:
714 if result.status != "expired":
715 return None
716 message = self._messages.get(message_id)
717 return message.work_pool_id if message is not None else None
719 @staticmethod
720 def _queue_preference_passes(
721 *,
722 preferred_work_queue_ids: Iterable[UUID] | None,
723 allow_fallback_to_any_queue: bool,
724 ) -> tuple[set[UUID] | None, ...]:
725 if preferred_work_queue_ids is None:
726 return (None,)
728 preferred_queue_filter = set(preferred_work_queue_ids)
729 if allow_fallback_to_any_queue:
730 return (preferred_queue_filter, None)
731 return (preferred_queue_filter,)
733 @staticmethod
734 def _policy() -> _QueuePolicy:
735 worker_channel_settings = get_current_settings().server.worker_channel
736 retention_seconds = (
737 worker_channel_settings.cleanup_completed_idempotency_retention_seconds
738 )
739 return _QueuePolicy(
740 lease_duration=timedelta(
741 seconds=worker_channel_settings.cleanup_lease_seconds
742 ),
743 max_delivery_attempts=worker_channel_settings.cleanup_max_delivery_attempts,
744 completed_idempotency_retention=(
745 None
746 if retention_seconds is None
747 else timedelta(seconds=retention_seconds)
748 ),
749 )
751 def _prune_completed_idempotency_locked(
752 self, *, current_time: DateTime, policy: _QueuePolicy
753 ) -> None:
754 retention = policy.completed_idempotency_retention
755 if retention is None:
756 return
758 expired_message_ids = [
759 message_id
760 for message_id, acked_message in self._acked_messages.items()
761 if acked_message.completed_at + retention <= current_time
762 ]
763 for message_id in expired_message_ids:
764 acked_message = self._acked_messages.pop(message_id)
765 idempotency_key = (
766 acked_message.message.work_pool_id,
767 acked_message.message.idempotency_key,
768 )
769 if self._idempotency_keys.get(idempotency_key) == message_id:
770 self._idempotency_keys.pop(idempotency_key, None)