--- a/mercurial/bookmarks.py	Mon Jan 31 08:44:48 2022 +0100
+++ b/mercurial/bookmarks.py	Mon Jan 31 18:38:15 2022 +0100
@@ -212,7 +212,11 @@
         The transaction is then responsible for updating the file content."""
         location = b'' if bookmarksinstore(self._repo) else b'plain'
         tr.addfilegenerator(
-            b'bookmarks', (b'bookmarks',), self._write, location=location
+            b'bookmarks',
+            (b'bookmarks',),
+            self._write,
+            location=location,
+            post_finalize=True,
         )
         tr.hookargs[b'bookmark_moved'] = b'1'
 

--- a/mercurial/dirstate.py	Mon Jan 31 08:44:48 2022 +0100
+++ b/mercurial/dirstate.py	Mon Jan 31 18:38:15 2022 +0100
@@ -730,12 +730,14 @@
                     (self._filename_tk,),
                     lambda f: self._write_tracked_key(tr, f),
                     location=b'plain',
+                    post_finalize=True,
                 )
             tr.addfilegenerator(
                 b'dirstate-1-main',
                 (self._filename,),
                 lambda f: self._writedirstate(tr, f),
                 location=b'plain',
+                post_finalize=True,
             )
             if write_key:
                 tr.addfilegenerator(
@@ -743,6 +745,7 @@
                     (self._filename_tk,),
                     lambda f: self._write_tracked_key(tr, f),
                     location=b'plain',
+                    post_finalize=True,
                 )
             return
 
@@ -1425,6 +1428,7 @@
                 (self._filename,),
                 lambda f: self._writedirstate(tr, f),
                 location=b'plain',
+                post_finalize=True,
             )
 
             # ensure that pending file written above is unlinked at

--- a/mercurial/transaction.py	Mon Jan 31 08:44:48 2022 +0100
+++ b/mercurial/transaction.py	Mon Jan 31 18:38:15 2022 +0100
@@ -25,16 +25,6 @@
 
 version = 2
 
-# These are the file generators that should only be executed after the
-# finalizers are done, since they rely on the output of the finalizers (like
-# the changelog having been written).
-postfinalizegenerators = {
-    b'bookmarks',
-    b'dirstate-0-key-pre',
-    b'dirstate-1-main',
-    b'dirstate-2-key-post',
-}
-
 GEN_GROUP_ALL = b'all'
 GEN_GROUP_PRE_FINALIZE = b'prefinalize'
 GEN_GROUP_POST_FINALIZE = b'postfinalize'
@@ -339,7 +329,13 @@
 
     @active
     def addfilegenerator(
-        self, genid, filenames, genfunc, order=0, location=b''
+        self,
+        genid,
+        filenames,
+        genfunc,
+        order=0,
+        location=b'',
+        post_finalize=False,
     ):
         """add a function to generates some files at transaction commit
 
@@ -362,10 +358,14 @@
         The `location` arguments may be used to indicate the files are located
         outside of the the standard directory for transaction. It should match
         one of the key of the `transaction.vfsmap` dictionary.
+
+        The `post_finalize` argument can be set to `True` for file generation
+        that must be run after the transaction has been finalized.
         """
         # For now, we are unable to do proper backup and restore of custom vfs
         # but for bookmarks that are handled outside this mechanism.
-        self._filegenerators[genid] = (order, filenames, genfunc, location)
+        entry = (order, filenames, genfunc, location, post_finalize)
+        self._filegenerators[genid] = entry
 
     @active
     def removefilegenerator(self, genid):
@@ -385,13 +385,12 @@
 
         for id, entry in sorted(pycompat.iteritems(self._filegenerators)):
             any = True
-            order, filenames, genfunc, location = entry
+            order, filenames, genfunc, location, post_finalize = entry
 
             # for generation at closing, check if it's before or after finalize
-            is_post = id in postfinalizegenerators
-            if skip_post and is_post:
+            if skip_post and post_finalize:
                 continue
-            elif skip_pre and not is_post:
+            elif skip_pre and not post_finalize:
                 continue
 
             vfs = self._vfsmap[location]

--- a/tests/testlib/crash_transaction_late.py	Mon Jan 31 08:44:48 2022 +0100
+++ b/tests/testlib/crash_transaction_late.py	Mon Jan 31 18:38:15 2022 +0100
@@ -9,7 +9,6 @@
 
 from mercurial import (
     error,
-    transaction,
 )
 
 
@@ -18,14 +17,15 @@
 
 
 def reposetup(ui, repo):
-
-    transaction.postfinalizegenerators.add(b'late-abort')
-
     class LateAbortRepo(repo.__class__):
         def transaction(self, *args, **kwargs):
             tr = super(LateAbortRepo, self).transaction(*args, **kwargs)
             tr.addfilegenerator(
-                b'late-abort', [b'late-abort'], abort, order=9999999
+                b'late-abort',
+                [b'late-abort'],
+                abort,
+                order=9999999,
+                post_finalize=True,
             )
             return tr