Prevent concurrent reuse of SQLite statements

Author: mangod12

Lib/test/test_sqlite3/test_dbapi.py

@@ -26,6 +26,7 @@
 import sqlite3 as sqlite
 import subprocess
 import sys
+import textwrap
 import threading
 import unittest
 import urllib.parse
@@ -1734,6 +1735,45 @@ def run(con, err):
             t.join()
             self.assertEqual(len(err), 0, "\n".join(err))
 
+    @requires_subprocess()
+    def test_concurrent_queries_on_shared_connection(self):
+        code = textwrap.dedent("""\
+            import hashlib
+            import os
+            import sqlite3
+            import threading
+            from concurrent.futures import ThreadPoolExecutor
+
+            connection = sqlite3.connect(":memory:", check_same_thread=False)
+            connection.execute("CREATE TABLE entries (key INTEGER, value BLOB)")
+            data = os.urandom(1_400_000)
+            digest = hashlib.blake2b(data).digest()
+            keys = range(4)
+            connection.executemany(
+                "INSERT INTO entries VALUES (?, ?)", ((key, data) for key in keys)
+            )
+            barrier = threading.Barrier(len(keys))
+
+            def load(key):
+                barrier.wait()
+                for _ in range(100):
+                    row = connection.execute(
+                        "SELECT value FROM entries WHERE key = ?", (key,)
+                    ).fetchone()
+                    assert row is not None
+                    assert hashlib.blake2b(row[0]).digest() == digest
+
+            with ThreadPoolExecutor(max_workers=len(keys)) as executor:
+                list(executor.map(load, keys))
+        """)
+        proc = subprocess.run(
+            [sys.executable, "-X", "faulthandler", "-c", code],
+            capture_output=True,
+            text=True,
+            timeout=SHORT_TIMEOUT,
+        )
+        self.assertEqual(proc.returncode, 0, proc.stderr)
+
 
 class ConstructorTests(unittest.TestCase):
     def test_date(self):

Misc/NEWS.d/next/Library/gh-160699.bugfix.rst

@@ -0,0 +1 @@
+Fix a crash when multiple threads execute queries on the same SQLite connection.

Modules/_sqlite/cursor.c

@@ -141,14 +141,18 @@ pysqlite_cursor_init_impl(pysqlite_Cursor *self,
 }
 
 static inline int
-stmt_reset(pysqlite_Statement *self)
+stmt_reset(pysqlite_Statement *self, int keep_in_use)
 {
     int rc = SQLITE_OK;
 
     if (self->st != NULL) {
         Py_BEGIN_ALLOW_THREADS
         rc = sqlite3_reset(self->st);
         Py_END_ALLOW_THREADS
+
+        if (rc == SQLITE_OK && !keep_in_use) {
+            pysqlite_statement_mark_in_use(self, 0);
+        }
     }
 
     return rc;
@@ -179,7 +183,7 @@ cursor_clear(PyObject *op)
     pysqlite_set_ref(op, &self->row_factory, NULL);
     if (self->statement) {
         /* Reset the statement if the user has not closed the cursor */
-        int rc = stmt_reset(self->statement);
+        int rc = stmt_reset(self->statement, 0);
         Py_CLEAR(self->statement);
         if (rc != SQLITE_OK) {
             cursor_cannot_reset_stmt_error(self, 0);
@@ -854,7 +858,7 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
 
     if (self->statement) {
         // Reset pending statements on this cursor.
-        if (stmt_reset(self->statement) != SQLITE_OK) {
+        if (stmt_reset(self->statement, 0) != SQLITE_OK) {
             goto reset_failure;
         }
     }
@@ -872,15 +876,18 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
         goto error;
     }
 
-    if (sqlite3_stmt_busy(self->statement->st)) {
+    /* Reserve cached statements before releasing the GIL. SQLite only marks
+       a statement busy once sqlite3_step() starts, leaving a race window. */
+    if (!pysqlite_statement_try_mark_in_use(self->statement)) {
         Py_SETREF(self->statement,
                   pysqlite_statement_create(self->connection, operation));
         if (self->statement == NULL) {
             goto error;
         }
+        pysqlite_statement_mark_in_use(self->statement, 1);
     }
 
-    if (stmt_reset(self->statement) != SQLITE_OK) {
+    if (stmt_reset(self->statement, 1) != SQLITE_OK) {
         goto reset_failure;
     }
     self->rowcount = self->statement->is_dml ? 0L : -1L;
@@ -966,7 +973,7 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
             if (self->statement->is_dml) {
                 self->rowcount += (long)sqlite3_changes(self->connection->db);
             }
-            if (stmt_reset(self->statement) != SQLITE_OK) {
+            if (stmt_reset(self->statement, 0) != SQLITE_OK) {
                 goto reset_failure;
             }
         }
@@ -1001,7 +1008,7 @@ _pysqlite_query_execute(pysqlite_Cursor* self, int multiple, PyObject* operation
             int sqlite3_state = sqlite3_errcode(db);
             // stmt_reset() may return a previously set exception,
             // either triggered because of Python or sqlite3.
-            rc = stmt_reset(self->statement);
+            rc = stmt_reset(self->statement, 0);
             Py_CLEAR(self->statement);
             if (sqlite3_state == SQLITE_OK && rc != SQLITE_OK) {
                 cursor_cannot_reset_stmt_error(self, 1);
@@ -1170,15 +1177,15 @@ pysqlite_cursor_iternext(PyObject *op)
         if (self->statement->is_dml) {
             self->rowcount = (long)sqlite3_changes(self->connection->db);
         }
-        rc = stmt_reset(self->statement);
+        rc = stmt_reset(self->statement, 0);
         Py_CLEAR(self->statement);
         if (rc != SQLITE_OK) {
             goto reset_failure;
         }
     }
     else if (rc != SQLITE_ROW) {
         rc = set_error_from_db(self->connection->state, self->connection->db);
-        int reset_rc = stmt_reset(self->statement);
+        int reset_rc = stmt_reset(self->statement, 0);
         Py_CLEAR(self->statement);
         Py_DECREF(row);
         if (rc == SQLITE_OK && reset_rc != SQLITE_OK) {
@@ -1352,7 +1359,7 @@ pysqlite_cursor_close_impl(pysqlite_Cursor *self)
     }
 
     if (self->statement) {
-        int rc = stmt_reset(self->statement);
+        int rc = stmt_reset(self->statement, 0);
         // Force self->statement to be NULL even if stmt_reset() may have
         // failed to avoid a possible double-free if someone calls close()
         // twice as a leak here would be better than a double-free.

Modules/_sqlite/statement.c

@@ -97,6 +97,7 @@ pysqlite_statement_create(pysqlite_Connection *connection, PyObject *sql)
     }
 
     self->st = stmt;
+    self->in_use = 0;
     self->is_dml = is_dml;
 
     PyObject_GC_Track(self);

Modules/_sqlite/statement.h

@@ -32,9 +32,31 @@ typedef struct
 {
     PyObject_HEAD
     sqlite3_stmt* st;
+    int in_use;
     int is_dml;
 } pysqlite_Statement;
 
+static inline int
+pysqlite_statement_try_mark_in_use(pysqlite_Statement *self)
+{
+    int in_use;
+    Py_BEGIN_CRITICAL_SECTION(self);
+    in_use = self->in_use;
+    if (!in_use) {
+        self->in_use = 1;
+    }
+    Py_END_CRITICAL_SECTION();
+    return !in_use;
+}
+
+static inline void
+pysqlite_statement_mark_in_use(pysqlite_Statement *self, int in_use)
+{
+    Py_BEGIN_CRITICAL_SECTION(self);
+    self->in_use = in_use;
+    Py_END_CRITICAL_SECTION();
+}
+
 pysqlite_Statement *pysqlite_statement_create(pysqlite_Connection *connection, PyObject *sql);
 
 int pysqlite_statement_setup_types(PyObject *module);