rbdc_oracle/connection/
worker.rs1use std::sync::Arc;
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::thread;
4
5use crate::connection::ConnectionState;
6use crate::connection::establish::EstablishParams;
7use crate::connection::execute;
8use crate::{OracleArguments, OracleQueryResult, OracleRow, OracleStatement};
9use crossfire::{AsyncTx, spsc};
10use either::Either;
11use futures_channel::oneshot;
12use rbdc::Error;
13use std::sync::Mutex;
14
15pub(crate) struct ConnectionWorker {
16 command_tx: AsyncTx<crossfire::spsc::Array<Command>>,
17 pub(crate) shared: Arc<WorkerSharedState>,
18}
19
20pub(crate) struct WorkerSharedState {
21 pub(crate) cached_statements_size: AtomicUsize,
22 pub(crate) conn: Mutex<ConnectionState>,
23}
24
25pub enum Command {
26 Prepare {
27 query: Box<str>,
28 tx: oneshot::Sender<Result<OracleStatement, Error>>,
29 },
30 Execute {
31 query: Box<str>,
32 arguments: Option<OracleArguments>,
33 persistent: bool,
34 tx: crossfire::Tx<
35 crossfire::spsc::Array<Result<Either<OracleQueryResult, OracleRow>, Error>>,
36 >,
37 },
38 ClearCache {
39 tx: oneshot::Sender<()>,
40 },
41 Ping {
42 tx: oneshot::Sender<Result<(), Error>>,
43 },
44 Shutdown {
45 tx: oneshot::Sender<()>,
46 },
47}
48
49impl ConnectionWorker {
50 pub(crate) async fn establish(params: EstablishParams) -> Result<Self, Error> {
51 let (establish_tx, establish_rx) = oneshot::channel();
52
53 thread::Builder::new()
54 .name(params.thread_name.clone())
55 .spawn(move || {
56 let (command_tx, command_rx) =
57 spsc::bounded_async_blocking(params.command_channel_size);
58
59 let conn = match params.establish() {
60 Ok(conn) => conn,
61 Err(e) => {
62 establish_tx.send(Err(e)).ok();
63 return;
64 }
65 };
66
67 let shared = Arc::new(WorkerSharedState {
68 cached_statements_size: AtomicUsize::new(0),
69 conn: Mutex::new(conn),
70 });
71 let mut conn = shared.conn.lock().unwrap();
72
73 if establish_tx
74 .send(Ok(Self {
75 command_tx,
76 shared: Arc::clone(&shared),
77 }))
78 .is_err()
79 {
80 return;
81 }
82
83 loop {
84 let cmd = match command_rx.recv() {
85 Ok(cmd) => cmd,
86 Err(_) => break,
87 };
88
89 match cmd {
90 Command::Prepare { query, tx } => {
91 tx.send(prepare(&mut conn, &query).map(|prepared| {
92 update_cached_statements_size(
93 &conn,
94 &shared.cached_statements_size,
95 );
96 prepared
97 }))
98 .ok();
99 }
100 Command::Execute {
101 query,
102 arguments,
103 persistent,
104 tx,
105 } => {
106 let iter = match execute::iter(&mut conn, &query, arguments, persistent)
107 {
108 Ok(iter) => iter,
109 Err(e) => {
110 tx.send(Err(e)).ok();
111 continue;
112 }
113 };
114
115 for res in iter {
116 if tx.send(res).is_err() {
117 break;
118 }
119 }
120
121 update_cached_statements_size(&conn, &shared.cached_statements_size);
122 }
123 Command::ClearCache { tx } => {
124 conn.statements.clear();
125 update_cached_statements_size(&conn, &shared.cached_statements_size);
126 tx.send(()).ok();
127 }
128 Command::Ping { tx } => {
129 let result = conn
130 .handle
131 .connection()
132 .ping()
133 .map_err(|e| Error::from(e.to_string()));
134 let should_stop = result.is_err();
135 tx.send(result).ok();
136 if should_stop {
137 return;
138 }
139 }
140 Command::Shutdown { tx } => {
141 let _ = conn.handle.connection().commit();
142 let _ = conn.handle.connection().close();
143 drop(conn);
144 drop(shared);
145 let _ = tx.send(());
146 return;
147 }
148 }
149 }
150 })
151 .map_err(|e| Error::from(e.to_string()))?;
152
153 establish_rx
154 .await
155 .map_err(|_| Error::from("WorkerCrashed"))?
156 }
157
158 pub(crate) async fn prepare(&mut self, query: &str) -> Result<OracleStatement, Error> {
159 self.oneshot_cmd(|tx| Command::Prepare {
160 query: query.into(),
161 tx,
162 })
163 .await?
164 }
165
166 pub(crate) async fn execute(
167 &mut self,
168 query: String,
169 args: Option<OracleArguments>,
170 chan_size: usize,
171 persistent: bool,
172 ) -> Result<
173 crossfire::AsyncRx<
174 crossfire::spsc::Array<Result<Either<OracleQueryResult, OracleRow>, Error>>,
175 >,
176 Error,
177 > {
178 let (tx, rx) = spsc::bounded_blocking_async(chan_size);
179
180 self.command_tx
181 .send(Command::Execute {
182 query: query.into(),
183 arguments: args.map(OracleArguments::into_static),
184 persistent,
185 tx,
186 })
187 .await
188 .map_err(|_| Error::from("WorkerCrashed"))?;
189
190 Ok(rx)
191 }
192
193 pub(crate) async fn ping(&mut self) -> Result<(), Error> {
194 self.oneshot_cmd(|tx| Command::Ping { tx }).await?
195 }
196
197 pub(crate) async fn oneshot_cmd<F, T>(&mut self, command: F) -> Result<T, Error>
198 where
199 F: FnOnce(oneshot::Sender<T>) -> Command,
200 {
201 let (tx, rx) = oneshot::channel();
202
203 self.command_tx
204 .send(command(tx))
205 .await
206 .map_err(|_| Error::from("WorkerCrashed"))?;
207
208 rx.await.map_err(|_| Error::from("WorkerCrashed"))
209 }
210
211 pub(crate) async fn clear_cache(&mut self) -> Result<(), Error> {
212 self.oneshot_cmd(|tx| Command::ClearCache { tx }).await
213 }
214
215 pub(crate) async fn shutdown(&mut self) -> Result<(), Error> {
216 let (tx, rx) = oneshot::channel();
217
218 self.command_tx
219 .send(Command::Shutdown { tx })
220 .await
221 .map_err(|_| Error::from("WorkerCrashed"))?;
222
223 rx.await.map_err(|_| Error::from("WorkerCrashed"))
224 }
225}
226
227fn update_cached_statements_size(conn: &ConnectionState, size: &AtomicUsize) {
228 size.store(conn.statements.len(), Ordering::Release);
229}
230
231fn prepare(conn: &mut ConnectionState, query: &str) -> Result<OracleStatement, Error> {
232 super::executor::prepare(conn, query)
233}