odbc_api/cursor/concurrent_block_cursor.rs
1use std::{
2 mem::swap,
3 sync::mpsc::{sync_channel, Receiver, SyncSender},
4 thread::{self, JoinHandle},
5};
6
7use crate::{BlockCursor, Cursor, Error};
8
9use super::RowSetBuffer;
10
11/// A wrapper around block cursors which fetches data in a dedicated system thread. Intended to
12/// fetch data batch by batch while the application processes the batch last fetched. Works best
13/// with a double buffer strategy using two fetch buffers.
14///
15/// # Example
16///
17/// ```no_run
18/// use odbc_api::{
19/// Environment, buffers::{ColumnarAnyBuffer, BufferDesc}, Cursor, ConcurrentBlockCursor
20/// };
21/// use std::sync::OnceLock;
22///
23/// // We want to use the ODBC environment from another system thread without scope => Therefore it
24/// // needs to be static.
25/// static ENV: OnceLock<Environment> = OnceLock::new();
26/// let env = Environment::new()?;
27///
28/// let conn = ENV.get_or_init(|| env).connect_with_connection_string(
29/// "Driver={ODBC Driver 18 for SQL Server};Server=localhost;UID=SA;PWD=My@Test@Password1;",
30/// Default::default())?;
31///
32/// // We must use into_cursor to create a statement handle with static lifetime, which also owns
33/// // the connection. This way we can send it to another thread safely.
34/// let cursor = conn.into_cursor("SELECT * FROM very_big_table", ())?.unwrap();
35///
36/// // Batch size and buffer description. Here we assume there is only one integer column
37/// let buffer_a = ColumnarAnyBuffer::from_descs(1000, [BufferDesc::I32 { nullable: false }]);
38/// let mut buffer_b = ColumnarAnyBuffer::from_descs(1000, [BufferDesc::I32 { nullable: false }]);
39/// // And now we have a sendable block cursor with static lifetime
40/// let block_cursor = cursor.bind_buffer(buffer_a)?;
41///
42/// let mut cbc = ConcurrentBlockCursor::from_block_cursor(block_cursor);
43/// while cbc.fetch_into(&mut buffer_b)? {
44/// // Proccess batch in buffer b asynchronously to fetching it
45/// }
46///
47/// # Ok::<_, odbc_api::Error>(())
48/// ```
49pub struct ConcurrentBlockCursor<C, B> {
50 /// In order to avoid reallocating buffers over and over again, we use this channel to send the
51 /// buffers back to the fetch thread after we copied their contents into arrow arrays.
52 send_buffer: SyncSender<B>,
53 /// Receives filled batches from the fetch thread. Once the source is empty or if an error
54 /// occurs its associated sender is dropped, and receiving batches will return an error (which
55 /// we expect during normal operation and cleanup, and is not forwarded to the user).
56 receive_batch: Receiver<B>,
57 /// We join with the fetch thread if we stop receiving batches (i.e. receive_batch.recv()
58 /// returns an error) or `into_cursor` is called. `None` if the thread has already been joined.
59 /// In this case either an error has been reported to the user, or the cursor is stored in
60 /// `cursor`.
61 fetch_thread: Option<JoinHandle<Result<C, Error>>>,
62 /// Only `Some`, if the cursor has been consumed succesfully and `fetch_thread` has been joined.
63 /// Can only be `Some` if `fetch_thread` is `None`. If both `fetch_thread` and `cursor` are
64 /// `None`, it is implied that `fetch_thread` returned an error joining.
65 cursor: Option<C>,
66}
67
68impl<C, B> ConcurrentBlockCursor<C, B>
69where
70 C: Cursor + Send + 'static,
71 B: RowSetBuffer + Send + 'static,
72{
73 /// Construct a new concurrent block cursor.
74 ///
75 /// # Parameters
76 ///
77 /// * `block_cursor`: Taking a BlockCursor instead of a Cursor allows for better resource
78 /// stealing if constructing starting from a sequential Cursor, as we do not need to undbind
79 /// and bind the cursor.
80 pub fn from_block_cursor(block_cursor: BlockCursor<C, B>) -> Self {
81 let (send_buffer, receive_buffer) = sync_channel(1);
82 let (send_batch, receive_batch) = sync_channel(1);
83
84 let fetch_thread = thread::spawn(move || {
85 let mut block_cursor = block_cursor;
86 loop {
87 match block_cursor.fetch_with_truncation_check(true) {
88 Ok(Some(_batch)) => (),
89 Ok(None) => {
90 break block_cursor
91 .unbind()
92 .map(|(undbound_cursor, _buffer)| undbound_cursor);
93 }
94 Err(odbc_error) => {
95 drop(send_batch);
96 break Err(odbc_error);
97 }
98 }
99 // There has been another row group fetched by the cursor. We unbind the buffers so
100 // we can pass ownership of it to the application and bind a new buffer to the
101 // cursor in order to start fetching the next batch.
102 let (cursor, buffer) = block_cursor.unbind()?;
103 if send_batch.send(buffer).is_err() {
104 // Should the main thread stop receiving buffers, this thread should
105 // also stop fetching batches.
106 break Ok(cursor);
107 }
108 // Wait for the application thread to give us a buffer to fill.
109 match receive_buffer.recv() {
110 Err(_) => {
111 // Application thread dropped sender and does not want more buffers to be
112 // filled. Let's stop this thread and return the cursor
113 break Ok(cursor);
114 }
115 Ok(next_buffer) => {
116 block_cursor = cursor.bind_buffer(next_buffer).unwrap();
117 }
118 }
119 }
120 });
121
122 Self {
123 send_buffer,
124 receive_batch,
125 fetch_thread: Some(fetch_thread),
126 cursor: None,
127 }
128 }
129
130 /// Join fetch thread and yield the cursor back.
131 pub fn into_cursor(self) -> Result<C, Error> {
132 drop(self.receive_batch);
133 // Dropping the send buffer is necessary to avoid deadlocks, in case there would not be any
134 // buffer in the channel waiting for the fetch thread. Since we consume the cursor here, it
135 // is also impossible for the application to send another buffer.
136 drop(self.send_buffer);
137 if let Some(cursor) = self.cursor {
138 Ok(cursor)
139 } else {
140 self.fetch_thread.unwrap().join().unwrap()
141 }
142 }
143}
144
145impl<C, B> ConcurrentBlockCursor<C, B> {
146 /// Receive the current batch and take ownership of its buffer. `None` if the cursor is already
147 /// consumed, or had an error previously. This method blocks until a new batch available. In
148 /// order for new batches available new buffers must be send to the thread in order for it to
149 /// fill them. So calling fetch repeatedly without calling [`Self::fill`] in between may
150 /// deadlock.
151 pub fn fetch(&mut self) -> Result<Option<B>, Error> {
152 match self.receive_batch.recv() {
153 // We successfully fetched a batch from the database.
154 Ok(batch) => Ok(Some(batch)),
155 // Fetch thread stopped sending batches. Either because we consumed the result set
156 // completly or we hit an error.
157 Err(_receive_error) => {
158 if let Some(join_handle) = self.fetch_thread.take() {
159 // If there has been an error returning the batch, or unbinding the buffer `?`
160 // will raise it.
161 self.cursor = Some(join_handle.join().unwrap()?);
162 // We ran out of batches in the result set. End the stream.
163 Ok(None)
164 } else {
165 // This only happen if this method is called after it returned either `false` or
166 // `Err` once. Let us treat this scenario like a result set which is consumed
167 // completly.
168 Ok(None)
169 }
170 }
171 }
172 }
173
174 /// Send a buffer to the thread fetching in order for it to be filled and to be retrieved later
175 /// using either `fetch`, or `fetch_into`.
176 pub fn fill(&mut self, buffer: B) {
177 let _ = self.send_buffer.send(buffer);
178 }
179
180 /// Fetches values from the ODBC datasource into buffer. Values are streamed batch by batch in
181 /// order to avoid reallocation of the buffers used for tranistion. This call blocks until a new
182 /// batch is ready. This method combines both [`Self::fetch`] and [`Self::fill`].
183 ///
184 /// # Parameters
185 ///
186 /// * `buffer`: A columnar any buffer which can bind to the cursor wrapped by this instance.
187 /// After the method call the reference will not point to the same instance which had been
188 /// passed into the function call, but to the one which was bound to the cursor in order to
189 /// fetch the last batch. The buffer passed into this method, is then used to fetch the next
190 /// batch. As such this method is ideal to implement concurrent fetching using two buffers.
191 /// One which is written to, and one that is read, which flip their roles between batches.
192 /// Also called double buffering.
193 ///
194 /// # Return
195 ///
196 /// * `true`: Fetched a batch from the data source. The contents of that batch are now in
197 /// `buffer`.
198 /// * `false`: No batch could be fetched. The result set is consumed completly.
199 pub fn fetch_into(&mut self, buffer: &mut B) -> Result<bool, Error> {
200 if let Some(mut batch) = self.fetch()? {
201 swap(buffer, &mut batch);
202 self.fill(batch);
203 Ok(true)
204 } else {
205 Ok(false)
206 }
207 }
208}