Fix three bugs in the pooling layer

* PooledDB.connection() raised an IndexError when a thread waiting for a
  shared connection that was in a transaction was woken up after that
  connection had become idle and left the shared cache.  The waiter now
  re-evaluates the whole situation and falls back to the idle cache.
* PooledPgConnection.reopen() wrapped another pooled connection instead of
  the underlying steady connection, so every close()/reopen() cycle nested
  the proxies one level deeper.  The pool got a pooled_connection() method
  that hands out the bare steady connection for this purpose.
* SharedDBConnection compared by value but hashed by identity.  Equality is
  now identity, which also makes the list.remove() in unshare() remove the
  connection being unshared rather than the first equally ranked one.

Author: Cito

dbutils/pooled_db.py

@@ -315,30 +315,37 @@ def connection(self, shareable=True):
         """
         if shareable and self._maxshared:
             with self._lock:
-                while (not self._shared_cache and self._maxconnections
-                        and self._connections >= self._maxconnections):
-                    self._wait_lock()
-                if len(self._shared_cache) < self._maxshared:
-                    # shared cache is not full, get a dedicated connection
-                    try:  # first try to get it from the idle cache
-                        con = self._idle_cache.pop(0)
-                    except IndexError:  # else get a fresh connection
-                        con = self.steady_connection()
-                    else:
-                        con._ping_check()  # check this connection
-                    con = SharedDBConnection(con)
-                    self._connections += 1
-                else:  # shared cache full or no more connections allowed
+                while True:
+                    while (not self._shared_cache and self._maxconnections
+                            and self._connections >= self._maxconnections):
+                        self._wait_lock()
+                    if len(self._shared_cache) < self._maxshared:
+                        # shared cache is not full, get a dedicated connection
+                        try:  # first try to get it from the idle cache
+                            con = self._idle_cache.pop(0)
+                        except IndexError:  # else get a fresh connection
+                            con = self.steady_connection()
+                        else:
+                            con._ping_check()  # check this connection
+                        con = SharedDBConnection(con)
+                        self._connections += 1
+                        break
+                    # shared cache full or no more connections allowed
                     self._shared_cache.sort()  # least shared connection first
-                    con = self._shared_cache.pop(0)  # get it
-                    while con.con._transaction:
-                        # do not share connections which are in a transaction
-                        self._shared_cache.insert(0, con)
+                    # only look at it, but leave it in the shared cache
+                    # as long as we may still have to wait for it
+                    con = self._shared_cache[0]
+                    if con.con._transaction:
+                        # do not share connections which are in a transaction,
+                        # wait until the situation has changed and start over,
+                        # since by then the connection may have become idle
+                        # and been removed from the shared cache altogether
                         self._wait_lock()
-                        self._shared_cache.sort()
-                        con = self._shared_cache.pop(0)
+                        continue
+                    del self._shared_cache[0]  # get it
                     con.con._ping_check()  # check the underlying connection
                     con.share()  # increase share of this connection
+                    break
                 # put the connection (back) into the shared cache
                 self._shared_cache.append(con)
                 self._lock.notify()
@@ -487,12 +494,17 @@ def __lt__(self, other):
 
     def __eq__(self, other):
         """Check whether this connection is the same as the other one."""
-        return (self.con._transaction == other.con._transaction
-                and self.shared == other.shared)
+        # The ordering above only serves to pick the least shared connection,
+        # it does not make different connections interchangeable.  Therefore
+        # equality is identity, so that the shared cache can be searched for
+        # one particular connection (see the unshare method of the pool).
+        return self is other
 
     def __hash__(self):
         """Get hash value of this connection."""
-        return hash((self.con, self.shared))
+        # must be based on identity as well, and must stay constant even
+        # though the number of shares of the connection can change
+        return id(self)
 
     def share(self):
         """Increase the share of this connection."""

dbutils/pooled_pg.py

@@ -214,15 +214,23 @@ def steady_connection(self):
         return SteadyPgConnection(self._maxusage, self._setsession, True,
                                   *self._args, **self._kwargs)
 
-    def connection(self):
-        """Get a steady, cached PostgreSQL connection from the pool."""
+    def pooled_connection(self):
+        """Get a steady PostgreSQL connection from the pool cache.
+
+        This is the unwrapped connection that is kept in the pool cache.
+        It reserves one of the generally allowed connections, so it must
+        be balanced with a corresponding call of the cache() method.
+        """
         if self._connections and not self._connections.acquire(self._blocking):
             raise TooManyConnectionsError
         try:
-            con = self._cache.get_nowait()
+            return self._cache.get_nowait()
         except Empty:
-            con = self.steady_connection()
-        return PooledPgConnection(self, con)
+            return self.steady_connection()
+
+    def connection(self):
+        """Get a steady, cached PostgreSQL connection from the pool."""
+        return PooledPgConnection(self, self.pooled_connection())
 
     def cache(self, con):
         """Put a connection back into the pool cache."""
@@ -289,7 +297,10 @@ def reopen(self):
         if self._con:
             self._con.reopen()
         else:
-            self._con = self._pool.connection()
+            # take the steady connection itself out of the pool, since
+            # wrapping another pooled connection would nest the proxies
+            # and put a proxy instead of a connection into the pool cache
+            self._con = self._pool.pooled_connection()
 
     def __getattr__(self, name):
         """Proxy all members of the class."""

docs/changelog.html

@@ -33,6 +33,14 @@ <h2>3.2.0</h2>
 <li><p>Fixed the transaction flag being reset on the cursor instead of the
 connection when a connection could not be reopened during a transaction,
 which kept the failover mechanism suspended.</p></li>
+<li><p>Fixed an <span class="docutils literal">IndexError</span> in <span class="docutils literal">PooledDB.connection()</span> when a thread waiting
+for a shared connection that was in a transaction got woken up after that
+connection had become idle and left the shared cache.</p></li>
+<li><p>Fixed <span class="docutils literal">PooledPgConnection.reopen()</span> wrapping another pooled connection
+instead of the underlying steady connection when it was called after the
+connection had already been returned to the pool.</p></li>
+<li><p>Fixed <span class="docutils literal">pooled_db.SharedDBConnection</span> violating the contract between
+<span class="docutils literal">__eq__</span> and <span class="docutils literal">__hash__</span>; shared connections now compare by identity.</p></li>
 </ul>
 </section>
 <section id="section-2">

docs/changelog.rst

@@ -25,6 +25,14 @@ Bugfixes:
 * Fixed the transaction flag being reset on the cursor instead of the
   connection when a connection could not be reopened during a transaction,
   which kept the failover mechanism suspended.
+* Fixed an ``IndexError`` in ``PooledDB.connection()`` when a thread waiting
+  for a shared connection that was in a transaction got woken up after that
+  connection had become idle and left the shared cache.
+* Fixed ``PooledPgConnection.reopen()`` wrapping another pooled connection
+  instead of the underlying steady connection when it was called after the
+  connection had already been returned to the pool.
+* Fixed ``pooled_db.SharedDBConnection`` violating the contract between
+  ``__eq__`` and ``__hash__``; shared connections now compare by identity.
 
 3.1.2
 =====

tests/test_pooled_db.py

@@ -1379,6 +1379,40 @@ def connection():
     assert shared[0]._con is con
 
 
+def test_shared_in_transaction_becoming_idle(dbapi):
+    """Check that a thread waiting for a shared connection survives unshare."""
+    pool = PooledDB(dbapi, 0, 0, 1, 0, True)
+    db = pool.connection()
+    con = db._con
+    db.begin()
+    # the only connection that may be shared is in a transaction now
+    shared = []
+    errors = []
+
+    def connection():
+        try:
+            shared.append(pool.connection())
+        except Exception as error:
+            errors.append(error)
+
+    thread = Thread(target=connection)
+    thread.start()
+    thread.join(0.1)
+    # the thread cannot share that connection and blocks instead of failing
+    assert thread.is_alive()
+    assert not shared
+    # closing the connection takes it out of the shared cache completely,
+    # so the waiting thread is woken up to find an empty shared cache;
+    # it must then get the connection from the idle cache instead of
+    # trying to share a connection that is no longer there
+    db.close()
+    thread.join(0.1)
+    assert not thread.is_alive()
+    assert not errors
+    assert len(shared) == 1
+    assert shared[0]._con is con
+
+
 def test_shared_in_transaction_with_two_connections(dbapi):
     """Check that sharing prefers connections without a transaction."""
     pool = PooledDB(dbapi, 0, 2, 2)
@@ -1490,26 +1524,45 @@ def test_shared_db_connection_compare(dbapi):
     con1.con._transaction = False
     con2 = SharedDBConnection(dbapi.connect())
     con2.con._transaction = False
-    assert con1 == con2
-    assert con1 <= con2
-    assert con1 >= con2
-    assert not con1 != con2  # noqa: SIM202
+    # connections with the same number of shares are ranked equally
     assert not con1 < con2
-    assert not con1 > con2
+    assert not con2 < con1
+    # the connection with fewer shares is preferred
     con2.share()
-    assert not con1 == con2  # noqa: SIM201
+    assert con1 < con2
+    assert not con2 < con1
     assert con1 <= con2
     assert not con1 >= con2
-    assert con1 != con2
-    assert con1 < con2
-    assert not con1 > con2
+    # but connections in a transaction always come last
     con1.con._transaction = True
-    assert not con1 == con2  # noqa: SIM201
-    assert not con1 <= con2
-    assert con1 >= con2
-    assert con1 != con2
     assert not con1 < con2
+    assert con2 < con1
     assert con1 > con2
+    assert not con1 <= con2
+
+
+def test_shared_db_connection_equality(dbapi):
+    """Check that shared connections are only equal to themselves."""
+    con1 = SharedDBConnection(dbapi.connect())
+    con1.con._transaction = False
+    con2 = SharedDBConnection(dbapi.connect())
+    con2.con._transaction = False
+    same_as_con1 = con1
+    # the ordering only serves to pick the least shared connection,
+    # it does not make equally ranked connections interchangeable
+    assert con1 == same_as_con1
+    assert con1 != con2
+    # equal connections must therefore also have equal hash values
+    # (this used to be violated, since the equality was based on the
+    # values while the hash value was based on the identity)
+    assert hash(con1) == hash(same_as_con1)
+    assert len({con1, con2, same_as_con1}) == 2
+    # and removing one connection from a list of connections
+    # must not remove another connection that is ranked the same
+    cache = [con1, con2]
+    cache.remove(con2)
+    assert len(cache) == 1
+    assert cache[0] is con1
 
 
 def test_shared_db_connection_hash(dbapi):
@@ -1518,9 +1571,10 @@ def test_shared_db_connection_hash(dbapi):
     con = SharedDBConnection(dbapi.connect())
     hashed = hash(con)
     assert hash(con) == hashed
-    # the hash is derived from the underlying connection and the shares
+    # the hash is based on the identity, so it does not change
+    # when the number of shares of the connection changes
     con.share()
-    assert hash(con) != hashed
+    assert hash(con) == hashed
     con.unshare()
     assert hash(con) == hashed
 

tests/test_pooled_pg.py

@@ -370,6 +370,30 @@ def test_reopen_connection_after_closing():
     assert con.num_queries == 1
 
 
+def test_reopen_connection_does_not_nest_the_proxy():
+    """Check that reopening does not wrap another pooled connection."""
+    pool = PooledPg(1, 2, 3)
+    db = pool.connection()
+    con = db._con
+    db.close()
+    db.reopen()
+    # the proxy must wrap the steady connection itself, not another proxy,
+    # since otherwise the pool cache would fill up with nested proxies
+    assert db._con is con
+    assert isinstance(db._con, SteadyPgConnection)
+    db.close()
+    cached = pool._cache.get_nowait()
+    assert cached is con
+    assert isinstance(cached, SteadyPgConnection)
+    pool._cache.put_nowait(cached)
+    # reopening must not disturb the accounting of the allowed connections
+    dbs = [pool.connection() for _ in range(3)]
+    with pytest.raises(TooManyConnectionsError):
+        pool.connection()
+    while dbs:
+        dbs.pop().close()
+
+
 def test_context_manager():
     """Check that the connection can be used as a context manager."""
     pool = PooledPg(1, 1, 1)