Adding batch CIPD version resolution to gclient.

At the moment on repo's with many CIPD packages the number
of CLI calls to 'cipd describe' can make it very slow.

This mechanism will call cipd ensure-file-resolve with the entire
list in one call then reference items from the list before trying
individual cipd describe commands.

Bug:b/279097790
Change-Id: Ieb4bc77853357cfadb3ea6311fa0013aa7999ac6
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/tools/depot_tools/+/4551125
Reviewed-by: Gavin Mak <gavinmak@google.com>
Reviewed-by: Rachael Newitt <renewitt@google.com>
Commit-Queue: Dan Le Febvre <dlf@google.com>

gclient_scm.py

@@ -1553,6 +1553,7 @@ class CipdRoot(object):
     self._packages_by_subdir = collections.defaultdict(list)
     self._root_dir = root_dir
     self._service_url = service_url
+    self._resolved_packages = None
 
   def add_package(self, subdir, package, version):
     """Adds a package to this CIPD root.
@@ -1582,6 +1583,11 @@ class CipdRoot(object):
     """Get the list of configured packages for the given subdir."""
     return list(self._packages_by_subdir[subdir])
 
+  def resolved_packages(self):
+    if not self._resolved_packages:
+      self._resolved_packages = self.ensure_file_resolve()
+    return self._resolved_packages
+
   def clobber(self):
     """Remove the .cipd directory.
 
@@ -1644,6 +1650,52 @@ class CipdRoot(object):
         gclient_utils.CheckCallAndFilter(
             cmd, print_stdout=True, show_header=True)
 
+  @contextlib.contextmanager
+  def _create_ensure_file_for_resolve(self):
+    try:
+      contents = '$ResolvedVersions %s\n' % os.devnull
+      for subdir, packages in sorted(self._packages_by_subdir.items()):
+        contents += '@Subdir %s\n' % subdir
+        for package in sorted(packages, key=lambda p: p.name):
+          contents += '%s %s\n' % (package.name, package.version)
+        contents += '\n'
+      ensure_file = None
+      with tempfile.NamedTemporaryFile(suffix='.ensure',
+                                       delete=False,
+                                       mode='wb') as ensure_file:
+        ensure_file.write(contents.encode('utf-8', 'replace'))
+      yield ensure_file.name
+    finally:
+      if ensure_file is not None and os.path.exists(ensure_file.name):
+        os.remove(ensure_file.name)
+
+  def _create_resolved_file(self):
+    return tempfile.NamedTemporaryFile(suffix='.resolved',
+                                       delete=False,
+                                       mode='wb')
+
+  def ensure_file_resolve(self):
+    """Run `cipd ensure-file-resolve`."""
+    with self._mutator_lock:
+      with self._create_resolved_file() as output_file:
+        with self._create_ensure_file_for_resolve() as ensure_file:
+          cmd = [
+              'cipd',
+              'ensure-file-resolve',
+              '-log-level',
+              'error',
+              '-ensure-file',
+              ensure_file,
+              '-json-output',
+              output_file.name,
+          ]
+          gclient_utils.CheckCallAndFilter(cmd,
+                                           print_stdout=False,
+                                           show_header=False)
+          with open(output_file.name) as f:
+            output_json = json.load(f)
+            return output_json.get('result', {})
+
   def run(self, command):
     if command == 'update':
       self.ensure()
@@ -1716,6 +1768,22 @@ class CipdWrapper(SCMWrapper):
     """Grab the instance ID."""
     try:
       tmpdir = tempfile.mkdtemp()
+      # Attempt to get instance_id from the root resolved cache.
+      # Resolved cache will not match on any CIPD packages with
+      # variables such as ${platform}, they will fall back to
+      # the slower method below.
+      resolved = self._root.resolved_packages()
+      if resolved:
+        # CIPD uses POSIX separators across all platforms, so
+        # replace any Windows separators.
+        path_split = self.relpath.replace(os.sep, "/").split(":")
+        if len(path_split) > 1:
+          src_path, package = path_split
+          if src_path in resolved:
+            for resolved_package in resolved[src_path]:
+              if package == resolved_package.get('pin', {}).get('package'):
+                return resolved_package.get('pin', {}).get('instance_id')
+
       describe_json_path = os.path.join(tmpdir, 'describe.json')
       cmd = [
           'cipd', 'describe',

testing_support/fake_cipd.py

@@ -19,6 +19,7 @@ PLATFORM_VAR = 'platform'
 CIPD_SUBDIR_RE = '@Subdir (.*)'
 CIPD_DESCRIBE = 'describe'
 CIPD_ENSURE = 'ensure'
+CIPD_ENSURE_FILE_RESOLVE = 'ensure-file-resolve'
 CIPD_EXPAND_PKG = 'expand-package-name'
 CIPD_EXPORT = 'export'
 
@@ -81,6 +82,8 @@ def parse_cipd(root, contents):
     if match:
       print('match')
       current_subdir = os.path.join(root, *match.group(1).split('/'))
+      if not root:
+        current_subdir = match.group(1)
     elif line and current_subdir:
       print('no match')
       tree.setdefault(current_subdir, []).append(line)
@@ -102,6 +105,30 @@ def expand_package_name_cmd(package_name):
   return package_name
 
 
+def ensure_file_resolve():
+  resolved = {"result": {}}
+  parser = argparse.ArgumentParser()
+  parser.add_argument('-ensure-file', required=True)
+  parser.add_argument('-json-output')
+  args, _ = parser.parse_known_args()
+  with io.open(args.ensure_file, 'r', encoding='utf-8') as f:
+    new_content = parse_cipd("", f.readlines())
+    for path, packages in new_content.items():
+      resolved_packages = []
+      for package in packages:
+        package_name = expand_package_name_cmd(package.split(" ")[0])
+        resolved_packages.append({
+            "package": package_name,
+            "pin": {
+                "package": package_name,
+                "instance_id": package_name + "-fake-resolved-id",
+            }
+        })
+      resolved["result"][path] = resolved_packages
+    with io.open(args.json_output, 'w', encoding='utf-8') as f:
+      f.write(json.dumps(resolved, indent=4))
+
+
 def describe_cmd(package_name):
   parser = argparse.ArgumentParser()
   parser.add_argument('-json-output')
@@ -132,7 +159,10 @@ def describe_cmd(package_name):
 
 def main():
   cmd = sys.argv[1]
-  assert cmd in [CIPD_DESCRIBE, CIPD_ENSURE, CIPD_EXPAND_PKG, CIPD_EXPORT]
+  assert cmd in [
+      CIPD_DESCRIBE, CIPD_ENSURE, CIPD_ENSURE_FILE_RESOLVE, CIPD_EXPAND_PKG,
+      CIPD_EXPORT
+  ]
   # Handle cipd expand-package-name
   if cmd == CIPD_EXPAND_PKG:
     # Expecting argument after cmd
@@ -144,6 +174,8 @@ def main():
     # Expecting argument after cmd
     assert len(sys.argv) >= 3
     return describe_cmd(sys.argv[2])
+  if cmd == CIPD_ENSURE_FILE_RESOLVE:
+    return ensure_file_resolve()
 
   parser = argparse.ArgumentParser()
   parser.add_argument('-ensure-file')

testing_support/fake_repos.py

@@ -747,6 +747,10 @@ hooks = [{
                 'package': 'package0/${{platform}}',
                 'version': 'package0/${{platform}}-fake-tag:1.0',
               },
+              {
+                'package': 'package1/another',
+                'version': 'package1/another-fake-tag:1.0',
+              },
             ],
             'dep_type': 'cipd',
           },

tests/gclient_cipd_smoketest.py

@@ -140,7 +140,8 @@ class GClientSmokeCipd(gclient_smoketest_base.GClientSmokeBase):
     results = self.gclient(['revinfo'])
     out = ('src: %(base)srepo_18\n'
            'src/cipd_dep:package0: %(instance_url1)s\n'
-           'src/cipd_dep:package0/${platform}: %(instance_url2)s\n' % {
+           'src/cipd_dep:package0/${platform}: %(instance_url2)s\n'
+           'src/cipd_dep:package1/another: %(instance_url3)s\n' % {
                'base':
                self.git_base,
                'instance_url1':
@@ -148,6 +149,9 @@ class GClientSmokeCipd(gclient_smoketest_base.GClientSmokeBase):
                'instance_url2':
                CHROME_INFRA_URL +
                '/package0/${platform}@package0/${platform}-fake-tag:1.0',
+               'instance_url3':
+               CHROME_INFRA_URL +
+               '/package1/another@package1/another-fake-tag:1.0',
            })
     self.check((out, '', 0), results)
 
@@ -155,19 +159,27 @@ class GClientSmokeCipd(gclient_smoketest_base.GClientSmokeBase):
     self.gclient(['config', self.git_base + 'repo_18', '--name', 'src'])
     self.gclient(['sync'])
     results = self.gclient(['revinfo', '--actual'])
-    out = ('src: %(base)srepo_18@%(hash1)s\n'
-           'src/cipd_dep:package0: %(instance_url1)s\n'
-           'src/cipd_dep:package0/${platform}: %(instance_url2)s\n' % {
-               'base':
-               self.git_base,
-               'hash1':
-               self.githash('repo_18', 1),
-               'instance_url1':
-               CHROME_INFRA_URL + '/p/package0/+/package0-fake-instance-id',
-               'instance_url2':
-               CHROME_INFRA_URL + '/p/package0/platform-expanded-test-only' +
-               '/+/package0/${platform}-fake-instance-id',
-           })
+    out = (
+        'src: %(base)srepo_18@%(hash1)s\n'
+        'src/cipd_dep:package0: %(instance_url1)s\n'
+        'src/cipd_dep:package0/${platform}: %(instance_url2)s\n'
+        'src/cipd_dep:package1/another: %(instance_url3)s\n' % {
+            'base':
+            self.git_base,
+            'hash1':
+            self.githash('repo_18', 1),
+            'instance_url1':
+            # The below 'fake-*' and 'platform-expanded-*' ID's are from:
+            # ../testing_support/fake_cipd.py. 'fake-resolved' represents
+            # the package being found in the batch resolution mechanism.
+            CHROME_INFRA_URL + '/p/package0/+/package0-fake-resolved-id',
+            'instance_url2':
+            CHROME_INFRA_URL + '/p/package0/platform-expanded-test-only' +
+            '/+/package0/${platform}-fake-instance-id',
+            'instance_url3':
+            CHROME_INFRA_URL + '/p/package1/another' +
+            '/+/package1/another-fake-resolved-id',
+        })
     self.check((out, '', 0), results)
 
 

tests/gclient_scm_test.py

@@ -954,6 +954,7 @@ class CipdWrapperTestCase(unittest.TestCase):
     mock.patch('tempfile.mkdtemp', lambda: self._workdir).start()
     mock.patch('gclient_scm.CipdRoot.add_package').start()
     mock.patch('gclient_scm.CipdRoot.clobber').start()
+    mock.patch('gclient_scm.CipdRoot.ensure_file_resolve').start()
     mock.patch('gclient_scm.CipdRoot.ensure').start()
     self.addCleanup(mock.patch.stopall)
     self.addCleanup(gclient_utils.rmtree, self._cipd_root_dir)