TransactionalUndo closes storages it opens.

Fixes #268.

CHANGES.rst

@@ -13,12 +13,17 @@
 
 - Drop support for Python 3.4.
 
+- Fix ``DB.undo()`` and ``DB.undoMultiple()`` to close the storage
+  they open behind the scenes when the transaction is committed or
+  rolled back. See `issue 268
+  <https://github.com/zopefoundation/ZODB/issues/268>`_.
+
 5.5.1 (2018-10-25)
 ==================
 
 - Fix KeyError on releasing resources of a Connection when closing the DB.
-  This requires at least version 2.4 of the `transaction` package.
-  See `issue 208 <https://github.com/zopefoundation/ZODB/issues/208>`.
+  This requires at least version 2.4 of the ``transaction`` package.
+  See `issue 208 <https://github.com/zopefoundation/ZODB/issues/208>`_.
 
 5.5.0 (2018-10-13)
 ==================

src/ZODB/DB.py

@@ -14,6 +14,8 @@
 """Database objects
 """
 from __future__ import print_function
+
+import contextlib
 import sys
 import logging
 import datetime
@@ -1056,19 +1058,38 @@ class TransactionalUndo(object):
 
     def __init__(self, db, tids):
         self._db = db
-        self._storage = getattr(
-            db._mvcc_storage, 'undo_instance', db._mvcc_storage.new_instance)()
         self._tids = tids
+        self._storage = None
 
     def abort(self, transaction):
         pass
 
+    def close(self):
+        if self._storage is not None:
+            self._storage.close()
+            self._storage = None
+
     def tpc_begin(self, transaction):
+        assert self._storage is None, "Already in an active transaction"
+
         tdata = TransactionMetaData(
             transaction.user,
             transaction.description,
             transaction.extension)
         transaction.set_data(self, tdata)
+        # `undo_instance` is not part of any IStorage interface;
+        # it is defined in our MVCCAdapter. Regardless, we're opening
+        # a new storage instance, and so we must close it to be sure
+        # to reclaim resources in a timely manner.
+        #
+        # Once the tpc_begin method has been called, the transaction manager will
+        # guarantee to call either `tpc_finish` or `tpc_abort`, so those are the only
+        # methods we need to be concerned about calling close() from.
+        self._storage = getattr(
+            self._db._mvcc_storage,
+            'undo_instance',
+            self._db._mvcc_storage.new_instance)()
+
         self._storage.tpc_begin(tdata)
 
     def commit(self, transaction):
@@ -1081,15 +1102,23 @@ def tpc_vote(self, transaction):
         self._storage.tpc_vote(transaction)
 
     def tpc_finish(self, transaction):
-        transaction = transaction.data(self)
-        self._storage.tpc_finish(transaction)
+        with contextlib.closing(self):
+            transaction = transaction.data(self)
+            self._storage.tpc_finish(transaction)
 
     def tpc_abort(self, transaction):
-        transaction = transaction.data(self)
-        self._storage.tpc_abort(transaction)
+        with contextlib.closing(self):
+            transaction = transaction.data(self)
+            self._storage.tpc_abort(transaction)
 
     def sortKey(self):
-        return "%s:%s" % (self._storage.sortKey(), id(self))
+        # The transaction sorts data managers first before it calls
+        # `tpc_begin`, so we can't use our own storage because it's
+        # not open yet. Fortunately new_instances of a storage are
+        # supposed to return the same sort key as the original storage
+        # did.
+        return "%s:%s" % (self._db._mvcc_storage.sortKey(), id(self))
+
 
 def connection(*args, **kw):
     """Create a database :class:`connection <ZODB.Connection.Connection>`.

src/ZODB/tests/testDB.py

@@ -110,6 +110,107 @@ def check(info, text):
         check(db.undoLog(0, 3)  , True)
         check(db.undoInfo(0, 3) , True)
 
+
+class TransactionalUndoTests(unittest.TestCase):
+
+    def _makeOne(self):
+        from ZODB.DB import TransactionalUndo
+
+        class MockStorage(object):
+            instance_count = 0
+            close_count = 0
+            begin_count = 0
+            abort_count = 0
+
+            def new_instance(self):
+                self.instance_count += 1
+                return self
+
+            def tpc_begin(self, tx):
+                self.begin_count += 1
+
+            def close(self):
+                self.close_count += 1
+
+            def sortKey(self):
+                return 'MockStorage'
+
+        class MockDB(object):
+
+            def __init__(self):
+                self.storage = self._mvcc_storage = MockStorage()
+
+        return TransactionalUndo(MockDB(), [1])
+
+    def test_only_new_instance_on_begin(self):
+        undo = self._makeOne()
+        self.assertIsNone(undo._storage)
+        self.assertEqual(0, undo._db.storage.instance_count)
+
+        undo.tpc_begin(transaction.get())
+        self.assertIsNotNone(undo._storage)
+        self.assertEqual(1, undo._db.storage.instance_count)
+        self.assertEqual(1, undo._db.storage.begin_count)
+        self.assertIsNotNone(undo._storage)
+
+        # And we can't begin again
+        with self.assertRaises(AssertionError):
+            undo.tpc_begin(transaction.get())
+
+    def test_close_many(self):
+        undo = self._makeOne()
+        self.assertIsNone(undo._storage)
+        self.assertEqual(0, undo._db.storage.instance_count)
+
+        undo.close()
+        # Not open, didn't go through
+        self.assertEqual(0, undo._db.storage.close_count)
+
+        undo.tpc_begin(transaction.get())
+        undo.close()
+        undo.close()
+        self.assertEqual(1, undo._db.storage.close_count)
+        self.assertIsNone(undo._storage)
+
+    def test_sortKey(self):
+        # We get the same key whether or not we're open
+        undo = self._makeOne()
+        key = undo.sortKey()
+        self.assertIn('MockStorage', key)
+
+        undo.tpc_begin(transaction.get())
+        key2 = undo.sortKey()
+        self.assertEqual(key, key2)
+
+    def test_tpc_abort_closes(self):
+        undo = self._makeOne()
+        undo.tpc_begin(transaction.get())
+        undo._db.storage.tpc_abort = lambda tx: None
+        undo.tpc_abort(transaction.get())
+        self.assertEqual(1, undo._db.storage.close_count)
+
+    def test_tpc_abort_closes_on_exception(self):
+        undo = self._makeOne()
+        undo.tpc_begin(transaction.get())
+        with self.assertRaises(AttributeError):
+            undo.tpc_abort(transaction.get())
+        self.assertEqual(1, undo._db.storage.close_count)
+
+    def test_tpc_finish_closes(self):
+        undo = self._makeOne()
+        undo.tpc_begin(transaction.get())
+        undo._db.storage.tpc_finish = lambda tx: None
+        undo.tpc_finish(transaction.get())
+        self.assertEqual(1, undo._db.storage.close_count)
+
+    def test_tpc_finish_closes_on_exception(self):
+        undo = self._makeOne()
+        undo.tpc_begin(transaction.get())
+        with self.assertRaises(AttributeError):
+            undo.tpc_finish(transaction.get())
+        self.assertEqual(1, undo._db.storage.close_count)
+
+
 def test_invalidateCache():
     """The invalidateCache method invalidates a connection caches for all of
     the connections attached to a database::
@@ -423,7 +524,7 @@ def cleanup_on_close():
 """
 
 def test_suite():
-    s = unittest.makeSuite(DBTests)
+    s = unittest.defaultTestLoader.loadTestsFromName(__name__)
     s.addTest(doctest.DocTestSuite(
         setUp=ZODB.tests.util.setUp, tearDown=ZODB.tests.util.tearDown,
         checker=checker