From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Sat, 23 May 2020 04:50:41 +0000 (+1200)
Subject: Close the connection on GC only if created in the same process
X-Git-Tag: 3.0.dev0~492
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=ae8a43bc527b1b06d724386576e07623b49826ca;p=thirdparty%2Fpsycopg.git

Close the connection on GC only if created in the same process

Fix the same issue reported in psycopg/psycopg2#829.
---

diff --git a/psycopg3/pq/pq_ctypes.py b/psycopg3/pq/pq_ctypes.py
index b9bf534fa..ce29ea58d 100644
--- a/psycopg3/pq/pq_ctypes.py
+++ b/psycopg3/pq/pq_ctypes.py
@@ -8,6 +8,7 @@ implementation.
 
 # Copyright (C) 2020 The Psycopg Team
 
+import os
 import logging
 from weakref import ref
 from functools import partial
@@ -62,6 +63,7 @@ class PGconn:
         "pgconn_ptr",
         "notice_handler",
         "_notice_receiver",
+        "_procpid",
         "__weakref__",
     )
 
@@ -74,8 +76,13 @@ class PGconn:
         )
         impl.PQsetNoticeReceiver(pgconn_ptr, self._notice_receiver, None)
 
+        self._procpid = os.getpid()
+
     def __del__(self) -> None:
-        self.finish()
+        # Close the connection only if it was created in this process,
+        # not if this object is being GC'd after fork.
+        if os.getpid() == self._procpid:
+            self.finish()
 
     @classmethod
     def connect(cls, conninfo: bytes) -> "PGconn":
diff --git a/psycopg3/pq/pq_cython.pxd b/psycopg3/pq/pq_cython.pxd
index 479b85c4d..bf61c01d2 100644
--- a/psycopg3/pq/pq_cython.pxd
+++ b/psycopg3/pq/pq_cython.pxd
@@ -1,3 +1,4 @@
+from posix.fcntl cimport pid_t
 from psycopg3.pq cimport libpq as impl
 
 ctypedef char *(*conn_bytes_f) (const impl.PGconn *)
@@ -7,12 +8,12 @@ ctypedef int(*conn_int_f) (const impl.PGconn *)
 cdef class PGconn:
     cdef impl.PGconn* pgconn_ptr
     cdef object __weakref__
+    cdef public object notice_handler
+    cdef pid_t _procpid
 
     @staticmethod
     cdef PGconn _from_ptr(impl.PGconn *ptr)
 
-    cdef public object notice_handler
-
     cdef int _ensure_pgconn(self) except 0
     cdef char *_call_bytes(self, conn_bytes_f func) except NULL
     cdef int _call_int(self, conn_int_f func) except -1
diff --git a/psycopg3/pq/pq_cython.pyx b/psycopg3/pq/pq_cython.pyx
index d982dd47e..e1061cf77 100644
--- a/psycopg3/pq/pq_cython.pyx
+++ b/psycopg3/pq/pq_cython.pyx
@@ -4,6 +4,7 @@ libpq Python wrapper using cython bindings.
 
 # Copyright (C) 2020 The Psycopg Team
 
+from posix.unistd cimport getpid
 from cpython.mem cimport PyMem_Malloc, PyMem_Free
 
 import logging
@@ -59,9 +60,13 @@ cdef class PGconn:
 
     def __cinit__(self):
         self.pgconn_ptr = NULL
+        self._procpid = getpid()
 
     def __dealloc__(self):
-        self.finish()
+        # Close the connection only if it was created in this process,
+        # not if this object is being GC'd after fork.
+        if self._procpid == getpid():
+            self.finish()
 
     @classmethod
     def connect(cls, conninfo: bytes) -> PGconn:
diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py
index 66cdd9c6e..5aa126ca7 100644
--- a/tests/test_concurrency.py
+++ b/tests/test_concurrency.py
@@ -2,10 +2,15 @@
 Tests dealing with concurrency issues.
 """
 
+import os
+import sys
 import time
 import queue
 import pytest
+import shutil
+import tempfile
 import threading
+import subprocess as sp
 
 import psycopg3
 
@@ -54,3 +59,49 @@ def test_commit_concurrency(conn):
     stop = True
 
     assert notices.empty(), "%d notices raised" % notices.qsize()
+
+
+@pytest.mark.slow
+def test_multiprocess_close(dsn):
+    # Check the problem reported in psycopg2#829
+    # Subprocess gcs the copy of the fd after fork so it closes connection.
+    module = f"""\
+import time
+import psycopg3
+
+def thread():
+    conn = psycopg3.connect({repr(dsn)})
+    curs = conn.cursor()
+    for i in range(10):
+        curs.execute("select 1")
+        time.sleep(0.1)
+
+def process():
+    time.sleep(0.2)
+"""
+
+    script = """\
+import time
+import threading
+import multiprocessing
+import mptest
+
+t = threading.Thread(target=mptest.thread, name='mythread')
+t.start()
+time.sleep(0.2)
+multiprocessing.Process(target=mptest.process, name='myprocess').start()
+t.join()
+"""
+
+    dir = tempfile.mkdtemp()
+    try:
+        with open(os.path.join(dir, "mptest.py"), "w") as f:
+            f.write(module)
+        env = dict(os.environ)
+        env["PYTHONPATH"] = dir + os.pathsep + env.get("PYTHONPATH", "")
+        out = sp.check_output(
+            [sys.executable, "-c", script], stderr=sp.STDOUT, env=env
+        ).decode("utf8", "replace")
+        assert out == "", out.strip().splitlines()[-1]
+    finally:
+        shutil.rmtree(dir, ignore_errors=True)