D3559: narrow: only wrap dirstate functions once, instead of per-reposetup

[spectral (Kyle Lippincott)]

spectral created this revision.
Herald added a reviewer: durin42.
Herald added a subscriber: mercurial-devel.
Herald added a reviewer: hg-reviewers.

REVISION SUMMARY
  chg will call reposetup multiple times, and we will end up double-wrapping (or
  worse) the dirstate functions; this can cause issues like OSError 'No such file
  or directory' during rebase operations, when we go to double-delete our
  narrowspec backup file.

REPOSITORY
  rHG Mercurial

REVISION DETAIL
  https://phab.mercurial-scm.org/D3559

AFFECTED FILES
  hgext/narrow/__init__.py
  hgext/narrow/narrowdirstate.py
  tests/test-narrow-expanddirstate.t

CHANGE DETAILS

diff --git a/tests/test-narrow-expanddirstate.t b/tests/test-narrow-expanddirstate.t
--- a/tests/test-narrow-expanddirstate.t
+++ b/tests/test-narrow-expanddirstate.t
@@ -56,6 +56,8 @@
   > from mercurial import patch
   > from mercurial import util as hgutil
   > 
+  > _repo = None
+  > 
   > def expandnarrowspec(ui, repo, newincludes=None):
   >   if not newincludes:
   >     return
@@ -72,28 +74,37 @@
   >   for f in repo[b'.'].manifest().walk(added):
   >     repo.dirstate.normallookup(f)
   > 
-  > def makeds(ui, repo):
+  > def makeds(ui):
   >   def wrapds(orig, self):
   >     ds = orig(self)
   >     class expandingdirstate(ds.__class__):
   >       @hgutil.propertycache
   >       def _map(self):
   >         ret = super(expandingdirstate, self)._map
-  >         with repo.wlock(), repo.lock(), repo.transaction(
+  >         with _repo.wlock(), _repo.lock(), _repo.transaction(
   >             b'expandnarrowspec'):
-  >           expandnarrowspec(ui, repo,
+  >           expandnarrowspec(ui, _repo,
   >                            encoding.environ.get(b'DIRSTATEINCLUDES'))
   >         return ret
   >     ds.__class__ = expandingdirstate
   >     return ds
   >   return wrapds
   > 
   > def reposetup(ui, repo):
+  >   # Normally there'd be checks that this is a narrowed localrepo, and uses
+  >   # of _repo would skip it if it's not specified; we aren't doing that for
+  >   # this toy version.
+  >   global _repo
+  >   assert _repo is None or _repo.root == repo.root, "%s %s" % (
+  >       _repo.root, repo.root)
+  >   _repo = repo
+  > 
+  > def extsetup(ui):
   >   extensions.wrapfilecache(localrepo.localrepository, b'dirstate',
-  >                            makeds(ui, repo))
+  >                            makeds(ui))
   >   def overridepatch(orig, *args, **kwargs):
-  >     with repo.wlock():
-  >       expandnarrowspec(ui, repo, encoding.environ.get(b'PATCHINCLUDES'))
+  >     with _repo.wlock():
+  >       expandnarrowspec(ui, _repo, encoding.environ.get(b'PATCHINCLUDES'))
   >       return orig(*args, **kwargs)
   > 
   >   extensions.wrapfunction(patch, b'patch', overridepatch)
diff --git a/hgext/narrow/narrowdirstate.py b/hgext/narrow/narrowdirstate.py
--- a/hgext/narrow/narrowdirstate.py
+++ b/hgext/narrow/narrowdirstate.py
@@ -9,20 +9,45 @@
 
 from mercurial.i18n import _
 from mercurial import (
+    changegroup,
     dirstate,
     error,
     extensions,
     match as matchmod,
     narrowspec,
     util as hgutil,
 )
 
-def setup(repo):
+# Mapping of root:str to repo for repos that have the narrow requirement
+# specified.
+_rootrepomap = {}
+
+def _getrepo(ds):
+    """Check if narrow is enabled for repo associated with `ds`; return repo."""
+    return _rootrepomap.get(ds._root, None)
+
+def reposetup(repo):
+    """Record that we want to do narrowing on this repo.
+
+    We need to wrap the dirstate functions only once - when using chg, reposetup
+    is potentially called multiple times on the same repo for the lifetime of
+    the process, so we wrap in extsetup() and do a check in every function for
+    whether or not it's enabled for that repo, in case the same chg process is
+    being used in a repo that's *not* narrowed.
+    """
+    # This should have been done by __init__.py already, but let's be paranoid.
+    assert changegroup.NARROW_REQUIREMENT in repo.requirements
+
+    global _rootrepomap
+    _rootrepomap[repo.root] = repo
+
+def extsetup():
     """Add narrow spec dirstate ignore, block changes outside narrow spec."""
 
     def walk(orig, self, match, subrepos, unknown, ignored, full=True,
              narrowonly=True):
-        if narrowonly:
+        repo = _getrepo(self)
+        if repo and narrowonly:
             # hack to not exclude explicitly-specified paths so that they can
             # be warned later on e.g. dirstate.add()
             em = matchmod.exact(match._root, match._cwd, match.files())
@@ -36,6 +61,10 @@
     editfuncs = ['normal', 'add', 'normallookup', 'copy', 'remove', 'merge']
     for func in editfuncs:
         def _wrapper(orig, self, *args):
+            repo = _getrepo(self)
+            if not repo:
+                return orig(self, *args)
+
             dirstate = repo.dirstate
             narrowmatch = repo.narrowmatch()
             for f in args:
@@ -46,7 +75,8 @@
         extensions.wrapfunction(dirstate.dirstate, func, _wrapper)
 
     def filterrebuild(orig, self, parent, allfiles, changedfiles=None):
-        if changedfiles is None:
+        repo = _getrepo(self)
+        if repo and changedfiles is None:
             # Rebuilding entire dirstate, let's filter allfiles to match the
             # narrowspec.
             allfiles = [f for f in allfiles if repo.narrowmatch()(f)]
@@ -59,24 +89,27 @@
         return backupname.replace('dirstate', narrowspec.FILENAME)
 
     def restorebackup(orig, self, tr, backupname):
-        self._opener.rename(_narrowbackupname(backupname), narrowspec.FILENAME,
-                            checkambig=True)
+        if _getrepo(self):
+            self._opener.rename(_narrowbackupname(backupname),
+                                narrowspec.FILENAME, checkambig=True)
         orig(self, tr, backupname)
 
     extensions.wrapfunction(dirstate.dirstate, 'restorebackup', restorebackup)
 
     def savebackup(orig, self, tr, backupname):
         orig(self, tr, backupname)
 
-        narrowbackupname = _narrowbackupname(backupname)
-        self._opener.tryunlink(narrowbackupname)
-        hgutil.copyfile(self._opener.join(narrowspec.FILENAME),
-                        self._opener.join(narrowbackupname), hardlink=True)
+        if _getrepo(self):
+            narrowbackupname = _narrowbackupname(backupname)
+            self._opener.tryunlink(narrowbackupname)
+            hgutil.copyfile(self._opener.join(narrowspec.FILENAME),
+                            self._opener.join(narrowbackupname), hardlink=True)
 
     extensions.wrapfunction(dirstate.dirstate, 'savebackup', savebackup)
 
     def clearbackup(orig, self, tr, backupname):
         orig(self, tr, backupname)
-        self._opener.unlink(_narrowbackupname(backupname))
+        if _getrepo(self):
+            self._opener.unlink(_narrowbackupname(backupname))
 
     extensions.wrapfunction(dirstate.dirstate, 'clearbackup', clearbackup)
diff --git a/hgext/narrow/__init__.py b/hgext/narrow/__init__.py
--- a/hgext/narrow/__init__.py
+++ b/hgext/narrow/__init__.py
@@ -77,7 +77,7 @@
     narrowrepo.wraprepo(repo)
     if changegroup.NARROW_REQUIREMENT in repo.requirements:
         narrowcopies.setup(repo)
-        narrowdirstate.setup(repo)
+        narrowdirstate.reposetup(repo)
         narrowpatch.setup(repo)
         narrowwirepeer.reposetup(repo)
 
@@ -92,6 +92,7 @@
     extensions.wrapfunction(verifymod.verifier, '__init__', _verifierinit)
     extensions.wrapfunction(hg, 'postshare', narrowrepo.wrappostshare)
     extensions.wrapfunction(hg, 'copystore', narrowrepo.unsharenarrowspec)
+    narrowdirstate.extsetup()
 
 templatekeyword = narrowtemplates.templatekeyword
 revsetpredicate = narrowtemplates.revsetpredicate



To: spectral, durin42, #hg-reviewers
Cc: mercurial-devel