Avoid a deadlock using concurrent green threads on the same connection

Use the async_cursor property to store an indication that something is
running (even if it is not necessarily a cursor running the query).

4 files changed, 24 insertions, 6 deletions

diff --git a/psycopg/connection.h b/psycopg/connection.h
index d79392b..24b3be3 100644
--- a/psycopg/connection.h
+++ b/psycopg/connection.h
@@ -91,7 +91,10 @@ typedef struct {
     PGconn *pgconn;           /* the postgresql connection */
     PGcancel *cancel;         /* the cancellation structure */
 
-    PyObject *async_cursor;   /* weakref to a cursor executing an asynchronous query */
+    /* Weakref to the object executing an asynchronous query. The object
+     * is a cursor for async connections, but it may be something else
+     * for a green connection. If NULL, the connection is idle. */
+    PyObject *async_cursor;
     int async_status;         /* asynchronous execution status */
 
     /* notice processing */
diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index 24d424d..6a1d9c5 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -881,7 +881,7 @@ conn_poll(connectionObject *self)
     case CONN_STATUS_PREPARED:
         res = _conn_poll_query(self);
 
-        if (res == PSYCO_POLL_OK && self->async_cursor) {
+        if (res == PSYCO_POLL_OK && self->async && self->async_cursor) {
             /* An async query has just finished: parse the tuple in the
              * target cursor. */
             cursorObject *curs;
diff --git a/psycopg/cursor_type.c b/psycopg/cursor_type.c
index 8ede845..c18fb71 100644
--- a/psycopg/cursor_type.c
+++ b/psycopg/cursor_type.c
@@ -739,7 +739,6 @@ psyco_curs_fetchone(cursorObject *self, PyObject *args)
     PyObject *res;
 
     EXC_IF_CURS_CLOSED(self);
-    EXC_IF_ASYNC_IN_PROGRESS(self, fetchone);
     if (_psyco_curs_prefetch(self) < 0) return NULL;
     EXC_IF_NO_TUPLES(self);
 
@@ -747,6 +746,7 @@ psyco_curs_fetchone(cursorObject *self, PyObject *args)
         char buffer[128];
 
         EXC_IF_NO_MARK(self);
+        EXC_IF_ASYNC_IN_PROGRESS(self, fetchone);
         EXC_IF_TPC_PREPARED(self->conn, fetchone);
         PyOS_snprintf(buffer, 127, "FETCH FORWARD 1 FROM \"%s\"", self->name);
         if (pq_execute(self, buffer, 0) == -1) return NULL;
@@ -853,7 +853,6 @@ psyco_curs_fetchmany(cursorObject *self, PyObject *args, PyObject *kwords)
     }
 
     EXC_IF_CURS_CLOSED(self);
-    EXC_IF_ASYNC_IN_PROGRESS(self, fetchmany);
     if (_psyco_curs_prefetch(self) < 0) return NULL;
     EXC_IF_NO_TUPLES(self);
 
@@ -861,6 +860,7 @@ psyco_curs_fetchmany(cursorObject *self, PyObject *args, PyObject *kwords)
         char buffer[128];
 
         EXC_IF_NO_MARK(self);
+        EXC_IF_ASYNC_IN_PROGRESS(self, fetchmany);
         EXC_IF_TPC_PREPARED(self->conn, fetchone);
         PyOS_snprintf(buffer, 127, "FETCH FORWARD %d FROM \"%s\"",
             (int)size, self->name);
@@ -924,7 +924,6 @@ psyco_curs_fetchall(cursorObject *self, PyObject *args)
     PyObject *list, *res;
 
     EXC_IF_CURS_CLOSED(self);
-    EXC_IF_ASYNC_IN_PROGRESS(self, fetchall);
     if (_psyco_curs_prefetch(self) < 0) return NULL;
     EXC_IF_NO_TUPLES(self);
 
@@ -932,6 +931,7 @@ psyco_curs_fetchall(cursorObject *self, PyObject *args)
         char buffer[128];
 
         EXC_IF_NO_MARK(self);
+        EXC_IF_ASYNC_IN_PROGRESS(self, fetchall);
         EXC_IF_TPC_PREPARED(self->conn, fetchall);
         PyOS_snprintf(buffer, 127, "FETCH FORWARD ALL FROM \"%s\"", self->name);
         if (pq_execute(self, buffer, 0) == -1) return NULL;
@@ -1112,7 +1112,6 @@ psyco_curs_scroll(cursorObject *self, PyObject *args, PyObject *kwargs)
         return NULL;
 
     EXC_IF_CURS_CLOSED(self);
-    EXC_IF_ASYNC_IN_PROGRESS(self, scroll)
 
     /* if the cursor is not named we have the full result set and we can do
        our own calculations to scroll; else we just delegate the scrolling
@@ -1141,6 +1140,7 @@ psyco_curs_scroll(cursorObject *self, PyObject *args, PyObject *kwargs)
         char buffer[128];
 
         EXC_IF_NO_MARK(self);
+        EXC_IF_ASYNC_IN_PROGRESS(self, scroll)
         EXC_IF_TPC_PREPARED(self->conn, scroll);
 
         if (strcmp(mode, "absolute") == 0) {
diff --git a/psycopg/green.c b/psycopg/green.c
index c9b6e07..65578f5 100644
--- a/psycopg/green.c
+++ b/psycopg/green.c
@@ -152,6 +152,20 @@ psyco_exec_green(connectionObject *conn, const char *command)
 {
     PGresult *result = NULL;
 
+    /* Check that there is a single concurrently executing query */
+    if (conn->async_cursor) {
+        PyErr_SetString(ProgrammingError,
+            "a single async query can be executed on the same connection");
+        goto end;
+    }
+    /* we don't care about which cursor is executing the query, and
+     * it may also be that no cursor is involved at all and this is
+     * an internal query. So just store anything in the async_cursor,
+     * respecting the code expecting it to be a weakref */
+    if (!(conn->async_cursor = PyWeakref_NewRef((PyObject*)conn, NULL))) {
+        goto end;
+    }
+
     /* Send the query asynchronously */
     if (0 == pq_send_query(conn, command)) {
         goto end;
@@ -173,6 +187,7 @@ psyco_exec_green(connectionObject *conn, const char *command)
 
 end:
     conn->async_status = ASYNC_DONE;
+    Py_CLEAR(conn->async_cursor);
     return result;
 }