NOWAIT lock: raise ProgrammingError on concurrent connection use

Author: Copilot

src/MySQLdb/_mysql.c

@@ -122,8 +122,26 @@ _mysql_ConnectionObject_AllocateLock(_mysql_ConnectionObject *self)
     return 0;
 }
 
-static void
+/* Try to acquire the connection lock without blocking.
+ * Returns 0 on success, -1 if the lock is already held by another thread
+ * (ProgrammingError is set). */
+static int
 _mysql_ConnectionObject_Lock(_mysql_ConnectionObject *self)
+{
+    PyLockStatus status = PyThread_acquire_lock(self->lock, NOWAIT_LOCK);
+    if (status != PY_LOCK_ACQUIRED) {
+        PyErr_SetString(_mysql_ProgrammingError,
+            "This connection is already in use from another thread. "
+            "Do not use the same connection object from multiple threads simultaneously.");
+        return -1;
+    }
+    return 0;
+}
+
+/* Blocking variant used only during deallocation where we must acquire the
+ * lock to call mysql_free_result() even if it means waiting. */
+static void
+_mysql_ConnectionObject_LockWait(_mysql_ConnectionObject *self)
 {
     Py_BEGIN_ALLOW_THREADS
     PyThread_acquire_lock(self->lock, WAIT_LOCK);
@@ -140,12 +158,12 @@ _mysql_ConnectionObject_Unlock(_mysql_ConnectionObject *self)
 #define END_CONNECTION_LOCK(c) _mysql_ConnectionObject_Unlock(c)
 
 #define BEGIN_RESULT_CONNECTION_LOCK(r) \
-    BEGIN_CONNECTION_LOCK(result_connection(r))
+    _mysql_ConnectionObject_LockWait(result_connection(r))
 #define END_RESULT_CONNECTION_LOCK(r) \
     END_CONNECTION_LOCK(result_connection(r))
 #define BEGIN_CONNECTION_OPERATION(c, on_closed) \
     do { \
-        BEGIN_CONNECTION_LOCK(c); \
+        if (_mysql_ConnectionObject_Lock(c) < 0) return NULL; \
         if (!(c)->open) { \
             END_CONNECTION_LOCK(c); \
             on_closed; \
@@ -320,7 +338,9 @@ _mysql_ResultObject_Initialize(
     self->conn = (PyObject *) conn;
     Py_INCREF(conn);
     self->use = use;
-    BEGIN_CONNECTION_LOCK(conn);
+    if (BEGIN_CONNECTION_LOCK(conn) < 0) {
+        return -1;
+    }
     if (!conn->open) {
         END_CONNECTION_LOCK(conn);
         _mysql_Exception(conn);
@@ -1135,7 +1155,10 @@ _mysql_escape_string(
     if (self && PyModule_Check((PyObject*)self))
         self = NULL;
     if (self) {
-        BEGIN_CONNECTION_LOCK(self);
+        if (BEGIN_CONNECTION_LOCK(self) < 0) {
+            Py_DECREF(str);
+            return NULL;
+        }
         use_connection = self->open;
     }
     if (use_connection) {
@@ -1177,7 +1200,7 @@ _mysql_string_literal(
     if (self && PyModule_Check((PyObject*)self))
         self = NULL;
     if (self) {
-        BEGIN_CONNECTION_LOCK(self);
+        if (BEGIN_CONNECTION_LOCK(self) < 0) return NULL;
         use_connection = self->open;
     }
 
@@ -1305,7 +1328,7 @@ _mysql_escape(
                     "argument 2 must be a mapping");
             return NULL;
         }
-        BEGIN_CONNECTION_LOCK((_mysql_ConnectionObject *)self);
+        if (BEGIN_CONNECTION_LOCK((_mysql_ConnectionObject *)self) < 0) return NULL;
         converter = ((_mysql_ConnectionObject *) self)->converter;
         Py_XINCREF(converter);
         END_CONNECTION_LOCK((_mysql_ConnectionObject *)self);

tests/test_connection.py

@@ -29,68 +29,78 @@ def test_multi_statements_false():
     assert rows == ((17,),)
 
 
-def _assert_thread_id_blocked_during_operation(conn, func, min_wait=0.15):
-    error = None
-    done = threading.Event()
-    thread_id_done = threading.Event()
-    results = {}
-
-    def run():
-        nonlocal error
-        try:
-            func()
-        except Exception as exc:  # pragma: no cover - error checked below
-            error = exc
-        finally:
-            done.set()
-
-    def read_thread_id():
-        try:
-            results["thread_id"] = conn.thread_id()
-        except Exception as exc:  # pragma: no cover - assertion checked below
-            results["error"] = exc
-        finally:
-            thread_id_done.set()
-
-    thread = threading.Thread(target=run)
-    thread.start()
-    time.sleep(0.05)
-    assert not done.is_set()
-
-    blocker = threading.Thread(target=read_thread_id)
-    blocker.start()
-
-    assert not thread_id_done.wait(min_wait)
-    thread.join()
-    blocker.join()
-    assert error is None
-    assert done.is_set()
-    assert "error" not in results
-    assert isinstance(results["thread_id"], int)
-
-
-def test_connection_methods_are_serialized():
+def test_connection_concurrent_use_raises():
+    """While a slow query holds the connection lock, any other access from
+    a second thread must raise ProgrammingError immediately (not block)."""
     conn = connection_factory()
     try:
-        def run_query():
-            conn.query("SELECT SLEEP(0.2)")
-            result = conn.store_result()
-            assert result.fetch_row() == ((0,),)
-
-        _assert_thread_id_blocked_during_operation(conn, run_query)
+        thread_error = None
+        done = threading.Event()
+
+        def run_slow_query():
+            nonlocal thread_error
+            try:
+                conn.query("SELECT SLEEP(0.5)")
+                result = conn.store_result()
+                result.fetch_row()
+            except Exception as exc:  # pragma: no cover
+                thread_error = exc
+            finally:
+                done.set()
+
+        thread = threading.Thread(target=run_slow_query)
+        thread.start()
+
+        # Give the background thread time to acquire the lock and enter SLEEP.
+        time.sleep(0.1)
+
+        start = time.monotonic()
+        with pytest.raises(ProgrammingError, match="already in use"):
+            conn.thread_id()
+        # Should fail immediately, not wait for the SLEEP to finish.
+        assert time.monotonic() - start < 0.1
+
+        done.wait()
+        thread.join()
+        assert thread_error is None
     finally:
         conn.close()
 
 
-def test_result_methods_share_connection_lock():
+def test_result_concurrent_use_raises():
+    """While fetch_row holds the connection lock streaming a slow result,
+    any other access from a second thread must raise ProgrammingError immediately."""
     conn = connection_factory()
     try:
-        conn.query("SELECT 1 UNION ALL SELECT SLEEP(0.2)")
+        conn.query("SELECT 1 UNION ALL SELECT SLEEP(0.5)")
         result = conn.use_result()
 
-        def fetch_all_rows():
-            assert result.fetch_row(maxrows=0) == ((1,), (0,))
+        thread_error = None
+        done = threading.Event()
 
-        _assert_thread_id_blocked_during_operation(conn, fetch_all_rows)
+        def fetch_all_rows():
+            nonlocal thread_error
+            try:
+                assert result.fetch_row(maxrows=0) == ((1,), (0,))
+            except Exception as exc:  # pragma: no cover
+                thread_error = exc
+            finally:
+                done.set()
+
+        thread = threading.Thread(target=fetch_all_rows)
+        thread.start()
+
+        # Give the background thread time to acquire the lock and enter SLEEP.
+        time.sleep(0.1)
+
+        start = time.monotonic()
+        with pytest.raises(ProgrammingError, match="already in use"):
+            conn.thread_id()
+        # Should fail immediately, not wait for the SLEEP to finish.
+        assert time.monotonic() - start < 0.1
+
+        done.wait()
+        thread.join()
+        assert thread_error is None
     finally:
         conn.close()