third-party/sqlx-sqlite/src/connection/worker.rs

use std::future::Future;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::thread;

use futures_channel::oneshot;
use futures_intrusive::sync::{Mutex, MutexGuard};
use sqlx_core::sql_str::SqlStr;
use tracing::span::Span;

use sqlx_core::error::Error;
use sqlx_core::transaction::{
    begin_ansi_transaction_sql, commit_ansi_transaction_sql, rollback_ansi_transaction_sql,
};
use sqlx_core::Either;

use crate::connection::establish::EstablishParams;
use crate::connection::execute;
use crate::connection::{ConnectionState, SqliteWorkerConnection};
use crate::{SqliteArguments, SqliteQueryResult, SqliteRow, SqliteStatement};

#[cfg(feature = "deserialize")]
use crate::connection::deserialize::{deserialize, serialize, SchemaName, SqliteOwnedBuf};

// Each SQLite connection has a dedicated thread.

// TODO: Tweak this so that we can use a thread pool per pool of SQLite3 connections to reduce
//       OS resource usage. Low priority because a high concurrent load for SQLite3 is very
//       unlikely.

pub(crate) struct ConnectionWorker {
    command_tx: flume::Sender<(Command, tracing::Span)>,
    /// Mutex for locking access to the database.
    pub(crate) shared: Arc<WorkerSharedState>,
}

pub(crate) struct WorkerSharedState {
    transaction_depth: AtomicUsize,
    cached_statements_size: AtomicUsize,
    pub(crate) conn: Mutex<ConnectionState>,
}

impl WorkerSharedState {
    pub(crate) fn get_transaction_depth(&self) -> usize {
        self.transaction_depth.load(Ordering::Acquire)
    }

    pub(crate) fn get_cached_statements_size(&self) -> usize {
        self.cached_statements_size.load(Ordering::Acquire)
    }
}

enum Command {
    Call {
        operation: Box<dyn FnOnce(&mut ConnectionState) + Send + 'static>,
    },
    Prepare {
        query: SqlStr,
        tx: oneshot::Sender<Result<SqliteStatement, Error>>,
    },
    #[cfg(feature = "offline")]
    Describe {
        query: SqlStr,
        tx: oneshot::Sender<Result<sqlx_core::describe::Describe<crate::Sqlite>, Error>>,
    },
    Execute {
        query: SqlStr,
        arguments: Option<SqliteArguments>,
        persistent: bool,
        tx: flume::Sender<Result<Either<SqliteQueryResult, SqliteRow>, Error>>,
        limit: Option<usize>,
    },
    #[cfg(feature = "deserialize")]
    Serialize {
        schema: Option<SchemaName>,
        tx: oneshot::Sender<Result<SqliteOwnedBuf, Error>>,
    },
    #[cfg(feature = "deserialize")]
    Deserialize {
        schema: Option<SchemaName>,
        data: SqliteOwnedBuf,
        read_only: bool,
        tx: oneshot::Sender<Result<(), Error>>,
    },
    Begin {
        tx: rendezvous_oneshot::Sender<Result<(), Error>>,
        statement: Option<SqlStr>,
    },
    Commit {
        tx: rendezvous_oneshot::Sender<Result<(), Error>>,
    },
    Rollback {
        tx: Option<rendezvous_oneshot::Sender<Result<(), Error>>>,
    },
    UnlockDb,
    ClearCache {
        tx: oneshot::Sender<()>,
    },
    Ping {
        tx: oneshot::Sender<()>,
    },
    Shutdown {
        tx: oneshot::Sender<()>,
    },
}

impl ConnectionWorker {
    pub(crate) async fn establish(params: EstablishParams) -> Result<Self, Error> {
        let (establish_tx, establish_rx) = oneshot::channel();

        thread::Builder::new()
            .name(params.thread_name.clone())
            .spawn(move || {
                let (command_tx, command_rx) = flume::bounded(params.command_channel_size);

                let conn = match params.establish() {
                    Ok(conn) => conn,
                    Err(e) => {
                        establish_tx.send(Err(e)).ok();
                        return;
                    }
                };

                let shared = Arc::new(WorkerSharedState {
                    transaction_depth: AtomicUsize::new(0),
                    cached_statements_size: AtomicUsize::new(0),
                    // note: must be fair because in `Command::UnlockDb` we unlock the mutex
                    // and then immediately try to relock it; an unfair mutex would immediately
                    // grant us the lock even if another task is waiting.
                    conn: Mutex::new(conn, true),
                });
                let mut conn = shared.conn.try_lock().unwrap();

                if establish_tx
                    .send(Ok(Self {
                        command_tx,
                        shared: Arc::clone(&shared),
                    }))
                    .is_err()
                {
                    return;
                }

                // If COMMIT or ROLLBACK is processed but not acknowledged, there would be another
                // ROLLBACK sent when the `Transaction` drops. We need to ignore it otherwise we
                // would rollback an already completed transaction.
                let mut ignore_next_start_rollback = false;

                for (cmd, span) in command_rx {
                    let _guard = span.enter();
                    match cmd {
                        Command::Call { operation } => {
                            operation(&mut conn);
                            update_cached_statements_size(&conn, &shared.cached_statements_size);
                        }
                        Command::Prepare { query, tx } => {
                            tx.send(prepare(&mut conn, query)).ok();

                            // This may issue an unnecessary write on failure,
                            // but it doesn't matter in the grand scheme of things.
                            update_cached_statements_size(
                                &conn,
                                &shared.cached_statements_size,
                            );
                        }
                        #[cfg(feature = "offline")]
                        Command::Describe { query, tx } => {
                            tx.send(crate::connection::describe::describe(&mut conn, query)).ok();
                        }
                        Command::Execute {
                            query,
                            arguments,
                            persistent,
                            tx,
                            limit
                        } => {
                            let iter = match execute::iter(&mut conn, query, arguments, persistent)
                            {
                                Ok(iter) => iter,
                                Err(e) => {
                                    tx.send(Err(e)).ok();
                                    continue;
                                }
                            };

                            match limit {
                                None => {
                                    for res in iter {
                                        let has_error = res.is_err();
                                        if tx.send(res).is_err() || has_error {
                                            break;
                                        }
                                    }
                                },
                                Some(limit) => {
                                    let mut iter = iter;
                                    let mut rows_returned = 0;

                                    while let Some(res) = iter.next() {
                                        if let Ok(ok) = &res {
                                            if ok.is_right() {
                                                rows_returned += 1;
                                                if rows_returned >= limit {
                                                    drop(iter);
                                                    let _ = tx.send(res);
                                                    break;
                                                }
                                            }
                                        }
                                        let has_error = res.is_err();
                                        if tx.send(res).is_err() || has_error {
                                            break;
                                        }
                                    }
                                },
                            }

                            update_cached_statements_size(&conn, &shared.cached_statements_size);
                        }
                        Command::Begin { tx, statement } => {
                            let depth = shared.transaction_depth.load(Ordering::Acquire);

                            let is_custom_statement = statement.is_some();
                            let statement = match statement {
                                // custom `BEGIN` statements are not allowed if
                                // we're already in a transaction (we need to
                                // issue a `SAVEPOINT` instead)
                                Some(_) if depth > 0 => {
                                    if tx.blocking_send(Err(Error::InvalidSavePointStatement)).is_err() {
                                        break;
                                    }
                                    continue;
                                },
                                Some(statement) => statement,
                                None => begin_ansi_transaction_sql(depth),
                            };
                            let res =
                                conn.handle
                                    .exec(statement.as_str())
                                    .and_then(|res| {
                                        if is_custom_statement && !conn.handle.in_transaction() {
                                            return Err(Error::BeginFailed)
                                        }

                                        shared.transaction_depth.fetch_add(1, Ordering::Release);

                                        Ok(res)
                                    });
                            let res_ok = res.is_ok();

                            if tx.blocking_send(res).is_err() && res_ok {
                                // The BEGIN was processed but not acknowledged. This means no
                                // `Transaction` was created and so there is no way to commit /
                                // rollback this transaction. We need to roll it back
                                // immediately otherwise it would remain started forever.
                                if let Err(error) = conn
                                    .handle
                                    .exec(rollback_ansi_transaction_sql(depth + 1).as_str())
                                    .map(|_| {
                                        shared.transaction_depth.fetch_sub(1, Ordering::Release);
                                    })
                                {
                                    // The rollback failed. To prevent leaving the connection
                                    // in an inconsistent state we shutdown this worker which
                                    // causes any subsequent operation on the connection to fail.
                                    tracing::error!(%error, "failed to rollback cancelled transaction");
                                    break;
                                }
                            }
                        }
                        Command::Commit { tx } => {
                            let depth = shared.transaction_depth.load(Ordering::Acquire);

                            let res = if depth > 0 {
                                conn.handle
                                    .exec(commit_ansi_transaction_sql(depth).as_str())
                                    .map(|_| {
                                        shared.transaction_depth.fetch_sub(1, Ordering::Release);
                                    })
                            } else {
                                Ok(())
                            };
                            let res_ok = res.is_ok();

                            if tx.blocking_send(res).is_err() && res_ok {
                                // The COMMIT was processed but not acknowledged. This means that
                                // the `Transaction` doesn't know it was committed and will try to
                                // rollback on drop. We need to ignore that rollback.
                                ignore_next_start_rollback = true;
                            }
                        }
                        Command::Rollback { tx } => {
                            if ignore_next_start_rollback && tx.is_none() {
                                ignore_next_start_rollback = false;
                                continue;
                            }

                            let depth = shared.transaction_depth.load(Ordering::Acquire);

                            let res = if depth > 0 {
                                conn.handle
                                    .exec(rollback_ansi_transaction_sql(depth).as_str())
                                    .map(|_| {
                                        shared.transaction_depth.fetch_sub(1, Ordering::Release);
                                    })
                            } else {
                                Ok(())
                            };

                            let res_ok = res.is_ok();

                            if let Some(tx) = tx {
                                if tx.blocking_send(res).is_err() && res_ok {
                                    // The ROLLBACK was processed but not acknowledged. This means
                                    // that the `Transaction` doesn't know it was rolled back and
                                    // will try to rollback again on drop. We need to ignore that
                                    // rollback.
                                    ignore_next_start_rollback = true;
                                }
                            }
                        }
                        #[cfg(feature = "deserialize")]
                        Command::Serialize { schema, tx } => {
                            tx.send(serialize(&mut conn, schema)).ok();
                        }
                        #[cfg(feature = "deserialize")]
                        Command::Deserialize { schema, data, read_only, tx } => {
                            tx.send(deserialize(&mut conn, schema, data, read_only)).ok();
                        }
                        Command::ClearCache { tx } => {
                            conn.statements.clear();
                            update_cached_statements_size(&conn, &shared.cached_statements_size);
                            tx.send(()).ok();
                        }
                        Command::UnlockDb => {
                            drop(conn);
                            conn = futures_executor::block_on(shared.conn.lock());
                        }
                        Command::Ping { tx } => {
                            tx.send(()).ok();
                        }
                        Command::Shutdown { tx } => {
                            // drop the connection references before sending confirmation
                            // and ending the command loop
                            drop(conn);
                            drop(shared);
                            let _ = tx.send(());
                            return;
                        }
                    }
                }
            })?;

        establish_rx.await.map_err(|_| Error::WorkerCrashed)?
    }

    pub(crate) async fn run_with_connection<F, R>(&mut self, operation: F) -> Result<R, Error>
    where
        F: FnOnce(&mut SqliteWorkerConnection<'_>) -> R + Send + 'static,
        R: Send + 'static,
    {
        let outcome: std::thread::Result<R> = self
            .oneshot_cmd(move |reply| Command::Call {
                operation: Box::new(move |conn| {
                    let outcome =
                        std::panic::catch_unwind(std::panic::AssertUnwindSafe(move || {
                            operation(&mut SqliteWorkerConnection { conn })
                        }));
                    // Accepted callbacks run even after their observer disappears.
                    // Failed delivery drops the owned outcome on this worker.
                    let _ = reply.send(outcome);
                }),
            })
            .await?;
        match outcome {
            Ok(value) => Ok(value),
            Err(payload) => std::panic::resume_unwind(payload),
        }
    }

    pub(crate) async fn prepare(&mut self, query: SqlStr) -> Result<SqliteStatement, Error> {
        self.oneshot_cmd(|tx| Command::Prepare { query, tx })
            .await?
    }

    #[cfg(feature = "offline")]
    pub(crate) async fn describe(
        &mut self,
        query: SqlStr,
    ) -> Result<sqlx_core::describe::Describe<crate::Sqlite>, Error> {
        self.oneshot_cmd(|tx| Command::Describe { query, tx })
            .await?
    }

    pub(crate) async fn execute(
        &mut self,
        query: SqlStr,
        args: Option<SqliteArguments>,
        chan_size: usize,
        persistent: bool,
        limit: Option<usize>,
    ) -> Result<flume::Receiver<Result<Either<SqliteQueryResult, SqliteRow>, Error>>, Error> {
        let (tx, rx) = flume::bounded(chan_size);

        self.command_tx
            .send_async((
                Command::Execute {
                    query,
                    arguments: args,
                    persistent,
                    tx,
                    limit,
                },
                Span::current(),
            ))
            .await
            .map_err(|_| Error::WorkerCrashed)?;

        Ok(rx)
    }

    pub(crate) async fn begin(&mut self, statement: Option<SqlStr>) -> Result<(), Error> {
        self.oneshot_cmd_with_ack(|tx| Command::Begin { tx, statement })
            .await?
    }

    pub(crate) async fn commit(&mut self) -> Result<(), Error> {
        self.oneshot_cmd_with_ack(|tx| Command::Commit { tx })
            .await?
    }

    pub(crate) async fn rollback(&mut self) -> Result<(), Error> {
        self.oneshot_cmd_with_ack(|tx| Command::Rollback { tx: Some(tx) })
            .await?
    }

    pub(crate) fn start_rollback(&mut self) -> Result<(), Error> {
        self.command_tx
            .send((Command::Rollback { tx: None }, Span::current()))
            .map_err(|_| Error::WorkerCrashed)
    }

    pub(crate) async fn ping(&mut self) -> Result<(), Error> {
        self.oneshot_cmd(|tx| Command::Ping { tx }).await
    }

    #[cfg(feature = "deserialize")]
    pub(crate) async fn deserialize(
        &mut self,
        schema: Option<SchemaName>,
        data: SqliteOwnedBuf,
        read_only: bool,
    ) -> Result<(), Error> {
        self.oneshot_cmd(|tx| Command::Deserialize {
            schema,
            data,
            read_only,
            tx,
        })
        .await?
    }

    #[cfg(feature = "deserialize")]
    pub(crate) async fn serialize(
        &mut self,
        schema: Option<SchemaName>,
    ) -> Result<SqliteOwnedBuf, Error> {
        self.oneshot_cmd(|tx| Command::Serialize { schema, tx })
            .await?
    }

    async fn oneshot_cmd<F, T>(&mut self, command: F) -> Result<T, Error>
    where
        F: FnOnce(oneshot::Sender<T>) -> Command,
    {
        let (tx, rx) = oneshot::channel();

        self.command_tx
            .send_async((command(tx), Span::current()))
            .await
            .map_err(|_| Error::WorkerCrashed)?;

        rx.await.map_err(|_| Error::WorkerCrashed)
    }

    async fn oneshot_cmd_with_ack<F, T>(&mut self, command: F) -> Result<T, Error>
    where
        F: FnOnce(rendezvous_oneshot::Sender<T>) -> Command,
    {
        let (tx, rx) = rendezvous_oneshot::channel();

        self.command_tx
            .send_async((command(tx), Span::current()))
            .await
            .map_err(|_| Error::WorkerCrashed)?;

        rx.recv().await.map_err(|_| Error::WorkerCrashed)
    }

    pub(crate) async fn clear_cache(&mut self) -> Result<(), Error> {
        self.oneshot_cmd(|tx| Command::ClearCache { tx }).await
    }

    pub(crate) async fn unlock_db(&mut self) -> Result<MutexGuard<'_, ConnectionState>, Error> {
        let (guard, res) = futures_util::future::join(
            // we need to join the wait queue for the lock before we send the message
            self.shared.conn.lock(),
            self.command_tx
                .send_async((Command::UnlockDb, Span::current())),
        )
        .await;

        res.map_err(|_| Error::WorkerCrashed)?;

        Ok(guard)
    }

    /// Send a command to the worker to shut down the processing thread.
    ///
    /// A `WorkerCrashed` error may be returned if the thread has already stopped.
    pub(crate) fn shutdown(&mut self) -> impl Future<Output = Result<(), Error>> {
        let (tx, rx) = oneshot::channel();

        let send_res = self
            .command_tx
            .send((Command::Shutdown { tx }, Span::current()))
            .map_err(|_| Error::WorkerCrashed);

        async move {
            send_res?;

            // wait for the response
            rx.await.map_err(|_| Error::WorkerCrashed)
        }
    }
}

pub(super) fn prepare(conn: &mut ConnectionState, query: SqlStr) -> Result<SqliteStatement, Error> {
    // prepare statement object (or checkout from cache)
    let statement = conn.statements.get(query.as_str(), true)?;

    let mut parameters = 0;
    let mut columns = None;
    let mut column_names = None;

    while let Some(statement) = statement.prepare_next(&mut conn.handle)? {
        parameters += statement.handle.bind_parameter_count();

        // the first non-empty statement is chosen as the statement we pull columns from
        if !statement.columns.is_empty() && columns.is_none() {
            columns = Some(Arc::clone(statement.columns));
            column_names = Some(Arc::clone(statement.column_names));
        }
    }

    Ok(SqliteStatement {
        sql: query,
        columns: columns.unwrap_or_default(),
        column_names: column_names.unwrap_or_default(),
        parameters,
    })
}

fn update_cached_statements_size(conn: &ConnectionState, size: &AtomicUsize) {
    size.store(conn.statements.len(), Ordering::Release);
}

// A oneshot channel where send completes only after the receiver receives the value.
mod rendezvous_oneshot {
    use super::oneshot::{self, Canceled};

    pub fn channel<T>() -> (Sender<T>, Receiver<T>) {
        let (inner_tx, inner_rx) = oneshot::channel();
        (Sender { inner: inner_tx }, Receiver { inner: inner_rx })
    }

    pub struct Sender<T> {
        inner: oneshot::Sender<(T, oneshot::Sender<()>)>,
    }

    impl<T> Sender<T> {
        pub async fn send(self, value: T) -> Result<(), Canceled> {
            let (ack_tx, ack_rx) = oneshot::channel();
            self.inner.send((value, ack_tx)).map_err(|_| Canceled)?;
            ack_rx.await
        }

        pub fn blocking_send(self, value: T) -> Result<(), Canceled> {
            futures_executor::block_on(self.send(value))
        }
    }

    pub struct Receiver<T> {
        inner: oneshot::Receiver<(T, oneshot::Sender<()>)>,
    }

    impl<T> Receiver<T> {
        pub async fn recv(self) -> Result<T, Canceled> {
            let (value, ack_tx) = self.inner.await?;
            ack_tx.send(()).map_err(|_| Canceled)?;
            Ok(value)
        }
    }
}