@@ -117,64 +117,145 @@ impl VortexReadAt for ObjectStoreReadAt {
117117 let allocator = Arc :: clone ( & self . allocator ) ;
118118 let range = offset..( offset + length as u64 ) ;
119119
120- async move {
121- let mut buffer = allocator. allocate ( length, alignment) ?;
122-
123- let response = store
124- . get_opts (
125- & path,
126- GetOptions {
127- range : Some ( GetRange :: Bounded ( range. clone ( ) ) ) ,
128- ..Default :: default ( )
129- } ,
130- )
131- . await ?;
132-
133- let buffer = match response. payload {
134- #[ cfg( not( target_arch = "wasm32" ) ) ]
135- GetResultPayload :: File ( file, _) => {
136- handle
137- . spawn_blocking ( move || {
138- read_exact_at ( & file, buffer. as_mut_slice ( ) , range. start ) ?;
139- Ok :: < _ , io:: Error > ( buffer)
140- } )
141- . await
142- . map_err ( io:: Error :: other) ?
143- }
144- #[ cfg( target_arch = "wasm32" ) ]
145- GetResultPayload :: File ( ..) => {
146- unreachable ! ( "File payload not supported on wasm32" )
147- }
148- GetResultPayload :: Stream ( mut byte_stream) => {
149- let mut written = 0usize ;
150- while let Some ( bytes) = byte_stream. next ( ) . await {
151- let bytes = bytes?;
152- let end = written + bytes. len ( ) ;
153- vortex_ensure ! (
154- end <= length,
155- "Object store stream returned too many bytes: {} > expected {} (range: {:?})" ,
156- end,
157- length,
158- range
159- ) ;
160- buffer. as_mut_slice ( ) [ written..end] . copy_from_slice ( & bytes) ;
161- written = end;
162- }
163-
164- vortex_ensure ! (
165- written == length,
166- "Object store stream returned {} bytes but expected {} bytes (range: {:?})" ,
167- written,
168- length,
169- range
170- ) ;
171-
172- buffer
173- }
174- } ;
175-
176- Ok ( BufferHandle :: new_host ( buffer. freeze ( ) ) )
177- }
120+ // Requires to deal with borrowed lifetimes
121+ let io_handle = handle. clone ( ) ;
122+
123+ handle
124+ . spawn_io ( async move {
125+ let mut buffer = allocator. allocate ( length, alignment) ?;
126+
127+ let response = store
128+ . get_opts (
129+ & path,
130+ GetOptions {
131+ range : Some ( GetRange :: Bounded ( range. clone ( ) ) ) ,
132+ ..Default :: default ( )
133+ } ,
134+ )
135+ . await ?;
136+
137+ let buffer = match response. payload {
138+ #[ cfg( not( target_arch = "wasm32" ) ) ]
139+ GetResultPayload :: File ( file, _) => {
140+ io_handle
141+ . spawn_blocking ( move || {
142+ read_exact_at ( & file, buffer. as_mut_slice ( ) , range. start ) ?;
143+ Ok :: < _ , io:: Error > ( buffer)
144+ } )
145+ . await
146+ . map_err ( io:: Error :: other) ?
147+ }
148+ #[ cfg( target_arch = "wasm32" ) ]
149+ GetResultPayload :: File ( ..) => {
150+ unreachable ! ( "File payload not supported on wasm32" )
151+ }
152+ GetResultPayload :: Stream ( mut byte_stream) => {
153+ let mut written = 0usize ;
154+ while let Some ( bytes) = byte_stream. next ( ) . await {
155+ let bytes = bytes?;
156+ let end = written + bytes. len ( ) ;
157+ vortex_ensure ! (
158+ end <= length,
159+ "Object store stream returned too many bytes: {} > expected {} (range: {:?})" ,
160+ end,
161+ length,
162+ range
163+ ) ;
164+ buffer. as_mut_slice ( ) [ written..end] . copy_from_slice ( & bytes) ;
165+ written = end;
166+ }
167+
168+ vortex_ensure ! (
169+ written == length,
170+ "Object store stream returned {} bytes but expected {} bytes (range: {:?})" ,
171+ written,
172+ length,
173+ range
174+ ) ;
175+
176+ buffer
177+ }
178+ } ;
179+
180+ Ok ( BufferHandle :: new_host ( buffer. freeze ( ) ) )
181+ } )
178182 . boxed ( )
179183 }
180184}
185+
186+ #[ cfg( test) ]
187+ mod tests {
188+
189+ use std:: sync:: atomic:: AtomicUsize ;
190+ use std:: sync:: atomic:: Ordering ;
191+
192+ use object_store:: PutPayload ;
193+ use object_store:: memory:: InMemory ;
194+
195+ use super :: * ;
196+ use crate :: runtime:: AbortHandle ;
197+ use crate :: runtime:: AbortHandleRef ;
198+ use crate :: runtime:: Executor ;
199+
200+ const TEST_DATA : & [ u8 ] = b"object store test data" ;
201+
202+ #[ derive( Default ) ]
203+ struct CountingExecutor {
204+ spawn_count : AtomicUsize ,
205+ spawn_io_count : AtomicUsize ,
206+ }
207+
208+ impl Executor for CountingExecutor {
209+ fn spawn ( & self , fut : BoxFuture < ' static , ( ) > ) -> AbortHandleRef {
210+ self . spawn_count . fetch_add ( 1 , Ordering :: SeqCst ) ;
211+ TokioAbortHandle :: new_handle ( tokio:: spawn ( fut) . abort_handle ( ) )
212+ }
213+
214+ fn spawn_io ( & self , fut : BoxFuture < ' static , ( ) > ) -> AbortHandleRef {
215+ self . spawn_io_count . fetch_add ( 1 , Ordering :: SeqCst ) ;
216+ TokioAbortHandle :: new_handle ( tokio:: spawn ( fut) . abort_handle ( ) )
217+ }
218+
219+ fn spawn_cpu ( & self , task : Box < dyn FnOnce ( ) + Send + ' static > ) -> AbortHandleRef {
220+ TokioAbortHandle :: new_handle ( tokio:: spawn ( async move { task ( ) } ) . abort_handle ( ) )
221+ }
222+
223+ fn spawn_blocking_io ( & self , task : Box < dyn FnOnce ( ) + Send + ' static > ) -> AbortHandleRef {
224+ TokioAbortHandle :: new_handle ( tokio:: task:: spawn_blocking ( task) . abort_handle ( ) )
225+ }
226+ }
227+
228+ struct TokioAbortHandle ( tokio:: task:: AbortHandle ) ;
229+
230+ impl TokioAbortHandle {
231+ fn new_handle ( handle : tokio:: task:: AbortHandle ) -> AbortHandleRef {
232+ Box :: new ( Self ( handle) )
233+ }
234+ }
235+
236+ impl AbortHandle for TokioAbortHandle {
237+ fn abort ( self : Box < Self > ) {
238+ self . 0 . abort ( ) ;
239+ }
240+ }
241+
242+ #[ tokio:: test]
243+ async fn read_at_uses_spawn_io ( ) -> anyhow:: Result < ( ) > {
244+ let executor = Arc :: new ( CountingExecutor :: default ( ) ) ;
245+ let runtime = Arc :: clone ( & executor) as Arc < dyn Executor > ;
246+ let handle = Handle :: new ( Arc :: downgrade ( & runtime) ) ;
247+
248+ let store = Arc :: new ( InMemory :: new ( ) ) as Arc < dyn ObjectStore > ;
249+ let path = ObjectPath :: from ( "test.bin" ) ;
250+ store. put ( & path, PutPayload :: from_static ( TEST_DATA ) ) . await ?;
251+
252+ let reader = ObjectStoreReadAt :: new ( store, path, handle) ;
253+ let buffer = reader. read_at ( 7 , 5 , Alignment :: new ( 1 ) ) . await ?;
254+
255+ assert_eq ! ( buffer. to_host( ) . await . as_slice( ) , b"store" ) ;
256+ assert_eq ! ( executor. spawn_io_count. load( Ordering :: SeqCst ) , 1 ) ;
257+ assert_eq ! ( executor. spawn_count. load( Ordering :: SeqCst ) , 0 ) ;
258+
259+ Ok ( ( ) )
260+ }
261+ }
0 commit comments