Merge branch 'noleak' into python2

NEWS-2.3

@@ -6,6 +6,8 @@ What's new in psycopg 2.3.3
   - Added `register_composite()` function to cast PostgreSQL composite types
     into Python tuples/namedtuples.
   - The build script refuses to guess values if pg_config is not found.
+  - Connections and cursors are weakly referenceable.
+  - Fixed several reference leaks in less common code paths.
 
 * Bug fixes:
 

psycopg/adapter_datetime.c

@@ -348,20 +348,13 @@ psyco_Time(PyObject *self, PyObject *args)
     return res;
 }
 
-PyObject *
+static PyObject *
-psyco_Timestamp(PyObject *self, PyObject *args)
+_psyco_Timestamp(int year, int month, int day,
+                 int hour, int minute, double second, PyObject *tzinfo)
 {
+    double micro;
+    PyObject *obj;
     PyObject *res = NULL;
-    PyObject *tzinfo = NULL;
-    int year, month, day;
-    int hour=0, minute=0; /* default to midnight */
-    double micro, second=0.0;
-
-    PyObject* obj = NULL;
-
-    if (!PyArg_ParseTuple(args, "lii|iidO", &year, &month, &day,
-                          &hour, &minute, &second, &tzinfo))
-        return NULL;
 
     micro = (second - floor(second)) * 1000000.0;
     second = floor(second);
@@ -386,6 +379,21 @@ psyco_Timestamp(PyObject *self, PyObject *args)
     return res;
 }
 
+PyObject *
+psyco_Timestamp(PyObject *self, PyObject *args)
+{
+    PyObject *tzinfo = NULL;
+    int year, month, day;
+    int hour=0, minute=0; /* default to midnight */
+    double second=0.0;
+
+    if (!PyArg_ParseTuple(args, "lii|iidO", &year, &month, &day,
+                          &hour, &minute, &second, &tzinfo))
+        return NULL;
+
+    return _psyco_Timestamp(year, month, day, hour, minute, second, tzinfo);
+}
+
 PyObject *
 psyco_DateFromTicks(PyObject *self, PyObject *args)
 {
@@ -446,20 +454,12 @@ psyco_TimestampFromTicks(PyObject *self, PyObject *args)
     t = (time_t)floor(ticks);
     ticks -= (double)t;
     if (localtime_r(&t, &tm)) {
-        PyObject *value = Py_BuildValue("iiiiidO",
+        res = _psyco_Timestamp(
-            tm.tm_year+1900, tm.tm_mon+1, tm.tm_mday,
+            tm.tm_year + 1900, tm.tm_mon + 1, tm.tm_mday,
-            tm.tm_hour, tm.tm_min,
+            tm.tm_hour, tm.tm_min, (double)tm.tm_sec + ticks,
-            (double)tm.tm_sec + ticks,
             pyPsycopgTzLOCAL);
-        if (value) {
-            /* FIXME: not decref'ing the value here is a memory leak
-	       but, on the other hand, if we decref we get a clean nice
-	       segfault (on my 64 bit Python 2.4 box). So this leaks
-	       will stay until after 2.0.7 when we'll try to plug it */
-            res = psyco_Timestamp(self, value);
-        }
     }
-    
+
     return res;
 }
 

psycopg/connection.h

@@ -103,7 +103,7 @@ typedef struct {
     PGconn *pgconn;           /* the postgresql connection */
     PGcancel *cancel;         /* the cancellation structure */
 
-    PyObject *async_cursor;   /* a cursor executing an asynchronous query */
+    PyObject *async_cursor;   /* weakref to a cursor executing an asynchronous query */
     int async_status;         /* asynchronous execution status */
 
     /* notice processing */
@@ -119,6 +119,7 @@ typedef struct {
     PyObject *binary_types;   /* a set of typecasters for binary types */
 
     int equote;               /* use E''-style quotes for escaped strings */
+    PyObject *weakreflist;    /* list of weak references */
 
 } connectionObject;
 

psycopg/connection_int.c

@@ -752,7 +752,17 @@ conn_poll(connectionObject *self)
         if (res == PSYCO_POLL_OK && self->async_cursor) {
             /* An async query has just finished: parse the tuple in the
              * target cursor. */
-            cursorObject *curs = (cursorObject *)self->async_cursor;
+            cursorObject *curs;
+            PyObject *py_curs = PyWeakref_GetObject(self->async_cursor);
+            if (Py_None == py_curs) {
+                pq_clear_async(self);
+                PyErr_SetString(InterfaceError,
+                    "the asynchronous cursor has disappeared");
+                res = PSYCO_POLL_ERROR;
+                break;
+            }
+
+            curs = (cursorObject *)py_curs;
             IFCLEARPGRES(curs->pgres);
             curs->pgres = pq_get_last_result(self);
 
@@ -764,8 +774,7 @@ conn_poll(connectionObject *self)
             }
 
             /* We have finished with our async_cursor */
-            Py_XDECREF(self->async_cursor);
+            Py_CLEAR(self->async_cursor);
-            self->async_cursor = NULL;
         }
         break;
 

psycopg/connection_type.c

@@ -305,8 +305,6 @@ _psyco_conn_tpc_finish(connectionObject *self, PyObject *args,
             goto exit;
         }
     } else {
-        PyObject *tmp;
-
         /* committing/aborting our own transaction. */
         if (!self->tpc_xid) {
             PyErr_SetString(ProgrammingError,
@@ -332,11 +330,10 @@ _psyco_conn_tpc_finish(connectionObject *self, PyObject *args,
             goto exit;
         }
 
+        Py_CLEAR(self->tpc_xid);
+
         /* connection goes ready */
         self->status = CONN_STATUS_READY;
-        tmp = (PyObject *)self->tpc_xid;
-        self->tpc_xid = NULL;
-        Py_DECREF(tmp);
     }
 
     Py_INCREF(Py_None);
@@ -867,6 +864,7 @@ connection_setup(connectionObject *self, const char *dsn, long int async)
     self->binary_types = PyDict_New();
     self->notice_pending = NULL;
     self->encoding = NULL;
+    self->weakreflist = NULL;
 
     pthread_mutex_init(&(self->lock), NULL);
 
@@ -896,17 +894,22 @@ static void
 connection_dealloc(PyObject* obj)
 {
     connectionObject *self = (connectionObject *)obj;
-    
+
+    if (self->weakreflist) {
+        PyObject_ClearWeakRefs(obj);
+    }
+
     PyObject_GC_UnTrack(self);
 
     if (self->closed == 0) conn_close(self);
-    
+
     conn_notice_clean(self);
 
     if (self->dsn) free(self->dsn);
     if (self->encoding) free(self->encoding);
     if (self->critical) free(self->critical);
 
+    Py_CLEAR(self->tpc_xid);
     Py_CLEAR(self->async_cursor);
     Py_CLEAR(self->notice_list);
     Py_CLEAR(self->notice_filter);
@@ -960,6 +963,7 @@ connection_repr(connectionObject *self)
 static int
 connection_traverse(connectionObject *self, visitproc visit, void *arg)
 {
+    Py_VISIT(self->tpc_xid);
     Py_VISIT(self->async_cursor);
     Py_VISIT(self->notice_list);
     Py_VISIT(self->notice_filter);
@@ -1002,14 +1006,16 @@ PyTypeObject connectionType = {
     0,          /*tp_setattro*/
     0,          /*tp_as_buffer*/
 
-    Py_TPFLAGS_DEFAULT|Py_TPFLAGS_BASETYPE|Py_TPFLAGS_HAVE_GC, /*tp_flags*/
+    Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_GC |
+        Py_TPFLAGS_HAVE_WEAKREFS,
+                /*tp_flags*/
     connectionType_doc, /*tp_doc*/
 
     (traverseproc)connection_traverse, /*tp_traverse*/
     0,          /*tp_clear*/
 
     0,          /*tp_richcompare*/
-    0,          /*tp_weaklistoffset*/
+    offsetof(connectionObject, weakreflist), /* tp_weaklistoffset */
 
     0,          /*tp_iter*/
     0,          /*tp_iternext*/

psycopg/cursor.h

@@ -81,6 +81,8 @@ typedef struct {
     PyObject *string_types;   /* a set of typecasters for string types */
     PyObject *binary_types;   /* a set of typecasters for binary types */
 
+    PyObject *weakreflist;    /* list of weak references */
+
 } cursorObject;
 
 /* C-callable functions in cursor_int.c and cursor_ext.c */

psycopg/cursor_type.c

@@ -792,7 +792,8 @@ psyco_curs_fetchone(cursorObject *self, PyObject *args)
     /* if the query was async aggresively free pgres, to allow
        successive requests to reallocate it */
     if (self->row >= self->rowcount
-        && self->conn->async_cursor == (PyObject*)self)
+        && self->conn->async_cursor
+        && PyWeakref_GetObject(self->conn->async_cursor) == (PyObject*)self)
         IFCLEARPGRES(self->pgres);
 
     return res;
@@ -868,7 +869,8 @@ psyco_curs_fetchmany(cursorObject *self, PyObject *args, PyObject *kwords)
     /* if the query was async aggresively free pgres, to allow
        successive requests to reallocate it */
     if (self->row >= self->rowcount
-        && self->conn->async_cursor == (PyObject*)self)
+        && self->conn->async_cursor
+        && PyWeakref_GetObject(self->conn->async_cursor) == (PyObject*)self)
         IFCLEARPGRES(self->pgres);
 
     return list;
@@ -932,7 +934,8 @@ psyco_curs_fetchall(cursorObject *self, PyObject *args)
     /* if the query was async aggresively free pgres, to allow
        successive requests to reallocate it */
     if (self->row >= self->rowcount
-        && self->conn->async_cursor == (PyObject*)self)
+        && self->conn->async_cursor
+        && PyWeakref_GetObject(self->conn->async_cursor) == (PyObject*)self)
         IFCLEARPGRES(self->pgres);
 
     return list;
@@ -1642,6 +1645,7 @@ cursor_setup(cursorObject *self, connectionObject *conn, const char *name)
 
     self->string_types = NULL;
     self->binary_types = NULL;
+    self->weakreflist = NULL;
 
     Py_INCREF(Py_None);
     self->description = Py_None;
@@ -1667,7 +1671,11 @@ static void
 cursor_dealloc(PyObject* obj)
 {
     cursorObject *self = (cursorObject *)obj;
-    
+
+    if (self->weakreflist) {
+        PyObject_ClearWeakRefs(obj);
+    }
+
     PyObject_GC_UnTrack(self);
 
     if (self->name) PyMem_Free(self->name);
@@ -1769,14 +1777,15 @@ PyTypeObject cursorType = {
     0,          /*tp_as_buffer*/
 
     Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | Py_TPFLAGS_HAVE_ITER |
-      Py_TPFLAGS_HAVE_GC, /*tp_flags*/
+      Py_TPFLAGS_HAVE_GC | Py_TPFLAGS_HAVE_WEAKREFS ,
+                /*tp_flags*/
     cursorType_doc, /*tp_doc*/
 
     (traverseproc)cursor_traverse, /*tp_traverse*/
     0,          /*tp_clear*/
 
     0,          /*tp_richcompare*/
-    0,          /*tp_weaklistoffset*/
+    offsetof(cursorObject, weakreflist), /*tp_weaklistoffset*/
 
     cursor_iter, /*tp_iter*/
     cursor_next, /*tp_iternext*/

psycopg/pqpath.c

@@ -279,8 +279,7 @@ pq_clear_async(connectionObject *conn)
         Dprintf("pq_clear_async: clearing PGresult at %p", pgres);
         CLEARPGRES(pgres);
     }
-    Py_XDECREF(conn->async_cursor);
+    Py_CLEAR(conn->async_cursor);
-    conn->async_cursor = NULL;
 }
 
 
@@ -824,8 +823,11 @@ pq_execute(cursorObject *curs, const char *query, int async)
     }
     else {
         curs->conn->async_status = async_status;
-        Py_INCREF(curs);
+        curs->conn->async_cursor = PyWeakref_NewRef((PyObject *)curs, NULL);
-        curs->conn->async_cursor = (PyObject*)curs;
+        if (!curs->conn->async_cursor) {
+            /* weakref creation failed */
+            return -1;
+        }
     }
 
     return 1-async;

scripts/refcounter.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python
+"""Detect reference leaks after several unit test runs.
+
+The script runs the unit test and counts the objects alive after the run. If
+the object count differs between the last two runs, a report is printed and the
+script exits with error 1.
+"""
+
+# Copyright (C) 2011 Daniele Varrazzo <daniele.varrazzo@gmail.com>
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation; either version 2 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, write to the Free Software
+# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
+
+import gc
+import sys
+import difflib
+import unittest
+from pprint import pprint
+from collections import defaultdict
+
+def main():
+    opt = parse_args()
+
+    import psycopg2.tests
+    test = psycopg2.tests
+    if opt.suite:
+        test = getattr(test, opt.suite)
+
+    sys.stdout.write("test suite %s\n" % test.__name__)
+
+    for i in range(1, opt.nruns + 1):
+        sys.stdout.write("test suite run %d of %d\n" % (i, opt.nruns))
+        runner = unittest.TextTestRunner()
+        runner.run(test.test_suite())
+        dump(i, opt)
+
+    f1 = open('debug-%02d.txt' % (opt.nruns - 1)).readlines()
+    f2 = open('debug-%02d.txt' % (opt.nruns)).readlines()
+    for line in difflib.unified_diff(f1, f2,
+            "run %d" % (opt.nruns - 1), "run %d" % opt.nruns):
+        sys.stdout.write(line)
+
+    rv = f1 != f2 and 1 or 0
+
+    if opt.objs:
+        f1 = open('objs-%02d.txt' % (opt.nruns - 1)).readlines()
+        f2 = open('objs-%02d.txt' % (opt.nruns)).readlines()
+        for line in difflib.unified_diff(f1, f2,
+                "run %d" % (opt.nruns - 1), "run %d" % opt.nruns):
+            sys.stdout.write(line)
+
+    return rv
+
+def parse_args():
+    import optparse
+
+    parser = optparse.OptionParser(description=__doc__)
+    parser.add_option('--nruns', type='int', metavar="N", default=3,
+        help="number of test suite runs [default: %default]")
+    parser.add_option('--suite', metavar="NAME",
+        help="the test suite to run (e.g. 'test_cursor'). [default: all]")
+    parser.add_option('--objs', metavar="TYPE",
+        help="in case of leaks, print a report of object TYPE "
+            "(support still incomplete)")
+
+    opt, args = parser.parse_args()
+    return opt
+
+
+def dump(i, opt):
+    gc.collect()
+    objs = gc.get_objects()
+
+    c = defaultdict(int)
+    for o in objs:
+        c[type(o)] += 1
+
+    pprint(
+        sorted(((v,str(k)) for k,v in c.items()), reverse=True),
+        stream=open("debug-%02d.txt" % i, "w"))
+
+    if opt.objs:
+        co = []
+        t = getattr(__builtins__, opt.objs)
+        for o in objs:
+            if type(o) is t:
+                co.append(o)
+
+        # TODO: very incomplete
+        if t is dict:
+            co.sort(key = lambda d: d.items())
+        else:
+            co.sort()
+
+        pprint(co, stream=open("objs-%02d.txt" % i, "w"))
+
+
+if __name__ == '__main__':
+    sys.exit(main())
+

tests/test_async.py

@@ -415,6 +415,18 @@ class AsyncTests(unittest.TestCase):
         self.assertEqual("CREATE TABLE", cur.statusmessage)
         self.assert_(self.conn.notices)
 
+    def test_async_cursor_gone(self):
+        cur = self.conn.cursor()
+        cur.execute("select 42;");
+        del cur
+        self.assertRaises(psycopg2.InterfaceError, self.wait, self.conn)
+
+        # The connection is still usable
+        cur = self.conn.cursor()
+        cur.execute("select 42;");
+        self.wait(self.conn)
+        self.assertEqual(cur.fetchone(), (42,))
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)

tests/test_connection.py

@@ -111,6 +111,14 @@ class ConnectionTests(unittest.TestCase):
         self.assert_(time.time() - t0 < 3,
             "something broken in concurrency")
 
+    def test_weakref(self):
+        from weakref import ref
+        conn = psycopg2.connect(self.conn.dsn)
+        w = ref(conn)
+        conn.close()
+        del conn
+        self.assert_(w() is None)
+
 
 class IsolationLevelsTestCase(unittest.TestCase):
 

tests/test_cursor.py

@@ -99,6 +99,13 @@ class CursorTests(unittest.TestCase):
         curs2 = self.conn.cursor()
         self.assertEqual("foofoo", curs2.cast(705, 'foo'))
 
+    def test_weakref(self):
+        from weakref import ref
+        curs = self.conn.cursor()
+        w = ref(curs)
+        del curs
+        self.assert_(w() is None)
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)

tests/test_quote.py

@@ -71,7 +71,7 @@ class QuotingTestCase(unittest.TestCase):
             if not 0xD800 <= u <= 0xDFFF ]))    # surrogate area
         self.conn.set_client_encoding('UNICODE')
 
-        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
+        psycopg2.extensions.register_type(psycopg2.extensions.UNICODE, self.conn)
         curs.execute("SELECT %s::text;", (data,))
         res = curs.fetchone()[0]
 

tests/types_basic.py

@@ -166,7 +166,11 @@ class AdaptSubclassTest(unittest.TestCase):
 
         register_adapter(A, lambda a: AsIs("a"))
         register_adapter(B, lambda b: AsIs("b"))
-        self.assertEqual('b', adapt(C()).getquoted())
+        try:
+            self.assertEqual('b', adapt(C()).getquoted())
+        finally:
+           del psycopg2.extensions.adapters[A, psycopg2.extensions.ISQLQuote]
+           del psycopg2.extensions.adapters[B, psycopg2.extensions.ISQLQuote]
 
     def test_no_mro_no_joy(self):
         from psycopg2.extensions import adapt, register_adapter, AsIs
@@ -175,7 +179,11 @@ class AdaptSubclassTest(unittest.TestCase):
         class B(A): pass
 
         register_adapter(A, lambda a: AsIs("a"))
-        self.assertRaises(psycopg2.ProgrammingError, adapt, B())
+        try:
+            self.assertRaises(psycopg2.ProgrammingError, adapt, B())
+        finally:
+           del psycopg2.extensions.adapters[A, psycopg2.extensions.ISQLQuote]
+
 
 def test_suite():
     return unittest.TestLoader().loadTestsFromName(__name__)