X lib/zodb: Connection += onShutdownCallback

In the next patch we'll need this functionality to subscribe ZSync into
db.close to know when ZODB Connection is shut down even if it stays
alive referenced by some other objects.

lib/tests/test_zodb.py

@@ -28,6 +28,7 @@ from BTrees.IOBTree import IOBTree
 import transaction
 from transaction import TransactionManager
 from golang import defer, func
+import weakref, gc
 from pytest import raises
 import pytest; xfail = pytest.mark.xfail
 
@@ -354,6 +355,41 @@ def test_zodb_onresync():
     conn.close()
 
 
+# verify that ZODB.Connection.onShutdownCallback works
+@func
+def test_zodb_onshutdown():
+    stor = testdb.getZODBStorage()
+    defer(stor.close)
+    db  = DB(stor)
+
+    class T:
+        def __init__(t):
+            t.nshutdown = 0
+        def on_connection_shutdown(t):
+            t.nshutdown += 1
+
+    t1 = T()
+    t2 = T()
+    t3 = T()
+
+    # conn1 stays alive outside of db.pool
+    conn1 = db.open()
+    conn1.onShutdownCallback(t1)
+
+    # conn2 stays alive inside db.pool
+    conn2 = db.open()
+    conn2.onShutdownCallback(t2)
+    conn2.close()
+
+    assert t1.nshutdown == 0
+    assert t2.nshutdown == 0
+
+    # db.close triggers conn1 and conn2 shutdown
+    db.close()
+    assert t1.nshutdown == 1
+    assert t2.nshutdown == 1
+
+
 # test that zurl does not change from one open to another storage open.
 def test_zurlstable():
     if not isinstance(testdb, (testing.TestDB_FileStorage, testing.TestDB_ZEO, testing.TestDB_NEO)):

lib/zodb.py

@@ -296,6 +296,30 @@ else:
     raise AssertionError("ZODB3 is not supported anymore")
 
 
+# patch for ZODB.Connection to support callback on after database is closed
+ZODB.Connection.Connection._onShutdownCallbacks = None
+def Connection_onShutdownCallback(self, f):
+    if self._onShutdownCallbacks is None:
+        # NOTE WeakSet does not work for bound methods - they are always created
+        # anew for each obj.method access, and thus will go away almost immediately
+        self._onShutdownCallbacks = WeakSet()
+    self._onShutdownCallbacks.add(f)
+
+assert not hasattr(ZODB.Connection.Connection, 'onShutdownCallback')
+ZODB.Connection.Connection.onShutdownCallback = Connection_onShutdownCallback
+
+_orig_DB_close = ZODB.DB.close
+def _ZDB_close(self):
+    # the same code for ZODB3/4/5
+    @self._connectionMap
+    def _(conn):
+        if conn._onShutdownCallbacks:
+            for f in conn._onShutdownCallbacks:
+                f.on_connection_shutdown()
+    _orig_DB_close(self)
+
+ZODB.DB.close = _ZDB_close
+
 
 # zstor_2zurl converts a ZODB storage to URL to access it.
 def zstor_2zurl(zstor):