arrow_avro/reader/async_reader/
spawn.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18use std::future::Future;
19use std::ops::Range;
20
21use bytes::Bytes;
22use futures::future::BoxFuture;
23use futures::{FutureExt, TryFutureExt};
24use tokio::runtime::Handle;
25
26use crate::errors::AvroError;
27use crate::reader::async_reader::AsyncFileReader;
28
29/// An [`AsyncFileReader`] that performs I/O on a separate tokio runtime.
30///
31/// Tokio is a cooperative scheduler, and relies on tasks yielding in a timely
32/// manner to service IO. Therefore, running IO and CPU-bound tasks, such as
33/// avro decoding, on the same tokio runtime can lead to degraded throughput,
34/// dropped connections and other issues. For more information see [here].
35///
36/// This wrapper spawns each operation of the inner reader onto the provided
37/// runtime [`Handle`], so that the runtime driving the avro decoding does not
38/// also drive the I/O.
39///
40/// The inner reader must be [`Clone`] (typically an `Arc`'d handle to some
41/// shared resource) as each spawned task requires a `'static` copy of it.
42///
43/// [here]: https://www.influxdata.com/blog/using-rustlangs-async-tokio-runtime-for-cpu-bound-tasks/
44#[derive(Clone, Debug)]
45pub struct SpawnedReader<R> {
46    inner: R,
47    handle: Handle,
48}
49
50impl<R> SpawnedReader<R> {
51    /// Creates a new [`SpawnedReader`] that performs the I/O of `inner` on `handle`
52    pub fn new(inner: R, handle: Handle) -> Self {
53        Self { inner, handle }
54    }
55
56    /// Returns the inner reader
57    pub fn into_inner(self) -> R {
58        self.inner
59    }
60}
61
62/// Spawns `fut` on `handle`, propagating panics and mapping task cancellation
63/// to [`AvroError::External`]
64fn spawn<T>(
65    handle: &Handle,
66    fut: impl Future<Output = Result<T, AvroError>> + Send + 'static,
67) -> BoxFuture<'static, Result<T, AvroError>>
68where
69    T: Send + 'static,
70{
71    handle
72        .spawn(fut)
73        .map_ok_or_else(
74            |e| match e.try_into_panic() {
75                Err(e) => Err(AvroError::External(Box::new(e))),
76                Ok(p) => std::panic::resume_unwind(p),
77            },
78            |res| res,
79        )
80        .boxed()
81}
82
83impl<R> AsyncFileReader for SpawnedReader<R>
84where
85    R: AsyncFileReader + Clone + Send + 'static,
86{
87    fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>> {
88        let mut inner = self.inner.clone();
89        spawn(&self.handle, async move { inner.get_bytes(range).await })
90    }
91
92    fn get_byte_ranges(
93        &mut self,
94        ranges: Vec<Range<u64>>,
95    ) -> BoxFuture<'_, Result<Vec<Bytes>, AvroError>> {
96        let mut inner = self.inner.clone();
97        spawn(
98            &self.handle,
99            async move { inner.get_byte_ranges(ranges).await },
100        )
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use std::sync::{Arc, Mutex};
108    use std::thread::ThreadId;
109
110    /// An in-memory [`AsyncFileReader`] that records the thread each request ran on
111    #[derive(Clone)]
112    struct InMemoryReader {
113        data: Bytes,
114        threads: Arc<Mutex<Vec<ThreadId>>>,
115    }
116
117    impl AsyncFileReader for InMemoryReader {
118        fn get_bytes(&mut self, range: Range<u64>) -> BoxFuture<'_, Result<Bytes, AvroError>> {
119            self.threads
120                .lock()
121                .unwrap()
122                .push(std::thread::current().id());
123            let data = self.data.slice(range.start as usize..range.end as usize);
124            futures::future::ready(Ok(data)).boxed()
125        }
126    }
127
128    #[tokio::test]
129    async fn test_spawned_reader() {
130        let rt = tokio::runtime::Builder::new_multi_thread()
131            .worker_threads(1)
132            .build()
133            .unwrap();
134
135        let inner = InMemoryReader {
136            data: Bytes::from_static(b"hello world"),
137            threads: Default::default(),
138        };
139        let threads = inner.threads.clone();
140        let mut reader = SpawnedReader::new(inner, rt.handle().clone());
141
142        let bytes = reader.get_bytes(0..5).await.unwrap();
143        assert_eq!(bytes.as_ref(), b"hello");
144
145        let ranges = reader.get_byte_ranges(vec![0..5, 6..11]).await.unwrap();
146        assert_eq!(ranges[1].as_ref(), b"world");
147
148        // All I/O must have run on the spawned runtime, not the current one
149        let current_id = std::thread::current().id();
150        let threads = threads.lock().unwrap();
151        assert!(!threads.is_empty());
152        assert!(threads.iter().all(|id| *id != current_id));
153
154        // Runtimes have to be dropped in blocking contexts
155        tokio::runtime::Handle::current().spawn_blocking(move || drop(rt));
156    }
157
158    #[tokio::test]
159    async fn test_spawned_reader_fails_on_shutdown_runtime() {
160        let rt = tokio::runtime::Builder::new_multi_thread()
161            .worker_threads(1)
162            .build()
163            .unwrap();
164
165        let inner = InMemoryReader {
166            data: Bytes::from_static(b"hello world"),
167            threads: Default::default(),
168        };
169        let mut reader = SpawnedReader::new(inner, rt.handle().clone());
170
171        rt.shutdown_background();
172
173        let err = reader.get_bytes(0..1).await.unwrap_err().to_string();
174        assert!(err.contains("was cancelled"), "{err}");
175    }
176}
arrow_avro/reader/async_reader/spawn.rs

arrow_avro/reader/async_reader/
spawn.rs
