Merge lp:~jelmer/brz/write-group-context into lp:brz

[Jelmer Vernooij]

Add a WriteGroup() context manager.

Add a WriteGroup() context manager.

[Martin Packman (gz)]

Looks reasonable, thanks!

review: Approve

[The Breezy Bot (the-breezy-bot)]

Merging failed
https://ci.breezy-vcs.org/job/brz/job/brz-land/277/

[The Breezy Bot (the-breezy-bot)]

Merging failed
https://ci.breezy-vcs.org/job/brz/job/brz-land/284/

[The Breezy Bot (the-breezy-bot)]

Merging failed
https://ci.breezy-vcs.org/job/brz/job/brz-land/300/

[The Breezy Bot (the-breezy-bot)]

Running landing tests failed
https://ci.breezy-vcs.org/job/brz/job/brz-land/301/

=== modified file 'breezy/builtins.py'
--- breezy/builtins.py	2019-05-28 21:53:23 +0000
+++ breezy/builtins.py	2019-05-29 03:45:01 +0000
@@ -5260,29 +5260,18 @@
         return self._run(b, revision_id_list, revision)
 
     def _run(self, b, revision_id_list, revision):
+        from .repository import WriteGroup
         gpg_strategy = gpg.GPGStrategy(b.get_config_stack())
         if revision_id_list is not None:
-            b.repository.start_write_group()
-            try:
+            with WriteGroup(b.repository):
                 for revision_id in revision_id_list:
                     revision_id = cache_utf8.encode(revision_id)
                     b.repository.sign_revision(revision_id, gpg_strategy)
-            except BaseException:
-                b.repository.abort_write_group()
-                raise
-            else:
-                b.repository.commit_write_group()
         elif revision is not None:
             if len(revision) == 1:
                 revno, rev_id = revision[0].in_history(b)
-                b.repository.start_write_group()
-                try:
+                with WriteGroup(b.repository):
                     b.repository.sign_revision(rev_id, gpg_strategy)
-                except BaseException:
-                    b.repository.abort_write_group()
-                    raise
-                else:
-                    b.repository.commit_write_group()
             elif len(revision) == 2:
                 # are they both on rh- if so we can walk between them
                 # might be nice to have a range helper for arbitrary
@@ -5294,16 +5283,10 @@
                 if from_revno is None or to_revno is None:
                     raise errors.BzrCommandError(
                         gettext('Cannot sign a range of non-revision-history revisions'))
-                b.repository.start_write_group()
-                try:
+                with WriteGroup(b.repository):
                     for revno in range(from_revno, to_revno + 1):
                         b.repository.sign_revision(b.get_rev_id(revno),
                                                    gpg_strategy)
-                except BaseException:
-                    b.repository.abort_write_group()
-                    raise
-                else:
-                    b.repository.commit_write_group()
             else:
                 raise errors.BzrCommandError(
                     gettext('Please supply either one revision, or a range.'))
 
=== modified file 'breezy/bundle/serializer/v4.py'
--- breezy/bundle/serializer/v4.py	2018-11-11 04:08:32 +0000
+++ breezy/bundle/serializer/v4.py	2019-05-29 03:45:01 +0000
@@ -26,6 +26,7 @@
     lru_cache,
     multiparent,
     osutils,
+    repository as _mod_repository,
     revision as _mod_revision,
     trace,
     ui,
@@ -519,14 +520,8 @@
 
         Must be called with the Repository locked.
         """
-        self._repository.start_write_group()
-        try:
-            result = self._install_in_write_group()
-        except:
-            self._repository.abort_write_group()
-            raise
-        self._repository.commit_write_group()
-        return result
+        with _mod_repository.WriteGroup(self._repository):
+            return self._install_in_write_group()
 
     def _install_in_write_group(self):
         current_file = None
 
=== modified file 'breezy/bzr/vf_repository.py'
--- breezy/bzr/vf_repository.py	2019-02-04 19:39:30 +0000
+++ breezy/bzr/vf_repository.py	2019-05-29 03:45:01 +0000
@@ -68,6 +68,7 @@
     InterRepository,
     Repository,
     RepositoryFormat,
+    WriteGroup,
     )
 from .repository import (
     MetaDirRepository,
@@ -2863,8 +2864,7 @@
     Accepts an iterable of revision, tree, signature tuples.  The signature
     may be None.
     """
-    repository.start_write_group()
-    try:
+    with WriteGroup(repository):
         inventory_cache = lru_cache.LRUCache(10)
         for n, (revision, revision_tree, signature) in enumerate(iterable):
             _install_revision(repository, revision, revision_tree, signature,
@@ -2872,11 +2872,6 @@
             if pb is not None:
                 pb.update(gettext('Transferring revisions'),
                           n + 1, num_revisions)
-    except:
-        repository.abort_write_group()
-        raise
-    else:
-        repository.commit_write_group()
 
 
 def _install_revision(repository, rev, revision_tree, signature,
 
=== modified file 'breezy/commit_signature_commands.py'
--- breezy/commit_signature_commands.py	2018-11-12 01:41:38 +0000
+++ breezy/commit_signature_commands.py	2019-05-29 03:45:01 +0000
@@ -22,6 +22,7 @@
     controldir,
     errors,
     gpg,
+    repository as _mod_repository,
     revision as _mod_revision,
     )
 from .commands import Command
@@ -66,8 +67,7 @@
         count = 0
         with repo.lock_write():
             graph = repo.get_graph()
-            repo.start_write_group()
-            try:
+            with _mod_repository.WriteGroup(repo):
                 for rev_id, parents in graph.iter_ancestry(
                         [branch.last_revision()]):
                     if _mod_revision.is_null(rev_id):
@@ -86,11 +86,6 @@
                     count += 1
                     if not dry_run:
                         repo.sign_revision(rev_id, gpg_strategy)
-            except BaseException:
-                repo.abort_write_group()
-                raise
-            else:
-                repo.commit_write_group()
         self.outf.write(
             ngettext('Signed %d revision.\n', 'Signed %d revisions.\n',
                      count) % count)
 
=== modified file 'breezy/plugins/repodebug/missing_keys_for_stacking_fixer.py'
--- breezy/plugins/repodebug/missing_keys_for_stacking_fixer.py	2018-11-11 04:08:32 +0000
+++ breezy/plugins/repodebug/missing_keys_for_stacking_fixer.py	2019-05-29 03:45:01 +0000
@@ -21,6 +21,7 @@
 from ... import errors
 from ...sixish import viewvalues
 from ...bzr.vf_search import PendingAncestryResult
+from ...repository import WriteGroup
 from ...revision import NULL_REVISION
 
 
@@ -77,14 +78,8 @@
             assert raw_r._format.network_name() == b.repository._format.network_name()
             stream = b.repository.inventories.get_record_stream(
                 needed, 'topological', True)
-            raw_r.start_write_group()
-            try:
+            with WriteGroup(raw_r):
                 raw_r.inventories.insert_record_stream(stream)
-            except:
-                raw_r.abort_write_group()
-                raise
-            else:
-                raw_r.commit_write_group()
         finally:
             raw_r.unlock()
         b.unlock()
 
=== modified file 'breezy/repository.py'
--- breezy/repository.py	2019-02-04 19:39:30 +0000
+++ breezy/repository.py	2019-05-29 03:45:01 +0000
@@ -236,6 +236,28 @@
                                                       self.unlock)
 
 
+class WriteGroup(object):
+    """Context manager that manages a write group.
+
+    Raising an exception will result in the write group being aborted.
+    """
+
+    def __init__(self, repository, suppress_errors=False):
+        self.repository = repository
+        self._suppress_errors = suppress_errors
+
+    def __enter__(self):
+        self.repository.start_write_group()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if exc_type:
+            self.repository.abort_write_group(self._suppress_errors)
+            return False
+        else:
+            self.repository.commit_write_group()
+
+
 ######################################################################
 # Repositories
 
 
=== modified file 'breezy/tests/blackbox/test_reconcile.py'
--- breezy/tests/blackbox/test_reconcile.py	2018-11-11 04:08:32 +0000
+++ breezy/tests/blackbox/test_reconcile.py	2019-05-29 03:45:01 +0000
@@ -21,6 +21,7 @@
     controldir,
     tests,
     )
+from breezy.repository import WriteGroup
 from breezy.bzr import (
     inventory,
     )
@@ -53,10 +54,8 @@
         inv = inventory.Inventory(revision_id=b'missing')
         inv.root.revision = b'missing'
         repo.lock_write()
-        repo.start_write_group()
-        repo.add_inventory(b'missing', inv, [])
-        repo.commit_write_group()
-        repo.unlock()
+        with repo.lock_write(), WriteGroup(repo):
+            repo.add_inventory(b'missing', inv, [])
         (out, err) = self.run_bzr('reconcile')
         if repo._reconcile_backsup_inventory:
             does_backup_text = (
 
=== modified file 'breezy/tests/per_controldir/test_controldir.py'
--- breezy/tests/per_controldir/test_controldir.py	2018-11-11 04:08:32 +0000
+++ breezy/tests/per_controldir/test_controldir.py	2019-05-29 03:45:01 +0000
@@ -24,6 +24,7 @@
     errors,
     gpg,
     osutils,
+    repository as _mod_repository,
     revision as _mod_revision,
     transport,
     ui,
@@ -319,11 +320,8 @@
         tree_repo = tree.branch.repository
         if not tree_repo._format.supports_revision_signatures:
             self.skipTest('repository format does not support signing')
-        tree_repo.lock_write()
-        tree_repo.start_write_group()
-        tree_repo.sign_revision(rev1, gpg.LoopbackGPGStrategy(None))
-        tree_repo.commit_write_group()
-        tree_repo.unlock()
+        with tree_repo.lock_write(), _mod_repository.WriteGroup(tree_repo):
+            tree_repo.sign_revision(rev1, gpg.LoopbackGPGStrategy(None))
         target = self.make_branch('target')
         tree.branch.repository.copy_content_into(target.repository)
         tree.branch.copy_content_into(target)
 
=== modified file 'breezy/tests/per_interrepository/test_fetch.py'
--- breezy/tests/per_interrepository/test_fetch.py	2018-11-11 04:08:32 +0000
+++ breezy/tests/per_interrepository/test_fetch.py	2019-05-29 03:45:01 +0000
@@ -32,6 +32,7 @@
 from breezy.bzr.vf_search import (
     SearchResult,
     )
+from breezy.repository import WriteGroup
 from breezy.revision import (
     NULL_REVISION,
     Revision,
@@ -110,8 +111,7 @@
         source = tree.branch.repository
         source.lock_write()
         self.addCleanup(source.unlock)
-        source.start_write_group()
-        try:
+        with WriteGroup(source):
             # We need two revisions: OLD and NEW. NEW will claim to need a file
             # 'FOO' changed in 'OLD'. OLD will not have that file at all.
             source.texts.insert_record_stream([
@@ -133,11 +133,6 @@
                            revision_id=b'new',
                            parent_ids=[revid])
             source.add_revision(rev.revision_id, rev)
-        except:
-            source.abort_write_group()
-            raise
-        else:
-            source.commit_write_group()
         to_repo.fetch(source, b'new')
         to_repo.lock_read()
         self.addCleanup(to_repo.unlock)
@@ -448,17 +443,12 @@
         # We build a broken revision so that we can test the fetch code dies
         # properly. So copy the inventory and revision, but not the text.
         with to_repo.lock_write():
-            to_repo.start_write_group()
-            try:
+            with WriteGroup(to_repo, suppress_errors=True):
                 inv = tree.branch.repository.get_inventory(rev1)
                 to_repo.add_inventory(rev1, inv, [])
                 rev = tree.branch.repository.get_revision(rev1)
                 to_repo.add_revision(rev1, rev, inv=inv)
                 self.disable_commit_write_group_paranoia(to_repo)
-                to_repo.commit_write_group()
-            except:
-                to_repo.abort_write_group(suppress_errors=True)
-                raise
 
         # Implementations can either ensure that the target of the delta is
         # reconstructable, or raise an exception (which stream based copies
 
=== modified file 'breezy/tests/per_interrepository/test_interrepository.py'
--- breezy/tests/per_interrepository/test_interrepository.py	2018-11-11 04:08:32 +0000
+++ breezy/tests/per_interrepository/test_interrepository.py	2019-05-29 03:45:01 +0000
@@ -22,6 +22,7 @@
 import breezy.errors as errors
 import breezy.gpg
 from breezy.bzr.inventory import Inventory
+from breezy.repository import WriteGroup
 from breezy.revision import NULL_REVISION
 from breezy.tests import (
     TestNotApplicable,
@@ -65,13 +66,10 @@
         tree_a = self.make_branch_and_tree('a')
         self.controldir = tree_a.branch.controldir
         # add a corrupt inventory 'orphan'
-        tree_a.branch.repository.lock_write()
-        tree_a.branch.repository.start_write_group()
-        if tree_a.branch.repository._format.supports_ghosts:
-            inv_file = tree_a.branch.repository.inventories
-            inv_file.add_lines((b'orphan',), [], [])
-        tree_a.branch.repository.commit_write_group()
-        tree_a.branch.repository.unlock()
+        with tree_a.branch.repository.lock_write(), WriteGroup(tree_a.branch.repository):
+            if tree_a.branch.repository._format.supports_ghosts:
+                inv_file = tree_a.branch.repository.inventories
+                inv_file.add_lines((b'orphan',), [], [])
         # add a real revision 'rev1'
         self.rev1 = tree_a.commit('rev1', allow_pointless=True)
         # add a real revision 'rev2' based on rev1
@@ -149,12 +147,10 @@
                 'to repository does not support signatures')
         # and sign 'rev2'
         tree_a = WorkingTree.open('a')
-        tree_a.branch.repository.lock_write()
-        tree_a.branch.repository.start_write_group()
-        tree_a.branch.repository.sign_revision(self.rev2,
-                                               breezy.gpg.LoopbackGPGStrategy(None))
-        tree_a.branch.repository.commit_write_group()
-        tree_a.branch.repository.unlock()
+        with tree_a.branch.repository.lock_write(), \
+                WriteGroup(tree_a.branch.repository):
+            tree_a.branch.repository.sign_revision(
+                self.rev2, breezy.gpg.LoopbackGPGStrategy(None))
 
         from_repo = self.controldir.open_repository()
         from_signature = from_repo.get_signature_text(self.rev2)
 
=== modified file 'breezy/tests/per_pack_repository.py'
--- breezy/tests/per_pack_repository.py	2019-01-09 00:04:02 +0000
+++ breezy/tests/per_pack_repository.py	2019-05-29 03:45:01 +0000
@@ -473,26 +473,15 @@
         r1 = repository.Repository.open('.')
         r2 = repository.Repository.open('.')
         # add a pack to drop
-        r1.lock_write()
-        try:
-            r1.start_write_group()
-            try:
+        with r1.lock_write():
+            with repository.WriteGroup(r1):
                 self._add_text(r1, b'fileidr1')
-            except:
-                r1.abort_write_group()
-                raise
-            else:
-                r1.commit_write_group()
             r1._pack_collection.ensure_loaded()
             name_to_drop = r1._pack_collection.all_packs()[0].name
-        finally:
-            r1.unlock()
-        r1.lock_write()
-        try:
+        with r1.lock_write():
             # access enough data to load the names list
             list(r1.all_revision_ids())
-            r2.lock_write()
-            try:
+            with r2.lock_write():
                 # access enough data to load the names list
                 list(r2.all_revision_ids())
                 r1._pack_collection.ensure_loaded()
@@ -533,10 +522,6 @@
                                  r2._pack_collection.names())
                 self.assertEqual(1, len(r1._pack_collection.names()))
                 self.assertFalse(name_to_drop in r1._pack_collection.names())
-            finally:
-                r2.unlock()
-        finally:
-            r1.unlock()
 
     def test_concurrent_pack_triggers_reload(self):
         # create 2 packs, which we will then collapse
 
=== modified file 'breezy/tests/per_repository/test_signatures.py'
--- breezy/tests/per_repository/test_signatures.py	2018-11-29 23:42:41 +0000
+++ breezy/tests/per_repository/test_signatures.py	2019-05-29 03:45:01 +0000
@@ -22,7 +22,7 @@
     tests,
     urlutils,
     )
-
+from breezy.repository import WriteGroup
 from breezy.bzr.testament import Testament
 from breezy.tests import per_repository
 
@@ -58,24 +58,14 @@
     def test_store_signature(self):
         wt = self.make_branch_and_tree('.')
         branch = wt.branch
-        branch.lock_write()
-        try:
-            branch.repository.start_write_group()
+        with branch.lock_write(), WriteGroup(branch.repository):
             try:
                 branch.repository.store_revision_signature(
                     gpg.LoopbackGPGStrategy(None), b'FOO', b'A')
             except errors.NoSuchRevision:
-                branch.repository.abort_write_group()
                 raise tests.TestNotApplicable(
                     "repository does not support signing non-present"
                     "revisions")
-            except:
-                branch.repository.abort_write_group()
-                raise
-            else:
-                branch.repository.commit_write_group()
-        finally:
-            branch.unlock()
         # A signature without a revision should not be accessible.
         self.assertRaises(errors.NoSuchRevision,
                           branch.repository.has_signature_for_revision_id,
 
=== modified file 'breezy/tests/per_repository_chk/test_supported.py'
--- breezy/tests/per_repository_chk/test_supported.py	2018-11-16 11:37:47 +0000
+++ breezy/tests/per_repository_chk/test_supported.py	2019-05-29 03:45:01 +0000
@@ -38,24 +38,14 @@
 
     def test_add_bytes_to_chk_bytes_store(self):
         repo = self.make_repository('.')
-        repo.lock_write()
-        try:
-            repo.start_write_group()
-            try:
-                sha1, len, _ = repo.chk_bytes.add_lines((None,),
-                                                        None, [b"foo\n", b"bar\n"], random_id=True)
-                self.assertEqual(b'4e48e2c9a3d2ca8a708cb0cc545700544efb5021',
-                                 sha1)
-                self.assertEqual(
-                    {(b'sha1:4e48e2c9a3d2ca8a708cb0cc545700544efb5021',)},
-                    repo.chk_bytes.keys())
-            except:
-                repo.abort_write_group()
-                raise
-            else:
-                repo.commit_write_group()
-        finally:
-            repo.unlock()
+        with repo.lock_write(), repository.WriteGroup(repo):
+            sha1, len, _ = repo.chk_bytes.add_lines(
+                (None,), None, [b"foo\n", b"bar\n"], random_id=True)
+            self.assertEqual(
+                b'4e48e2c9a3d2ca8a708cb0cc545700544efb5021', sha1)
+            self.assertEqual(
+                {(b'sha1:4e48e2c9a3d2ca8a708cb0cc545700544efb5021',)},
+                repo.chk_bytes.keys())
         # And after an unlock/lock pair
         with repo.lock_read():
             self.assertEqual(
@@ -76,46 +66,25 @@
         node_sha1 = osutils.sha_strings(node_lines)
         expected_set = {(b'sha1:' + leaf_sha1,), (b'sha1:' + node_sha1,)}
         repo = self.make_repository('.')
-        repo.lock_write()
-        try:
-            repo.start_write_group()
-            try:
+        with repo.lock_write():
+            with repository.WriteGroup(repo):
                 # Internal node pointing at a leaf.
-                repo.chk_bytes.add_lines(
-                    (None,), None, node_lines, random_id=True)
-            except:
-                repo.abort_write_group()
-                raise
-            else:
-                repo.commit_write_group()
-            repo.start_write_group()
-            try:
+                repo.chk_bytes.add_lines((None,), None, node_lines, random_id=True)
+            with repository.WriteGroup(repo):
                 # Leaf in a separate pack.
-                repo.chk_bytes.add_lines(
-                    (None,), None, leaf_lines, random_id=True)
-            except:
-                repo.abort_write_group()
-                raise
-            else:
-                repo.commit_write_group()
+                repo.chk_bytes.add_lines((None,), None, leaf_lines, random_id=True)
             repo.pack()
             self.assertEqual(expected_set, repo.chk_bytes.keys())
-        finally:
-            repo.unlock()
         # and reopening
         repo = repo.controldir.open_repository()
-        repo.lock_read()
-        try:
+        with repo.lock_read():
             self.assertEqual(expected_set, repo.chk_bytes.keys())
-        finally:
-            repo.unlock()
 
     def test_chk_bytes_are_fully_buffered(self):
         repo = self.make_repository('.')
         repo.lock_write()
         self.addCleanup(repo.unlock)
-        repo.start_write_group()
-        try:
+        with repository.WriteGroup(repo):
             sha1, len, _ = repo.chk_bytes.add_lines((None,),
                                                     None, [b"foo\n", b"bar\n"], random_id=True)
             self.assertEqual(b'4e48e2c9a3d2ca8a708cb0cc545700544efb5021',
@@ -123,11 +92,6 @@
             self.assertEqual(
                 {(b'sha1:4e48e2c9a3d2ca8a708cb0cc545700544efb5021',)},
                 repo.chk_bytes.keys())
-        except:
-            repo.abort_write_group()
-            raise
-        else:
-            repo.commit_write_group()
         # This may not always be correct if we change away from BTreeGraphIndex
         # in the future. But for now, lets check that chk_bytes are fully
         # buffered
 
=== modified file 'breezy/tests/per_repository_reference/test_add_inventory.py'
--- breezy/tests/per_repository_reference/test_add_inventory.py	2018-11-11 04:08:32 +0000
+++ breezy/tests/per_repository_reference/test_add_inventory.py	2019-05-29 03:45:01 +0000
@@ -17,6 +17,7 @@
 """Tests for add_inventory on a repository with external references."""
 
 from breezy import errors
+from breezy.repository import WriteGroup
 from breezy.tests.per_repository_reference import (
     TestCaseWithExternalReferenceRepository,
     )
@@ -34,18 +35,8 @@
         self.addCleanup(tree.unlock)
         base = self.make_repository('base')
         repo = self.make_referring('referring', base)
-        repo.lock_write()
-        try:
-            repo.start_write_group()
-            try:
-                repo.add_inventory(revid, inv, [])
-            except:
-                repo.abort_write_group()
-                raise
-            else:
-                repo.commit_write_group()
-        finally:
-            repo.unlock()
+        with repo.lock_write(), WriteGroup(repo):
+            repo.add_inventory(revid, inv, [])
         repo.lock_read()
         self.addCleanup(repo.unlock)
         inv2 = repo.get_inventory(revid)
 
=== modified file 'breezy/tests/per_repository_reference/test_add_revision.py'
--- breezy/tests/per_repository_reference/test_add_revision.py	2017-05-21 18:10:28 +0000
+++ breezy/tests/per_repository_reference/test_add_revision.py	2019-05-29 03:45:01 +0000
@@ -17,6 +17,7 @@
 """Tests for add_revision on a repository with external references."""
 
 from breezy import errors
+from breezy.repository import WriteGroup
 from breezy.tests.per_repository_reference import (
     TestCaseWithExternalReferenceRepository,
     )
@@ -35,20 +36,10 @@
         rev = tree.branch.repository.get_revision(revid)
         base = self.make_repository('base')
         repo = self.make_referring('referring', base)
-        repo.lock_write()
-        try:
-            repo.start_write_group()
-            try:
-                rev = tree.branch.repository.get_revision(revid)
-                repo.texts.add_lines((inv.root.file_id, revid), [], [])
-                repo.add_revision(revid, rev, inv=inv)
-            except:
-                repo.abort_write_group()
-                raise
-            else:
-                repo.commit_write_group()
-        finally:
-            repo.unlock()
+        with repo.lock_write(), WriteGroup(repo):
+            rev = tree.branch.repository.get_revision(revid)
+            repo.texts.add_lines((inv.root.file_id, revid), [], [])
+            repo.add_revision(revid, rev, inv=inv)
         rev2 = repo.get_revision(revid)
         self.assertEqual(rev, rev2)
         self.assertRaises(errors.NoSuchRevision, base.get_revision, revid)
 
=== modified file 'breezy/tests/per_repository_reference/test_add_signature_text.py'
--- breezy/tests/per_repository_reference/test_add_signature_text.py	2018-11-11 04:08:32 +0000
+++ breezy/tests/per_repository_reference/test_add_signature_text.py	2019-05-29 03:45:01 +0000
@@ -17,6 +17,7 @@
 """Tests for add_signature_text on a repository with external references."""
 
 from breezy import errors
+from breezy.repository import WriteGroup
 from breezy.tests.per_repository_reference import (
     TestCaseWithExternalReferenceRepository,
     )
@@ -34,20 +35,11 @@
         self.addCleanup(tree.unlock)
         base = self.make_repository('base')
         repo = self.make_referring('referring', base)
-        repo.lock_write()
-        try:
-            repo.start_write_group()
-            try:
-                rev = tree.branch.repository.get_revision(revid)
-                repo.texts.add_lines((inv.root.file_id, revid), [], [])
-                repo.add_revision(revid, rev, inv=inv)
-                repo.add_signature_text(revid, b"text")
-                repo.commit_write_group()
-            except:
-                repo.abort_write_group()
-                raise
-        finally:
-            repo.unlock()
+        with repo.lock_write(), WriteGroup(repo):
+            rev = tree.branch.repository.get_revision(revid)
+            repo.texts.add_lines((inv.root.file_id, revid), [], [])
+            repo.add_revision(revid, rev, inv=inv)
+            repo.add_signature_text(revid, b"text")
         repo.get_signature_text(revid)
         self.assertRaises(errors.NoSuchRevision, base.get_signature_text,
                           revid)
 
=== modified file 'breezy/tests/per_repository_vf/helpers.py'
--- breezy/tests/per_repository_vf/helpers.py	2018-11-11 04:08:32 +0000
+++ breezy/tests/per_repository_vf/helpers.py	2019-05-29 03:45:01 +0000
@@ -20,6 +20,7 @@
     osutils,
     revision as _mod_revision,
     )
+from ...repository import WriteGroup
 from ...bzr import (
     inventory,
     )
@@ -50,9 +51,7 @@
                 "%s isn't a knit format" % self.repository_format)
 
         repo = self.make_repository('broken')
-        repo.lock_write()
-        repo.start_write_group()
-        try:
+        with repo.lock_write(), WriteGroup(repo):
             inv = inventory.Inventory(revision_id=b'revision-id')
             inv.root.revision = b'revision-id'
             inv_sha1 = repo.add_inventory(b'revision-id', inv, [])
@@ -66,16 +65,9 @@
             # Manually add the revision text using the RevisionStore API, with
             # bad parents.
             rev_text = repo._serializer.write_revision_to_string(revision)
-            repo.revisions.add_lines((revision.revision_id,),
-                                     [(b'incorrect-parent',)],
-                                     osutils.split_lines(rev_text))
-        except:
-            repo.abort_write_group()
-            repo.unlock()
-            raise
-        else:
-            repo.commit_write_group()
-            repo.unlock()
+            repo.revisions.add_lines(
+                (revision.revision_id,), [(b'incorrect-parent',)],
+                osutils.split_lines(rev_text))
 
         repo.lock_write()
         self.addCleanup(repo.unlock)
 
=== modified file 'breezy/tests/per_repository_vf/test_check_reconcile.py'
--- breezy/tests/per_repository_vf/test_check_reconcile.py	2018-11-11 04:08:32 +0000
+++ breezy/tests/per_repository_vf/test_check_reconcile.py	2019-05-29 03:45:01 +0000
@@ -25,6 +25,7 @@
     Inventory,
     InventoryFile,
     )
+from breezy.repository import WriteGroup
 from breezy.revision import (
     NULL_REVISION,
     Revision,
@@ -793,15 +794,9 @@
     def make_populated_repository(self, factory):
         """Create a new repository populated by the given factory."""
         repo = self.make_repository('broken-repo')
-        with repo.lock_write():
-            repo.start_write_group()
-            try:
-                factory(repo)
-                repo.commit_write_group()
-                return repo
-            except:
-                repo.abort_write_group()
-                raise
+        with repo.lock_write(), WriteGroup(repo):
+            factory(repo)
+            return repo
 
     def add_revision(self, repo, revision_id, inv, parent_ids):
         """Add a revision with a given inventory and parents to a repository.
 
=== modified file 'breezy/tests/per_repository_vf/test_repository.py'
--- breezy/tests/per_repository_vf/test_repository.py	2019-01-01 21:08:01 +0000
+++ breezy/tests/per_repository_vf/test_repository.py	2019-05-29 03:45:01 +0000
@@ -327,19 +327,9 @@
         self.controldir = tree_a.branch.controldir
         # add a corrupt inventory 'orphan'
         # this may need some generalising for knits.
-        tree_a.lock_write()
-        try:
-            tree_a.branch.repository.start_write_group()
-            try:
-                inv_file = tree_a.branch.repository.inventories
-                inv_file.add_lines((b'orphan',), [], [])
-            except:
-                tree_a.branch.repository.commit_write_group()
-                raise
-            else:
-                tree_a.branch.repository.abort_write_group()
-        finally:
-            tree_a.unlock()
+        with tree_a.lock_write(), _mod_repository.WriteGroup(tree_a.branch.repository):
+            inv_file = tree_a.branch.repository.inventories
+            inv_file.add_lines((b'orphan',), [], [])
         # add a real revision 'rev1'
         tree_a.commit('rev1', rev_id=b'rev1', allow_pointless=True)
         # add a real revision 'rev2' based on rev1
@@ -385,18 +375,13 @@
 
     def test_reserved_id(self):
         repo = self.make_repository('repository')
-        repo.lock_write()
-        repo.start_write_group()
-        try:
+        with repo.lock_write(), _mod_repository.WriteGroup(repo):
             self.assertRaises(errors.ReservedId, repo.add_inventory,
                               b'reserved:', None, None)
             self.assertRaises(errors.ReservedId, repo.add_inventory_by_delta,
                               "foo", [], b'reserved:', None)
-            self.assertRaises(errors.ReservedId, repo.add_revision,
-                              b'reserved:', None)
-        finally:
-            repo.abort_write_group()
-            repo.unlock()
+            self.assertRaises(
+                errors.ReservedId, repo.add_revision, b'reserved:', None)
 
 
 class TestCaseWithCorruptRepository(TestCaseWithRepository):
 
=== modified file 'breezy/tests/test_inv.py'
--- breezy/tests/test_inv.py	2018-11-16 11:37:47 +0000
+++ breezy/tests/test_inv.py	2019-05-29 03:45:01 +0000
@@ -141,8 +141,7 @@
 
 
 def _create_repo_revisions(repo, basis, delta, invalid_delta):
-    repo.start_write_group()
-    try:
+    with repository.WriteGroup(repo):
         rev = revision.Revision(b'basis', timestamp=0, timezone=None,
                                 message="", committer="[e-mail]")
         basis.revision_id = b'basis'
@@ -161,19 +160,13 @@
         rev = revision.Revision(b'result', timestamp=0, timezone=None,
                                 message="", committer="[e-mail]")
         repo.add_revision(b'result', rev, result_inv)
-        repo.commit_write_group()
-    except:
-        repo.abort_write_group()
-        raise
     return target_entries
 
 
 def _get_basis_entries(tree):
     basis_tree = tree.basis_tree()
-    basis_tree.lock_read()
-    basis_tree_entries = list(basis_tree.inventory.iter_entries_by_dir())
-    basis_tree.unlock()
-    return basis_tree_entries
+    with basis_tree.lock_read():
+        return list(basis_tree.inventory.iter_entries_by_dir())
 
 
 def _populate_different_tree(tree, basis, delta):
@@ -253,28 +246,16 @@
     format = self.format()
     control = self.make_controldir('tree', format=format._matchingcontroldir)
     repo = format.initialize(control)
-    with repo.lock_write():
-        repo.start_write_group()
-        try:
-            rev = revision.Revision(b'basis', timestamp=0, timezone=None,
-                                    message="", committer="[e-mail]")
-            basis.revision_id = b'basis'
-            create_texts_for_inv(repo, basis)
-            repo.add_revision(b'basis', rev, basis)
-            repo.commit_write_group()
-        except:
-            repo.abort_write_group()
-            raise
-    with repo.lock_write():
-        repo.start_write_group()
-        try:
-            inv_sha1 = repo.add_inventory_by_delta(b'basis', delta,
-                                                   b'result', [b'basis'])
-        except:
-            repo.abort_write_group()
-            raise
-        else:
-            repo.commit_write_group()
+    with repo.lock_write(), repository.WriteGroup(repo):
+        rev = revision.Revision(
+            b'basis', timestamp=0, timezone=None, message="",
+            committer="[e-mail]")
+        basis.revision_id = b'basis'
+        create_texts_for_inv(repo, basis)
+        repo.add_revision(b'basis', rev, basis)
+    with repo.lock_write(), repository.WriteGroup(repo):
+        inv_sha1 = repo.add_inventory_by_delta(
+            b'basis', delta, b'result', [b'basis'])
     # Fresh lock, reads disk again.
     repo = repo.controldir.open_repository()
     repo.lock_read()