SQLite's StatementWorker consumes 100% CPU inside its busy loop

[Patryk27]

Hi :-)

First of all, the code:

[dependencies]
sqlx = { version = "0.4.0-beta.1", default-features = false, features = ["runtime-tokio", "sqlite"] }
tokio = { version = "0.2", features = ["full"] }

use sqlx::Connection;

#[tokio::main]
async fn main() {
    let mut db = sqlx::SqliteConnection::connect(":memory:")
        .await
        .unwrap();

    let value: Option<(i64,)> = sqlx::query_as("SELECT 1")
        .fetch_optional(&mut db)
        .await
        .unwrap();

    // Keep the application alive
    loop {
        tokio::time::delay_for(tokio::time::Duration::from_secs(1)).await;
    }
}

This code makes the application instantaneously use 100% CPU.

Curious about it, I've conducted an investigation that led me here first:

impl<'c> Executor<'c> for &'c mut SqliteConnection {
    /* ... */

    // sqlx-core/src/sqlite/connection/executor.rs:157
    fn fetch_optional<'e, 'q: 'e, E: 'q>(
        self,
        query: E,
    ) -> BoxFuture<'e, Result<Option<SqliteRow>, Error>>
    where
        'c: 'e,
        E: Execute<'q, Self::Database>,
    {
        let mut s = self.fetch_many(query);

        Box::pin(async move {
            while let Some(v) = s.try_next().await? {
                if let Either::Right(r) = v {
                    return Ok(Some(r));
                }
            }

            Ok(None)
        })
    }

    /* ... */
}

... and later there:

impl StatementWorker {
    pub(crate) fn new() -> Self {
        /* ... */

        let handle = spawn({
            /* ... */

            move || {
                /* ... */

                // sqlx-core/src/sqlite/statement/worker.rs:49
                'run: while status.load(Ordering::Acquire) >= 0 {
                    'statement: loop {
                        match status.load(Ordering::Acquire) {
                            STATE_CLOSE => {
                                /* ... */
                            }

                            STATE_READY => {
                                /* ... */
                            }

                            _ => {
                                // waits for the receiving end to be ready to receive the rows
                                // this should take less than 1 microsecond under most conditions
                                spin_loop_hint();
                            }
                        }
                    }
                }
            }
        });

        /* ... */
    }

    /* ... */
}

The issue is that the current implementation of fetch_optional() retrieves only the first row from the stream and then silently discards the stream, without letting StatementWorker know about it. The worker then starts entering spin_loop_hint() millions times per second, waiting for response from a stream that doesn't exist anymore.

This forever-busy-loop can be reproduced even without referring to fetch_optional(), like so:

async fn main() {
    /* ... */

    let mut stream = sqlx::query_as::<_, (Option<i64>,)>("SELECT 1")
        .fetch_many(&mut db);

    // Retrieve just the first row and then discard rest of the stream
    while let Some(v) = s.try_next().await.unwrap() {
        break;
    }

    // Keep the application alive
    loop {
        tokio::time::delay_for(tokio::time::Duration::from_secs(1)).await;
    }
}

As for a possible solution, I'd suggest introducing a new type, called ExecutionToken:

pub(crate) struct ExecutionToken {
    status: Arc<AtomicI32>,
}

impl Drop for ExecutionToken {
    fn drop(&mut self) {
        self.status.store(STATE_READY, Ordering::Release);
    }
}

/* ... */

impl StatementWorker {
    pub(crate) fn execute(&self, statement: &StatementHandle) -> ExecutionToken {
        let token = ExecutionToken {
            status: Arc::clone(&self.status),
        };

        self.statement
            .store(statement.0.as_ptr(), Ordering::Release);

        token
    }
}

/* ... */

// and then later just:
// let _token = worker.execute(handle);

Since I'm proficient in neither sqlx's nor SQLite's internals, I'm not sure if that's the correct approach (although it works on my machine ™); if you think that's alright, I could prepare a merge request :-)

Thank you for your great work on this library,
Cheers.

[Patryk27]

I'm having second thoughts though: while using ExecutionToken does fix fetch_optional(), following code still causes 100% CPU usage, for the same reason:

async fn main() {
    /* ... */

    let mut stream = sqlx::query_as::<_, (i64,)>("SELECT 1")
        .fetch_many(&mut db);

    // Fetch first item to initialize the stream
    stream.try_next().await.unwrap();

    tokio::time::delay_for(tokio::time::Duration::from_secs(10)).await;

    drop(stream);
}

Maybe instead of a busy-wait loop with shared state we could use an async-aware mpsc channel or condvar?

[Patryk27]

Okie, I've managed to get a correct proof-of-concept working:

// worker.rs

use crate::error::Error;
use crate::sqlite::statement::StatementHandle;
use crossbeam_channel::{bounded, unbounded, Sender};
use either::Either;
use libsqlite3_sys::{sqlite3_step, SQLITE_DONE, SQLITE_ROW};
use sqlx_rt::yield_now;
use std::thread;

pub(crate) struct StatementWorker {
    tx: Sender<StatementWorkerCommand>,
}

enum StatementWorkerCommand {
    Step {
        statement: StatementHandle,
        tx: Sender<Result<Either<u64, ()>, Error>>,
    },
}

impl StatementWorker {
    pub(crate) fn new() -> Self {
        let (tx, rx) = unbounded();

        thread::spawn(move || {
            for cmd in rx.into_iter() {
                match cmd {
                    StatementWorkerCommand::Step { statement, tx } => {
                        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 _ = tx.send(resp);
                    }
                }
            }
        });

        Self { tx }
    }

    pub(crate) async fn step(
        &mut self,
        statement: &StatementHandle,
    ) -> Result<Either<u64, ()>, Error> {
        let (tx, rx) = bounded(1);

        self.tx
            .send(StatementWorkerCommand::Step {
                statement: *statement,
                tx,
            })
            .unwrap();

        while rx.is_empty() {
            yield_now().await;
        }

        rx.recv().unwrap()
    }
}

In terms of performance, this version is slightly slower than the original one, but - thanks to crossbeam's channels - not much: I've done a simple benchmark that launches the same query 100k times, and the original code took 350ms to run, while this one took 420ms, so I think it's in the ballpark of "alright" when it comes to SQLite.

Whaddya think?

[meteficha]

I'm being bit by this, too, and definitely wouldn't mind a 0.7 microsecond overhead per query. (BTW, thanks for investigating this, @Patryk27)

[mehcode]

Thanks a lot for digging into this. StatementWorker is definitely pretty experimental and if we can keep the perf similar I'm totally open to re-arranging it. What you suggest seems pretty good at first glance.

Are you open to sending in a PR with the changes?

[Patryk27]

Yeah, totally - I'll try to submit the patch today :-)