@@ -36,11 +36,16 @@ pub struct MergeableQueueItem {
36
36
}
37
37
38
38
#[ derive( Debug , Clone ) ]
39
+ enum QueueMessage {
40
+ Item ( MergeableQueueItem ) ,
41
+ Shutdown ,
42
+ }
43
+
39
44
struct Item {
40
45
/// When to process item (None = immediate).
41
46
/// Reversed to create min-heap for expirations.
42
47
expiration : Reverse < Option < Instant > > ,
43
- inner : MergeableQueueItem ,
48
+ inner : QueueMessage ,
44
49
}
45
50
46
51
impl PartialEq for Item {
@@ -94,6 +99,23 @@ pub fn create_mergeable_queue() -> (MergeableQueueSender, MergeableQueueReceiver
94
99
}
95
100
96
101
impl MergeableQueueSender {
102
+ pub fn is_closed ( & self ) -> bool {
103
+ self . inner . closed . load ( Ordering :: Relaxed )
104
+ }
105
+
106
+ pub fn shutdown ( & self ) {
107
+ let mut queue = self . inner . queue . lock ( ) . unwrap ( ) ;
108
+
109
+ // Send shutdown message
110
+ queue. push ( Item {
111
+ expiration : Reverse ( None ) ,
112
+ inner : QueueMessage :: Shutdown ,
113
+ } ) ;
114
+
115
+ // and wake receiver for immediate processing.
116
+ self . inner . notify . notify_one ( ) ;
117
+ }
118
+
97
119
pub fn enqueue ( & self , repo : GithubRepoName , pr_number : PullRequestNumber ) {
98
120
let expiration = Some ( Instant :: now ( ) + BASE_DELAY ) ;
99
121
@@ -148,7 +170,7 @@ impl MergeableQueueSender {
148
170
149
171
queue. push ( Item {
150
172
expiration : Reverse ( expiration) ,
151
- inner : item,
173
+ inner : QueueMessage :: Item ( item) ,
152
174
} ) ;
153
175
154
176
if should_notify {
@@ -157,19 +179,12 @@ impl MergeableQueueSender {
157
179
}
158
180
}
159
181
160
- impl Drop for MergeableQueueSender {
161
- fn drop ( & mut self ) {
162
- // Count is 2 if this is the last sender (this sender + receiver)
163
- if Arc :: strong_count ( & self . inner ) == 2 {
164
- self . inner . closed . store ( true , Ordering :: Relaxed ) ;
165
- // Notify the receiver in case it's waiting.
166
- self . inner . notify . notify_one ( ) ;
167
- }
182
+ impl MergeableQueueReceiver {
183
+ fn is_closed ( & self ) -> bool {
184
+ self . inner . closed . load ( Ordering :: Relaxed )
168
185
}
169
- }
170
186
171
- impl MergeableQueueReceiver {
172
- fn peek_inner ( & self ) -> Result < MergeableQueueItem , Option < Duration > > {
187
+ fn peek_inner ( & self ) -> Result < QueueMessage , Option < Duration > > {
173
188
let now = Instant :: now ( ) ;
174
189
let mut queue = self . inner . queue . lock ( ) . unwrap ( ) ;
175
190
@@ -206,30 +221,32 @@ impl MergeableQueueReceiver {
206
221
pub async fn dequeue ( & self ) -> Option < ( MergeableQueueItem , MergeableQueueSender ) > {
207
222
loop {
208
223
// Closed and empty queue, we're done.
209
- if self . inner . closed . load ( Ordering :: Relaxed )
210
- && self . inner . queue . lock ( ) . unwrap ( ) . is_empty ( )
211
- {
224
+ if self . is_closed ( ) && self . inner . queue . lock ( ) . unwrap ( ) . is_empty ( ) {
212
225
return None ;
213
226
}
214
227
215
228
match self . peek_inner ( ) {
216
229
// Item is ready.
217
- Ok ( item) => {
230
+ Ok ( QueueMessage :: Item ( item) ) => {
218
231
break Some ( (
219
232
item,
220
233
MergeableQueueSender {
221
234
inner : self . inner . clone ( ) ,
222
235
} ,
223
236
) ) ;
224
237
}
238
+ Ok ( QueueMessage :: Shutdown ) => {
239
+ self . inner . closed . store ( true , Ordering :: Relaxed ) ;
240
+ break None ;
241
+ }
225
242
// Item exists but not ready, wait until then or until notified of a higher priority item.
226
243
Err ( Some ( duration) ) => {
227
244
let _ = timeout ( duration, self . inner . notify . notified ( ) ) . await ;
228
245
}
229
246
// Queue is empty, wait until notified of a new item.
230
247
Err ( None ) => {
231
248
// If closed, we're done
232
- if self . inner . closed . load ( Ordering :: Relaxed ) {
249
+ if self . is_closed ( ) {
233
250
return None ;
234
251
}
235
252
// else, wait until next item.
@@ -309,7 +326,15 @@ pub async fn handle_mergeable_queue_item(
309
326
tracing:: info!(
310
327
"Retrying mergeable state check for PR: {pull_request} ({attempt}/{MAX_RETRIES})" ,
311
328
) ;
312
- mq_tx. enqueue_retry ( mq_item) ;
329
+
330
+ if mq_tx. is_closed ( ) {
331
+ tracing:: error!(
332
+ "Attempted to enqueue retry for PR {pull_request} but mergeable queue is shutdown"
333
+ ) ;
334
+ } else {
335
+ mq_tx. enqueue_retry ( mq_item) ;
336
+ }
337
+
313
338
return Ok ( ( ) ) ;
314
339
}
315
340
0 commit comments