Coverage for qreward / utils / scheduler / priority_queue.py: 100.00%

1"""Thread-safe priority queue for scheduling tasks by priority level.""" 

2 

3import heapq (empty)

4import itertools (empty)

5import threading (empty)

6import time (empty)

7from dataclasses import dataclass, field (empty)

8from enum import IntEnum (empty)

9from typing import Any, Optional (empty)

10 

11 

12class Priority(IntEnum): (empty)

13 """Task priority levels. Lower value = higher priority.""" 

14 

15 HIGH = 0 (empty)

16 NORMAL = 5 (empty)

17 LOW = 9 (empty)

18 

19 

20@dataclass(order=True) (empty)

21class _PriorityEntry: (empty)

22 """Internal entry for the priority heap. 

23 

24 Ordering: (effective_priority, sequence) ensures stable FIFO within 

25 the same priority level. 

26 

27 Attributes: 

28 effective_priority: Current priority (may be boosted by starvation 

29 protection). 

30 sequence: Monotonically increasing counter for FIFO ordering. 

31 enqueue_time: Timestamp when the entry was added (for starvation 

32 detection). 

33 item: The actual payload (excluded from ordering). 

34 cancelled: Whether this entry has been cancelled. 

35 """ 

36 

37 effective_priority: int (empty)

38 sequence: int (empty)

39 enqueue_time: float = field(compare=False) (empty)

40 item: Any = field(compare=False) (empty)

41 cancelled: bool = field(default=False, compare=False) (empty)

42 

43 

44class PriorityTaskQueue: (empty)

45 """Thread-safe priority queue with starvation protection. 

46 

47 Tasks are dequeued in priority order (lower value = higher priority). 

48 Within the same priority, tasks follow FIFO order. 

49 

50 Starvation protection: tasks waiting longer than ``starvation_threshold`` 

51 seconds are automatically promoted to HIGH priority on the next ``get``. 

52 

53 Example: 

54 >>> queue = PriorityTaskQueue(starvation_threshold=30.0) 

55 >>> queue.put("critical_task", priority=Priority.HIGH) 

56 >>> queue.put("normal_task", priority=Priority.NORMAL) 

57 >>> queue.get() # returns "critical_task" 

58 >>> queue.get() # returns "normal_task" 

59 """ 

60 

61 def __init__(self, starvation_threshold: float = 60.0) -> None: (empty)

62 """Initialize the priority queue. 

63 

64 Args: 

65 starvation_threshold: Seconds after which a waiting task is 

66 promoted to HIGH priority. Set to 0 to disable. 

67 """ 

68 if starvation_threshold < 0: (empty)

69 raise ValueError("starvation_threshold must be >= 0") (empty)

70 

71 self._heap: list[_PriorityEntry] = [] (empty)

72 self._counter = itertools.count() (empty)

73 self._lock = threading.Lock() (empty)

74 self._starvation_threshold = starvation_threshold (empty)

75 

76 def put(self, item: Any, priority: int = Priority.NORMAL) -> None: (empty)

77 """Add an item to the queue with the given priority. 

78 

79 Args: 

80 item: The item to enqueue. 

81 priority: Priority value (0-9, lower = higher priority). 

82 

83 Raises: 

84 ValueError: If priority is not in range [0, 9]. 

85 """ 

86 if not 0 <= priority <= 9: (empty)

87 raise ValueError(f"priority must be in [0, 9], got {priority}") (empty)

88 

89 entry = _PriorityEntry( (empty)

90 effective_priority=priority, 

91 sequence=next(self._counter), 

92 enqueue_time=time.monotonic(), 

93 item=item, 

94 ) 

95 with self._lock: (empty)

96 heapq.heappush(self._heap, entry) (empty)

97 

98 def get(self) -> Optional[Any]: (empty)

99 """Remove and return the highest-priority item. 

100 

101 Applies starvation protection before selecting: any entry waiting 

102 longer than the threshold is promoted to HIGH priority. 

103 

104 Returns: 

105 The item, or None if the queue is empty. 

106 """ 

107 with self._lock: (empty)

108 self._apply_starvation_protection() (empty)

109 while self._heap: (empty)

110 entry = heapq.heappop(self._heap) (empty)

111 if not entry.cancelled: (empty)

112 return entry.item (empty)

113 return None (empty)

114 

115 def peek(self) -> Optional[Any]: (empty)

116 """Return the highest-priority item without removing it. 

117 

118 Returns: 

119 The item, or None if the queue is empty. 

120 """ 

121 with self._lock: (empty)

122 self._apply_starvation_protection() (empty)

123 for entry in self._heap: (empty)

124 if not entry.cancelled: (empty)

125 return entry.item (empty)

126 return None (empty)

127 

128 @property (empty)

129 def queue_size(self) -> int: (empty)

130 """Return the number of non-cancelled items in the queue.""" 

131 with self._lock: (empty)

132 return sum(1 for entry in self._heap if not entry.cancelled) (empty)

133 

134 @property (empty)

135 def is_empty(self) -> bool: (empty)

136 """Check if the queue has no pending items.""" 

137 return self.queue_size == 0 (empty)

138 

139 def clear(self) -> None: (empty)

140 """Remove all items from the queue.""" 

141 with self._lock: (empty)

142 self._heap.clear() (empty)

143 

144 def snapshot(self) -> list[dict]: (empty)

145 """Return a snapshot of all pending entries for debugging. 

146 

147 Returns: 

148 List of dicts with priority, sequence, wait_seconds, item info. 

149 """ 

150 now = time.monotonic() (empty)

151 with self._lock: (empty)

152 return [ (empty)

153 { 

154 "priority": entry.effective_priority, 

155 "sequence": entry.sequence, 

156 "wait_seconds": round(now - entry.enqueue_time, 3), 

157 "cancelled": entry.cancelled, 

158 } 

159 for entry in sorted(self._heap) 

160 ] 

161 

162 def _apply_starvation_protection(self) -> None: (empty)

163 """Promote starved entries to HIGH priority. Must hold lock.""" 

164 if self._starvation_threshold <= 0: (empty)

165 return (empty)

166 

167 now = time.monotonic() (empty)

168 promoted = False (empty)

169 

170 for entry in self._heap: (empty)

171 if entry.cancelled: (empty)

172 continue (empty)

173 wait_time = now - entry.enqueue_time (empty)

174 if ( (empty)

175 wait_time >= self._starvation_threshold 

176 and entry.effective_priority > Priority.HIGH 

177 ): 

178 entry.effective_priority = Priority.HIGH (empty)

179 promoted = True (empty)

180 

181 if promoted: (empty)

182 heapq.heapify(self._heap) (empty)