From 1e698d4aaf8d871a1e192e0aca4e380fa5221dfa Mon Sep 17 00:00:00 2001
From: Arun Kulshreshtha <kulshrax@fb.com>
Date: Wed, 18 Mar 2020 18:21:16 -0700
Subject: [PATCH] treemanifest: add ondemandfetch parameter to _gettrees

Summary: Based on comments on D20382825, we need to make sure that `_gettrees()` knows for sure whether on-demand tree fetching is in use in order to properly identify missing nodes in the response.

Reviewed By: quark-zju

Differential Revision: D20520439

fbshipit-source-id: ffa6d62dbe8b6f641b1dacebcb6f94ceae714c1b
---
 .../edenscm/hgext/treemanifest/__init__.py    | 39 ++++++++++++++-----
 1 file changed, 29 insertions(+), 10 deletions(-)

diff --git a/eden/scm/edenscm/hgext/treemanifest/__init__.py b/eden/scm/edenscm/hgext/treemanifest/__init__.py
index 89709e23f6..c12e5df46a 100644
--- a/eden/scm/edenscm/hgext/treemanifest/__init__.py
+++ b/eden/scm/edenscm/hgext/treemanifest/__init__.py
@@ -612,7 +612,17 @@ def wraprepo(repo):
 
                 start = util.timer()
                 with self.ui.timesection("getdesignatednodes"):
-                    _gettrees(self, conn.peer, "", mfnodes, [], directories, start, 1)
+                    _gettrees(
+                        self,
+                        conn.peer,
+                        "",
+                        mfnodes,
+                        [],
+                        directories,
+                        start,
+                        depth=1,
+                        ondemandfetch=True,
+                    )
 
             return True
 
@@ -2002,7 +2012,17 @@ def _prefetchonlytrees(repo, opts):
     repo.prefetchtrees(mfnodes, basemfnodes=basemfnode)
 
 
-def _gettrees(repo, remote, rootdir, mfnodes, basemfnodes, directories, start, depth):
+def _gettrees(
+    repo,
+    remote,
+    rootdir,
+    mfnodes,
+    basemfnodes,
+    directories,
+    start,
+    depth,
+    ondemandfetch=False,
+):
     if "gettreepack" not in shallowutil.peercapabilities(remote):
         raise error.Abort(_("missing gettreepack capability on remote"))
     bundle = remote.gettreepack(rootdir, mfnodes, basemfnodes, directories, depth)
@@ -2014,15 +2034,14 @@ def _gettrees(repo, remote, rootdir, mfnodes, basemfnodes, directories, start, d
         count = 0
         missingnodes = set(mfnodes)
 
-        # If we're doing on-demand tree fetching, this means that we are not trying
-        # to fetch complete trees. Consequently, the set of mfnodes passed in are not
-        # all different versions of the same root directory -- instead they correspond
-        # to individual subdirectories within a single tree, which we are explicitly
-        # not downloading in its entirety. This means we should not check the directory
-        # path when checking for missing nodes in the response.
-        ondemandfetch = repo.ui.configbool("treemanifest", "ondemandfetch")
-
         for reply in receivednodes:
+            # If we're doing on-demand tree fetching, this means that we are not
+            # trying to fetch complete trees. Consequently, the set of mfnodes
+            # passed in are not all different versions of the same root
+            # directory -- instead they correspond to individual subdirectories
+            # within a single tree, which we are explicitly not downloading in
+            # its entirety. This means we should not check the directory path
+            # when checking for missing nodes in the response.
             if ondemandfetch:
                 missingnodes.difference_update(n for d, n in reply)
             else: