sqlite: use Weak poiter to StatementHandle in the worker

Otherwise some tests fail to close connection.
launchbadge · Apr 23, 2021 · 5461eee · 5461eee
1 parent 5eebc05
commit 5461eee
Showing 1 changed file with 14 additions and 9 deletions.
diff --git a/sqlx-core/src/sqlite/statement/worker.rs b/sqlx-core/src/sqlite/statement/worker.rs
@@ -4,7 +4,7 @@ use crossbeam_channel::{unbounded, Sender};
 use either::Either;
 use futures_channel::oneshot;
 use libsqlite3_sys::{sqlite3_step, SQLITE_DONE, SQLITE_ROW};
-use std::sync::Arc;
+use std::sync::{Arc, Weak};
 use std::thread;
 
 // Each SQLite connection has a dedicated thread.
@@ -19,7 +19,7 @@ pub(crate) struct StatementWorker {
 
 enum StatementWorkerCommand {
     Step {
-        statement: Arc<StatementHandle>,
+        statement: Weak<StatementHandle>,
         tx: oneshot::Sender<Result<Either<u64, ()>, Error>>,
     },
 }
@@ -32,14 +32,19 @@ impl StatementWorker {
             for cmd in rx {
                 match cmd {
                     StatementWorkerCommand::Step { statement, tx } => {
-                        let status = unsafe { sqlite3_step(statement.0.as_ptr()) };
+                        let resp = if let Some(statement) = statement.upgrade() {
+                            let status = unsafe { sqlite3_step(statement.0.as_ptr()) };
 
-                        let resp = match status {
-                            SQLITE_ROW => Ok(Either::Right(())),
-                            SQLITE_DONE => Ok(Either::Left(statement.changes())),
-                            _ => Err(statement.last_error().into()),
+                            let resp = match status {
+                                SQLITE_ROW => Ok(Either::Right(())),
+                                SQLITE_DONE => Ok(Either::Left(statement.changes())),
+                                _ => Err(statement.last_error().into()),
+                            };
+                            resp
+                        } else {
+                            // Statement is already finalized.
+                            Err(Error::WorkerCrashed)
                         };
-
                         let _ = tx.send(resp);
                     }
                 }
@@ -57,7 +62,7 @@ impl StatementWorker {
 
         self.tx
             .send(StatementWorkerCommand::Step {
-                statement: Arc::clone(statement),
+                statement: Arc::downgrade(statement),
                 tx,
             })
             .map_err(|_| Error::WorkerCrashed)?;