From 747cb649f7ecc681267e8b2d14c225b664546b80 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Tue, 1 Sep 2020 09:56:21 +0545
Subject: [PATCH] tree.repo: fix race condition when walking nested subrepos
 (#4505)

---
 dvc/tree/repo.py             | 13 ++++++---
 tests/unit/tree/test_repo.py | 53 +++++++++++++++++++++++++++++++++++-
 2 files changed, 61 insertions(+), 5 deletions(-)

diff --git a/dvc/tree/repo.py b/dvc/tree/repo.py
index 49c4b467b..aa1f5ad14 100644
--- a/dvc/tree/repo.py
+++ b/dvc/tree/repo.py
@@ -6,7 +6,7 @@ from contextlib import suppress
 from itertools import takewhile
 from typing import TYPE_CHECKING, Callable, Optional, Tuple, Union
 
-from funcy import wrap_with
+from funcy import lfilter, wrap_with
 from pygtrie import StringTrie
 
 from dvc.dvcfile import is_valid_filename
@@ -235,14 +235,19 @@ class RepoTree(BaseTree):  # pylint:disable=abstract-method
 
         yield repo_root, dirs, list(files)
 
+        def is_dvc_repo(d):
+            return self._is_dvc_repo(os.path.join(repo_root, d))
+
+        # remove subrepos to prevent it from being traversed
+        subrepos = set(filter(is_dvc_repo, repo_only))
         # set dir order for next recursion level - shared dirs first so that
         # next() for both generators recurses into the same shared directory
         dvc_dirs[:] = [dirname for dirname in dirs if dirname in dvc_set]
-        repo_dirs[:] = [dirname for dirname in dirs if dirname in repo_set]
+        repo_dirs[:] = lfilter(lambda d: d in (repo_set - subrepos), dirs)
 
         for dirname in dirs:
-            dir_path = os.path.join(repo_root, dirname)
-            if self._is_dvc_repo(dir_path):
+            if dirname in subrepos:
+                dir_path = os.path.join(repo_root, dirname)
                 yield from self._subrepo_walk(dir_path, dvcfiles=dvcfiles)
             elif dirname in shared:
                 yield from self._walk(repo_walk, dvc_walk, dvcfiles=dvcfiles)
diff --git a/tests/unit/tree/test_repo.py b/tests/unit/tree/test_repo.py
index 40998c106..1304fba82 100644
--- a/tests/unit/tree/test_repo.py
+++ b/tests/unit/tree/test_repo.py
@@ -182,7 +182,7 @@ def test_isdvc(tmp_dir, dvc):
 
 
 def make_subrepo(dir_, scm, config=None):
-    dir_.mkdir(parents=True)
+    dir_.mkdir(parents=True, exist_ok=True)
     with dir_.chdir():
         dir_.scm = scm
         dir_.init(dvc=True, subdir=True)
@@ -398,3 +398,54 @@ def test_get_hash_cached_granular(tmp_dir, dvc, mocker):
         "md5", "8d777f385d3dfec8815d20f7496026dc",
     )
     assert dvc_tree_spy.called
+
+
+@pytest.mark.parametrize("traverse_subrepos", [True, False])
+def test_walk_nested_subrepos(tmp_dir, dvc, scm, traverse_subrepos):
+    # generate a dvc and fs structure, with suffix based on repo's basename
+    def fs_structure(suffix):
+        return {
+            f"foo-{suffix}": f"foo-{suffix}",
+            f"dir-{suffix}": {f"bar-{suffix}": f"bar-{suffix}"},
+        }
+
+    def dvc_structure(suffix):
+        return {
+            f"lorem-{suffix}": f"lorem-{suffix}",
+            f"dvc-{suffix}": {f"ipsum-{suffix}": f"ipsum-{suffix}"},
+        }
+
+    paths = ["subrepo1", "subrepo2", "subrepo1/subrepo3"]
+    subrepos = [tmp_dir / path for path in paths]
+    for repo_dir in subrepos:
+        make_subrepo(repo_dir, scm)
+
+    extras = {".dvcignore", ".gitignore"}  # these files are always there
+    expected = {}
+    for repo_dir in subrepos + [tmp_dir]:
+        base = os.path.basename(repo_dir)
+        scm_files = fs_structure(base)
+        dvc_files = dvc_structure(base)
+        with repo_dir.chdir():
+            repo_dir.scm_gen(scm_files, commit=f"git add in {repo_dir}")
+            repo_dir.dvc_gen(dvc_files, commit=f"dvc add in {repo_dir}")
+
+        if traverse_subrepos or repo_dir == tmp_dir:
+            expected[str(repo_dir)] = set(
+                scm_files.keys() | dvc_files.keys() | extras
+            )
+            # files inside a dvc directory
+            expected[str(repo_dir / f"dvc-{base}")] = {f"ipsum-{base}"}
+            # files inside a git directory
+            expected[str(repo_dir / f"dir-{base}")] = {f"bar-{base}"}
+
+    if traverse_subrepos:
+        # update subrepos
+        expected[str(tmp_dir)].update(["subrepo1", "subrepo2"])
+        expected[str(tmp_dir / "subrepo1")].add("subrepo3")
+
+    actual = {}
+    tree = RepoTree(dvc, subrepos=traverse_subrepos)
+    for root, dirs, files in tree.walk(str(tmp_dir)):
+        actual[root] = set(dirs + files)
+    assert expected == actual