33
44use std:: collections:: BTreeMap ;
55use std:: collections:: BTreeSet ;
6+ use std:: collections:: HashSet ;
67use std:: pin:: Pin ;
78use std:: task:: Context ;
89use std:: task:: Poll ;
@@ -209,7 +210,10 @@ impl State {
209210 let mut current_end = requests[ 0 ] . offset + requests[ 0 ] . length as u64 ;
210211 let align = * self . coalesced_buffer_alignment as u64 ;
211212
212- let mut keys_to_remove = Vec :: new ( ) ;
213+ // Track requests that we've already decided to remove (or that were cancelled) so that
214+ // we don't repeatedly process them during range scans.
215+ let mut keys_to_remove: Vec < ( u64 , RequestId ) > = Vec :: new ( ) ;
216+ let mut ids_to_remove: HashSet < RequestId > = HashSet :: new ( ) ;
213217 let mut found_new_requests = true ;
214218
215219 // Keep expanding the window while we can find new requests within constraints
@@ -225,8 +229,8 @@ impl State {
225229 . requests_by_offset
226230 . range ( ( scan_start, RequestId :: MIN ) ..=( scan_end, RequestId :: MAX ) )
227231 {
228- // Skip if we've already marked this request for removal
229- if keys_to_remove . iter ( ) . any ( | & ( _ , id ) | id == req_id) {
232+ // Skip if we've already marked this request for removal.
233+ if ids_to_remove . contains ( & req_id) {
230234 continue ;
231235 }
232236
@@ -236,13 +240,15 @@ impl State {
236240 . or_else ( || self . requests . get ( & req_id) )
237241 . vortex_expect ( "Missing request in requests_by_offset" ) ;
238242
239- // Skip any cancelled requests
243+ // Skip any cancelled requests.
240244 if req. callback . is_closed ( ) {
241- keys_to_remove. push ( ( req_offset, req_id) ) ;
245+ if ids_to_remove. insert ( req_id) {
246+ keys_to_remove. push ( ( req_offset, req_id) ) ;
247+ }
242248 continue ;
243249 }
244250
245- // Check if this request is within coalescing distance of our current range
251+ // Check if this request is within coalescing distance of our current range.
246252 let req_end = req_offset + req. length as u64 ;
247253 if ( req_offset <= current_end + window. distance && req_end >= current_start)
248254 || ( req_end + window. distance >= current_start && req_offset <= current_end)
@@ -267,7 +273,9 @@ impl State {
267273 . vortex_expect ( "Missing request in requests_by_offset" ) ;
268274
269275 requests. push ( req) ;
270- keys_to_remove. push ( ( req_offset, req_id) ) ;
276+ if ids_to_remove. insert ( req_id) {
277+ keys_to_remove. push ( ( req_offset, req_id) ) ;
278+ }
271279 found_new_requests = true ;
272280 }
273281 }
0 commit comments