@@ -12,7 +12,6 @@ use futures::StreamExt;
1212use futures:: TryFutureExt ;
1313use futures:: future:: BoxFuture ;
1414use futures:: future:: Shared ;
15- use futures:: future:: WeakShared ;
1615use futures:: stream:: FuturesUnordered ;
1716use parking_lot:: Mutex ;
1817use vortex_array:: buffer:: BufferHandle ;
@@ -47,89 +46,26 @@ pub(super) fn segment_ranges(locations: &[SegmentLocation]) -> SegmentRanges {
4746}
4847
4948/// Serves splits from any [`SegmentSource`], by segment id.
50- ///
51- /// Its sessions share a segment's read while any of them still holds it, so splits that run at
52- /// once and need the same segment, such as chunks that each cut into one column segment, read
53- /// it once. A scan that retains its reads keeps every read for its lifetime, so splits that run
54- /// one after another share them too.
5549pub ( crate ) struct SegmentScanIo {
5650 segments : Arc < dyn SegmentSource > ,
5751 ranges : SegmentRanges ,
58- reads : SharedReads ,
5952}
6053
6154impl SegmentScanIo {
62- /// A scan's IO, which retains its reads when `retain` is set.
63- ///
64- /// Retaining suits a sparse selection, such as an engine fetching a few rows by id: its
65- /// reads are few, but each wide column segment it touches is needed by every chunk that
66- /// selects a row of it.
67- pub ( crate ) fn new (
68- segments : Arc < dyn SegmentSource > ,
69- ranges : SegmentRanges ,
70- retain : bool ,
71- ) -> Self {
72- Self {
73- segments,
74- ranges,
75- reads : Arc :: new ( Mutex :: new ( ScanReads {
76- retain,
77- reads : HashMap :: default ( ) ,
78- } ) ) ,
79- }
55+ pub ( crate ) fn new ( segments : Arc < dyn SegmentSource > , ranges : SegmentRanges ) -> Self {
56+ Self { segments, ranges }
8057 }
8158}
8259
8360impl IoService for SegmentScanIo {
8461 fn session ( & self ) -> Arc < dyn IoSource > {
85- Arc :: new ( SegmentIoSource :: with_shared_reads (
62+ Arc :: new ( SegmentIoSource :: new (
8663 Arc :: clone ( & self . segments ) ,
8764 Arc :: clone ( & self . ranges ) ,
88- Arc :: clone ( & self . reads ) ,
8965 ) )
9066 }
9167}
9268
93- /// The reads of a scan's sessions, by segment.
94- type SharedReads = Arc < Mutex < ScanReads > > ;
95-
96- /// The reads of a scan's sessions, by segment.
97- #[ derive( Default ) ]
98- struct ScanReads {
99- /// Whether reads are held for the scan's lifetime rather than while a session holds them.
100- retain : bool ,
101- reads : HashMap < SegmentId , ScanRead > ,
102- }
103-
104- /// A read of a scan, held weakly so that sharing it keeps no segment in memory longer, or
105- /// retained.
106- enum ScanRead {
107- Weak ( WeakShared < SegmentRead > ) ,
108- Retained ( Prefetched ) ,
109- }
110-
111- impl ScanReads {
112- /// The read of `segment` the scan still has.
113- fn get ( & self , segment : SegmentId ) -> Option < Prefetched > {
114- match self . reads . get ( & segment) ? {
115- ScanRead :: Weak ( read) => read. upgrade ( ) ,
116- ScanRead :: Retained ( read) => Some ( read. clone ( ) ) ,
117- }
118- }
119-
120- fn insert ( & mut self , segment : SegmentId , read : & Prefetched ) {
121- let read = if self . retain {
122- ScanRead :: Retained ( read. clone ( ) )
123- } else {
124- let Some ( read) = read. downgrade ( ) else {
125- return ;
126- } ;
127- ScanRead :: Weak ( read)
128- } ;
129- self . reads . insert ( segment, read) ;
130- }
131- }
132-
13369/// Serves a split's range requests from a segment source, without blocking.
13470///
13571/// Each requested range names one segment. [`submit`](IoSource::submit) starts its read at once,
@@ -149,57 +85,22 @@ pub(super) struct SegmentIoSource {
14985 prefetched : Mutex < HashMap < SegmentId , Prefetched > > ,
15086 /// Announced segments, requested and never polled.
15187 announced : Mutex < HashMap < SegmentId , VecDeque < SegmentFuture > > > ,
152- /// The reads of the scan's other sessions.
153- shared : SharedReads ,
15488}
15589
156- /// The read of a segment.
157- type SegmentRead = BoxFuture < ' static , Result < BufferHandle , Arc < VortexError > > > ;
158-
15990/// A read started ahead of the fetch that needs it.
160- type Prefetched = Shared < SegmentRead > ;
91+ type Prefetched = Shared < BoxFuture < ' static , Result < BufferHandle , Arc < VortexError > > > > ;
16192
16293impl SegmentIoSource {
163- #[ cfg( test) ]
16494 pub ( super ) fn new ( segments : Arc < dyn SegmentSource > , ranges : SegmentRanges ) -> Self {
165- Self :: with_shared_reads ( segments, ranges, SharedReads :: default ( ) )
166- }
167-
168- fn with_shared_reads (
169- segments : Arc < dyn SegmentSource > ,
170- ranges : SegmentRanges ,
171- shared : SharedReads ,
172- ) -> Self {
17395 Self {
17496 segments,
17597 ranges,
17698 reads : Mutex :: new ( FuturesUnordered :: new ( ) ) ,
17799 ready : Mutex :: default ( ) ,
178100 prefetched : Mutex :: default ( ) ,
179101 announced : Mutex :: default ( ) ,
180- shared,
181102 }
182103 }
183-
184- /// The read of `segment`: one another session of the scan still holds, else the read of an
185- /// announcement of it, else a new one.
186- fn read ( & self , segment : SegmentId ) -> Prefetched {
187- let mut shared = self . shared . lock ( ) ;
188- if let Some ( read) = shared. get ( segment) {
189- return read;
190- }
191- let read = self
192- . announced
193- . lock ( )
194- . get_mut ( & segment)
195- . and_then ( VecDeque :: pop_front)
196- . unwrap_or_else ( || self . segments . request ( segment) )
197- . map_err ( Arc :: new)
198- . boxed ( )
199- . shared ( ) ;
200- shared. insert ( segment, & read) ;
201- read
202- }
203104}
204105
205106impl IoSource for SegmentIoSource {
@@ -228,7 +129,15 @@ impl IoSource for SegmentIoSource {
228129 }
229130 if request. intent == IoIntent :: Prefetch {
230131 prefetched. entry ( segment) . or_insert_with ( || {
231- let read = self . read ( segment) ;
132+ let read = self
133+ . announced
134+ . lock ( )
135+ . get_mut ( & segment)
136+ . and_then ( VecDeque :: pop_front)
137+ . unwrap_or_else ( || self . segments . request ( segment) )
138+ . map_err ( Arc :: new)
139+ . boxed ( )
140+ . shared ( ) ;
232141 // Polling registers the read as wanted, so the source starts it now.
233142 drop (
234143 read. clone ( )
@@ -243,7 +152,16 @@ impl IoSource for SegmentIoSource {
243152 // first has completed.
244153 let mut bytes = prefetched
245154 . entry ( segment)
246- . or_insert_with ( || self . read ( segment) )
155+ . or_insert_with ( || {
156+ self . announced
157+ . lock ( )
158+ . get_mut ( & segment)
159+ . and_then ( VecDeque :: pop_front)
160+ . unwrap_or_else ( || self . segments . request ( segment) )
161+ . map_err ( Arc :: new)
162+ . boxed ( )
163+ . shared ( )
164+ } )
247165 . clone ( )
248166 . map_err ( move |err| vortex_err ! ( "read of segment {segment} failed: {err}" ) )
249167 . boxed ( ) ;
@@ -331,12 +249,10 @@ mod tests {
331249 use vortex_io:: request:: IoOwnerId ;
332250 use vortex_io:: request:: IoRequest ;
333251 use vortex_io:: request:: IoRequestId ;
334- use vortex_io:: request:: IoService ;
335252 use vortex_io:: request:: IoSource ;
336253 use vortex_io:: request:: IoTarget ;
337254
338255 use super :: SegmentIoSource ;
339- use super :: SegmentScanIo ;
340256 use super :: segment_ranges;
341257 use crate :: scan:: execution:: SegmentLocation ;
342258 use crate :: segments:: SegmentFuture ;
@@ -454,75 +370,6 @@ mod tests {
454370 Ok ( ( ) )
455371 }
456372
457- /// Sessions of one scan share a segment's read while one of them holds it, and read it again
458- /// once none does.
459- #[ tokio:: test]
460- async fn sessions_share_a_held_read ( ) -> VortexResult < ( ) > {
461- let segments = Arc :: new ( CountingSegments :: default ( ) ) ;
462- let location = SegmentLocation {
463- offset : 0 ,
464- length : 4 ,
465- alignment : Alignment :: none ( ) ,
466- } ;
467- let scan = SegmentScanIo :: new (
468- Arc :: clone ( & segments) as _ ,
469- segment_ranges ( & [ location] ) ,
470- false ,
471- ) ;
472- let first = scan. session ( ) ;
473- let second = scan. session ( ) ;
474- for ( io, id) in [ ( & first, 0 ) , ( & second, 1 ) ] {
475- io. submit ( IoOwnerId ( 0 ) , vec ! [ request( IoIntent :: Fetch , id) ] ) ?;
476- let completion = std:: future:: poll_fn ( |cx| io. poll_completion ( cx) ) . await ?;
477- assert_eq ! ( completion. request, IoRequestId ( id) ) ;
478- assert ! ( completion. result. is_ok( ) ) ;
479- }
480- assert_eq ! ( segments. 0 . load( Ordering :: Relaxed ) , 1 ) ;
481-
482- first. clear ( ) ;
483- second. clear ( ) ;
484- let third = scan. session ( ) ;
485- third. submit ( IoOwnerId ( 0 ) , vec ! [ request( IoIntent :: Fetch , 2 ) ] ) ?;
486- assert ! (
487- std:: future:: poll_fn( |cx| third. poll_completion( cx) )
488- . await ?
489- . result
490- . is_ok( )
491- ) ;
492- assert_eq ! ( segments. 0 . load( Ordering :: Relaxed ) , 2 ) ;
493- Ok ( ( ) )
494- }
495-
496- /// A scan that retains its reads shares a read with a later session after the earlier one
497- /// let it go.
498- #[ tokio:: test]
499- async fn retained_reads_outlive_their_sessions ( ) -> VortexResult < ( ) > {
500- let segments = Arc :: new ( CountingSegments :: default ( ) ) ;
501- let location = SegmentLocation {
502- offset : 0 ,
503- length : 4 ,
504- alignment : Alignment :: none ( ) ,
505- } ;
506- let scan = SegmentScanIo :: new (
507- Arc :: clone ( & segments) as _ ,
508- segment_ranges ( & [ location] ) ,
509- true ,
510- ) ;
511- for id in 0 ..2 {
512- let io = scan. session ( ) ;
513- io. submit ( IoOwnerId ( 0 ) , vec ! [ request( IoIntent :: Fetch , id) ] ) ?;
514- assert ! (
515- std:: future:: poll_fn( |cx| io. poll_completion( cx) )
516- . await ?
517- . result
518- . is_ok( )
519- ) ;
520- io. clear ( ) ;
521- }
522- assert_eq ! ( segments. 0 . load( Ordering :: Relaxed ) , 1 ) ;
523- Ok ( ( ) )
524- }
525-
526373 #[ tokio:: test]
527374 async fn announcements_keep_independent_consumers ( ) -> VortexResult < ( ) > {
528375 let segments = Arc :: new ( CountingSegments :: default ( ) ) ;
0 commit comments