Fixed leak of cancel key on connection.reset()

Moving the final free into the destructor is not necessary but looks
appropriate.

3 files changed, 19 insertions, 15 deletions

diff --git a/psycopg/connection.h b/psycopg/connection.h
index 4787e11..fc8c877 100644
--- a/psycopg/connection.h
+++ b/psycopg/connection.h
@@ -134,7 +134,6 @@ HIDDEN int  conn_get_standard_conforming_strings(PGconn *pgconn);
 RAISES_NEG HIDDEN int  conn_get_isolation_level(connectionObject *self);
 HIDDEN int  conn_get_protocol_version(PGconn *pgconn);
 HIDDEN int  conn_get_server_version(PGconn *pgconn);
-HIDDEN PGcancel *conn_get_cancel(PGconn *pgconn);
 HIDDEN void conn_notice_process(connectionObject *self);
 HIDDEN void conn_notice_clean(connectionObject *self);
 HIDDEN void conn_notifies_process(connectionObject *self);
diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index f36a0db..51c2381 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -439,10 +439,22 @@ conn_get_server_version(PGconn *pgconn)
     return (int)PQserverVersion(pgconn);
 }
 
-PGcancel *
-conn_get_cancel(PGconn *pgconn)
+/* set up the cancel key of the connection.
+ * On success return 0, else set an exception and return -1
+ */
+RAISES_NEG static int
+conn_setup_cancel(connectionObject *self, PGconn *pgconn)
 {
-    return PQgetCancel(pgconn);
+    if (self->cancel) {
+        PQfreeCancel(self->cancel);
+    }
+
+    if (!(self->cancel = PQgetCancel(self->pgconn))) {
+        PyErr_SetString(OperationalError, "can't get cancellation key");
+        return -1;
+    }
+
+    return 0;
 }
 
 
@@ -486,9 +498,7 @@ conn_setup(connectionObject *self, PGconn *pgconn)
         return -1;
     }
 
-    self->cancel = conn_get_cancel(self->pgconn);
-    if (self->cancel == NULL) {
-        PyErr_SetString(OperationalError, "can't get cancellation key");
+    if (0 > conn_setup_cancel(self, pgconn)) {
         return -1;
     }
 
@@ -788,10 +798,8 @@ _conn_poll_setup_async(connectionObject *self)
         if (0 > conn_read_encoding(self, self->pgconn)) {
             break;
         }
-        self->cancel = conn_get_cancel(self->pgconn);
-        if (self->cancel == NULL) {
-            PyErr_SetString(OperationalError, "can't get cancellation key");
-            break;
+        if (0 > conn_setup_cancel(self, self->pgconn)) {
+            return -1;
         }
 
         /* asynchronous connections always use isolation level 0, the user is
@@ -959,10 +967,6 @@ void conn_close_locked(connectionObject *self)
         self->pgconn = NULL;
         Dprintf("conn_close: PQfinish called");
     }
-    if (self->cancel) {
-        PQfreeCancel(self->cancel);
-        self->cancel = NULL;
-    }
 }
 
 /* conn_commit - commit on a connection */
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index 056be01..af8905d 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -1135,6 +1135,7 @@ connection_dealloc(PyObject* obj)
     PyMem_Free(self->encoding);
     PyMem_Free(self->codec);
     if (self->critical) free(self->critical);
+    if (self->cancel) PQfreeCancel(self->cancel);
 
     Py_CLEAR(self->tpc_xid);
     Py_CLEAR(self->async_cursor);