purge: use fsmonitor to speed up hg purge

Summary:
`hg purge --dirs` is slow because it asks for all the directories in the repo,
which disables the fsmonitor speed improvements to dirstate.walk.

To make purge faster, separately query watchman for all the dirs in the repo
(it can answer this fairly quickly).

Reviewed By: mjpieters

Differential Revision: D7067811

fbshipit-source-id: 18e76c685bd630862d12e7e59c2817a8f45ed073

hgext/fsmonitor/__init__.py

@@ -216,6 +216,53 @@ def _watchmantofsencoding(path):
 
     return encoded
 
+def _finddirs(dirstate):
+    '''Query watchman for all directories in the working copy'''
+    state = dirstate._fsmonitorstate
+    dirstate._watchmanclient.settimeout(state.timeout + 0.1)
+    result = dirstate._watchmanclient.command('query', {
+        'fields': ['name'],
+            'expression': [
+                'allof', ['type', 'd'],
+                [
+                    'not', [
+                        'anyof', ['dirname', '.hg'],
+                        ['name', '.hg', 'wholename']
+                    ]
+                ]
+            ],
+            'sync_timeout': int(state.timeout * 1000),
+            'empty_on_fresh_instance': state.walk_on_invalidate,
+        })
+    return result['files']
+
+def wrappurge(orig, repo, match, findfiles, finddirs, includeignored):
+    # If includeignored is set, we always need to do a full rewalk.
+    if includeignored:
+        return orig(repo, match, findfiles, finddirs, includeignored)
+
+    files = []
+    dirs = []
+    usefastdirs = True
+    if finddirs:
+        try:
+            fastdirs = _finddirs(repo.dirstate)
+        except Exception:
+            repo.ui.debug('fsmonitor: fallback to core purge, '
+                          'query dirs failed')
+            usefastdirs = False
+
+    if findfiles or not usefastdirs:
+        files, dirs = orig(repo, match, findfiles,
+                           finddirs and not usefastdirs, False)
+
+    if finddirs and usefastdirs:
+        dirs = (f for f in sorted(fastdirs, reverse=True)
+                if (match(f) and not os.listdir(repo.wjoin(f)) and
+                    not repo.dirstate._dirignore(f)))
+
+    return files, dirs
+
 def overridewalk(orig, self, match, subrepos, unknown, ignored, full=True):
     '''Replacement for dirstate.walk, hooking into Watchman.
 
@@ -629,6 +676,13 @@ def extsetup(ui):
 
     extensions.wrapfunction(merge, 'update', wrapupdate)
 
+    def purgeloaded(loaded=False):
+        if not loaded:
+            return
+        purge = extensions.find('purge')
+        extensions.wrapfunction(purge, 'findthingstopurge', wrappurge)
+    extensions.afterloaded('purge', purgeloaded)
+
 def wrapsymlink(orig, source, link_name):
     ''' if we create a dangling symlink, also touch the parent dir
     to encourage fsevents notifications to work more correctly '''

hgext/purge.py

@@ -44,6 +44,35 @@ command = registrar.command(cmdtable)
 # leave the attribute unspecified.
 testedwith = 'ships-with-hg-core'
 
+def findthingstopurge(repo, match, findfiles, finddirs, includeignored):
+    """Find files and/or directories that should be purged.
+
+    Returns a pair (files, dirs), where files is an iterable of files to
+    remove, and dirs is an iterable of directories to remove.
+    """
+    if finddirs:
+        directories = []
+        match.explicitdir = match.traversedir = directories.append
+
+    status = repo.status(match=match, ignored=includeignored, unknown=True)
+
+    if findfiles:
+        files = sorted(status.unknown + status.ignored)
+    else:
+        files = []
+
+    if finddirs:
+        # Use a generator expression to lazily test for directory contents,
+        # otherwise nested directories that are being removed would be counted
+        # when in reality they'd be removed already by the time the parent
+        # directory is to be removed.
+        dirs = (f for f in sorted(directories, reverse=True)
+                if (match(f) and not os.listdir(repo.wjoin(f))))
+    else:
+        dirs = []
+
+    return files, dirs
+
 @command('purge|clean',
     [('a', 'abort-on-err', None, _('abort if an error occurs')),
     ('',  'all', None, _('purge ignored files too')),
@@ -91,6 +120,7 @@ def purge(ui, repo, *dirs, **opts):
         act = False # --print0 implies --print
     removefiles = opts.get('files')
     removedirs = opts.get('dirs')
+    removeignored = opts.get('all')
     if not removefiles and not removedirs:
         removefiles = True
         removedirs = True
@@ -108,20 +138,15 @@ def purge(ui, repo, *dirs, **opts):
             ui.write('%s%s' % (name, eol))
 
     match = scmutil.match(repo[None], dirs, opts)
-    if removedirs:
-        directories = []
-        match.explicitdir = match.traversedir = directories.append
-    status = repo.status(match=match, ignored=opts.get('all'), unknown=True)
+    files, dirs = findthingstopurge(repo, match, removefiles, removedirs,
+                                    removeignored)
 
-    if removefiles:
-        for f in sorted(status.unknown + status.ignored):
-            if act:
-                ui.note(_('removing file %s\n') % f)
-            remove(util.unlink, f)
+    for f in files:
+        if act:
+            ui.note(_('removing file %s\n') % f)
+        remove(util.unlink, f)
 
-    if removedirs:
-        for f in sorted(directories, reverse=True):
-            if match(f) and not os.listdir(repo.wjoin(f)):
-                if act:
-                    ui.note(_('removing directory %s\n') % f)
-                remove(os.rmdir, f)
+    for f in dirs:
+        if act:
+            ui.note(_('removing directory %s\n') % f)
+        remove(os.rmdir, f)