Properly cleanup memory of broken connections

Fixed ticket #148.

5 files changed, 40 insertions, 4 deletions

diff --git a/NEWS b/NEWS
index c128cec..f344dec 100644
--- a/NEWS
+++ b/NEWS
@@ -13,6 +13,7 @@ What's new in psycopg 2.5
     Tobias Oberstein for the feature development.
   - connection.reset() implemented using DISCARD ALL on server versions
     supporting it.
+  - Properly cleanup memory of broken connections (ticket #142).
   - 'errorcodes' map updated to PostgreSQL 9.2.
 
 
diff --git a/psycopg/connection_int.c b/psycopg/connection_int.c
index ad49575..0e5cb3d 100644
--- a/psycopg/connection_int.c
+++ b/psycopg/connection_int.c
@@ -918,7 +918,8 @@ conn_poll(connectionObject *self)
 void
 conn_close(connectionObject *self)
 {
-    if (self->closed) {
+    /* a connection with closed == 2 still requires cleanup */
+    if (self->closed == 1) {
         return;
     }
 
@@ -936,7 +937,7 @@ conn_close(connectionObject *self)
 
 void conn_close_locked(connectionObject *self)
 {
-    if (self->closed) {
+    if (self->closed == 1) {
         return;
     }
 
@@ -957,6 +958,8 @@ void conn_close_locked(connectionObject *self)
         PQfinish(self->pgconn);
         self->pgconn = NULL;
         Dprintf("conn_close: PQfinish called");
+    }
+    if (self->cancel) {
         PQfreeCancel(self->cancel);
         self->cancel = NULL;
     }
diff --git a/psycopg/connection_type.c b/psycopg/connection_type.c
index b5fd078..63ac175 100644
--- a/psycopg/connection_type.c
+++ b/psycopg/connection_type.c
@@ -1127,7 +1127,7 @@ connection_dealloc(PyObject* obj)
 
     PyObject_GC_UnTrack(self);
 
-    if (self->closed == 0) conn_close(self);
+    conn_close(self);
 
     conn_notice_clean(self);
 
diff --git a/tests/test_connection.py b/tests/test_connection.py
index b3488f0..76e83cf 100755
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -26,10 +26,11 @@ import os
 import time
 import threading
 from testutils import unittest, decorate_all_tests
-from testutils import skip_before_postgres, skip_after_postgres
+from testutils import skip_before_postgres, skip_after_postgres, skip_if_no_superuser
 from operator import attrgetter
 
 import psycopg2
+import psycopg2.errorcodes
 import psycopg2.extensions
 from testconfig import dsn, dbname
 
@@ -66,6 +67,22 @@ class ConnectionTests(unittest.TestCase):
         conn.close()
         self.assertEqual(curs.closed, True)
 
+    @skip_before_postgres(8, 4)
+    @skip_if_no_superuser
+    def test_cleanup_on_badconn_close(self):
+        # ticket #148
+        conn = self.conn
+        cur = conn.cursor()
+        try:
+            cur.execute("select pg_terminate_backend(pg_backend_pid())")
+        except psycopg2.OperationalError, e:
+            if e.pgcode != psycopg2.errorcodes.ADMIN_SHUTDOWN:
+                raise
+
+        self.assertEqual(conn.closed, 2)
+        conn.close()
+        self.assertEqual(conn.closed, 1)
+
     def test_reset(self):
         conn = self.conn
         # switch isolation level, then reset
diff --git a/tests/testutils.py b/tests/testutils.py
index 26551d4..068a913 100644
--- a/tests/testutils.py
+++ b/tests/testutils.py
@@ -207,6 +207,21 @@ def skip_from_python(*ver):
         return skip_from_python__
     return skip_from_python_
 
+def skip_if_no_superuser(f):
+    """Skip a test if the database user running the test is not a superuser"""
+    def skip_if_no_superuser_(self):
+        from psycopg2 import ProgrammingError
+        try:
+            return f(self)
+        except ProgrammingError, e:
+            import psycopg2.errorcodes
+            if e.pgcode == psycopg2.errorcodes.INSUFFICIENT_PRIVILEGE:
+                self.skipTest("skipped because not superuser")
+            else:
+                raise
+
+    return skip_if_no_superuser_
+
 
 def script_to_py3(script):
     """Convert a script to Python3 syntax if required."""