arc: Add option to pin version for caches.

This is workaround for missing caches for release branches. I stil
have problems running PFQ with data collection on release branches. This
leads to some caches missing and noticable performance regression on
release branches.

The following scenarios are possible:
  * Generate caches using try_job and pin version manually for the
    release branch.
  * Run DataCollector on branch as ordinal TAST test. Once artifacts are
    ready, it automatically creates pin for the particular branch. So
    next builds will use last known good caches.

Using different versions are low risk on release branches.
  * GMS Core caches are GMS Core version dependent. We don't uprev GMS
  Core version in late stage of release branch.
  * ureadahead pack has risk to mismatch actual data however this should
  be minimal on release branches.

BUG=b:163823561
TEST=Locally, ./run_tests + monitor performance number from in-lab tests
on release branches.

Change-Id: I6cce5e4ed2271a0cab529d5326f3f3280148c20f
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2366100
Reviewed-by: Yusuke Sato <yusukes@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
Commit-Queue: Yury Khmel <khmel@google.com>
Commit-Queue: Mike Frysinger <vapier@chromium.org>
Tested-by: Yury Khmel <khmel@google.com>
Auto-Submit: Yury Khmel <khmel@google.com>
diff --git a/scripts/cros_mark_android_as_stable.py b/scripts/cros_mark_android_as_stable.py
index 714a188..3b33407 100644
--- a/scripts/cros_mark_android_as_stable.py
+++ b/scripts/cros_mark_android_as_stable.py
@@ -424,25 +424,34 @@
   }))
 
 
-def UpdateDataCollectorArtifacts(android_version,
-                                 runtime_artifacts_bucket_url):
+def FindDataCollectorArtifacts(gs_context,
+                               android_version,
+                               runtime_artifacts_bucket_url,
+                               version_reference):
   r"""Finds and includes into variables artifacts from arc.DataCollector.
 
+  This is used from UpdateDataCollectorArtifacts in order to check the
+  particular version.
+
   Args:
+    gs_context: context to execute gsutil
     android_version: The \d+ build id of Android.
     runtime_artifacts_bucket_url: root of runtime artifacts
+    build_branch: build branch. Used to determine the pinned version if exists.
+    version_reference: which version to use as a reference. Could be '${PV}' in
+                       case version of data collector artifacts matches the
+                       Android version or direct version in case of override.
 
   Returns:
-    dictionary with filled ebuild variables.
+    dictionary with filled ebuild variables. This dictionary is empty in case
+    no artificats are found.
   """
-
   variables = {}
+
   buckets = ['ureadahead_pack', 'gms_core_cache']
   archs = ['arm', 'arm64', 'x86', 'x86_64']
   build_types = ['user', 'userdebug']
 
-  version_reference = '${PV}'
-  gs_context = gs.GSContext()
   for bucket in buckets:
     for arch in archs:
       for build_type in build_types:
@@ -456,6 +465,59 @@
   return variables
 
 
+def UpdateDataCollectorArtifacts(android_version,
+                                 runtime_artifacts_bucket_url,
+                                 build_branch):
+  r"""Finds and includes into variables artifacts from arc.DataCollector.
+
+  This verifies default android version. In case artificts are not found for
+  default Android version it tries to find artifacts for pinned version. If
+  pinned version is provided, it is required artifacts exist for the pinned
+  version.
+
+  Args:
+    android_version: The \d+ build id of Android.
+    runtime_artifacts_bucket_url: root of runtime artifacts
+    build_branch: build branch. Used to determine the pinned version if exists.
+
+  Returns:
+    dictionary with filled ebuild variables.
+  """
+
+  gs_context = gs.GSContext()
+  # Check the existing version. If we find any artifacts, use them.
+  variables = FindDataCollectorArtifacts(gs_context,
+                                         android_version,
+                                         runtime_artifacts_bucket_url,
+                                         '${PV}')
+  if variables:
+    # Data artificts were found.
+    return variables
+
+  # Check pinned version for the current branch.
+  pin_path = (f'{runtime_artifacts_bucket_url}/{build_branch}_pin_version')
+  if not gs_context.Exists(pin_path):
+    # No pinned version.
+    logging.warning(
+        'No data collector artifacts were found for %s',
+        android_version)
+    return variables
+
+  pin_version = gs_context.Cat(pin_path, encoding='utf-8').rstrip()
+  logging.info('Pinned version %s overrides %s',
+               pin_version, android_version)
+  variables = FindDataCollectorArtifacts(gs_context,
+                                         pin_version,
+                                         runtime_artifacts_bucket_url,
+                                         pin_version)
+  if not variables:
+    # If pin version set it must contain data.
+    raise Exception('Pinned version %s:%s does not contain artificats' % (
+        build_branch, pin_version))
+
+  return variables
+
+
 def MarkAndroidEBuildAsStable(stable_candidate, unstable_ebuild,
                               android_package, android_version, package_dir,
                               build_branch, arc_bucket_url,
@@ -510,7 +572,7 @@
     variables[build + '_TARGET'] = '%s-%s' % (build_branch, target)
 
   variables.update(UpdateDataCollectorArtifacts(
-      android_version, runtime_artifacts_bucket_url))
+      android_version, runtime_artifacts_bucket_url, build_branch))
 
   portage_util.EBuild.MarkAsStable(
       unstable_ebuild.ebuild_path, new_ebuild_path,
diff --git a/scripts/cros_mark_android_as_stable_unittest.py b/scripts/cros_mark_android_as_stable_unittest.py
index 0bff23c..2297f13 100644
--- a/scripts/cros_mark_android_as_stable_unittest.py
+++ b/scripts/cros_mark_android_as_stable_unittest.py
@@ -259,6 +259,10 @@
                   f'{build_type}_{android_version}.tar')
           self.gs_mock.AddCmdResult(['stat', '--', path],
                                     side_effect=_RaiseGSNoSuchKey)
+    pin_path = (f'{self.runtime_artifacts_bucket_url}/'
+                f'{constants.ANDROID_PI_BUILD_BRANCH}_pin_version')
+    self.gs_mock.AddCmdResult(['stat', '--', pin_path],
+                              side_effect=_RaiseGSNoSuchKey)
 
   def makeSrcTargetUrl(self, target):
     """Helper to return the url for a target."""
@@ -565,15 +569,54 @@
                               stdout=(self.STAT_OUTPUT) % path2)
 
     variables = cros_mark_android_as_stable.UpdateDataCollectorArtifacts(
-        android_version, self.runtime_artifacts_bucket_url)
+        android_version,
+        self.runtime_artifacts_bucket_url,
+        constants.ANDROID_PI_BUILD_BRANCH)
 
-    self.assertEqual(2, len(variables))
-    self.assertIn('X86_64_USER_UREADAHEAD_PACK', variables)
     version_reference = '${PV}'
     expectation1 = (f'{self.runtime_artifacts_bucket_url}/'
                     f'ureadahead_pack_x86_64_user_{version_reference}.tar')
-    self.assertEqual(expectation1, variables['X86_64_USER_UREADAHEAD_PACK'])
-    self.assertIn('ARM_USERDEBUG_GMS_CORE_CACHE', variables)
     expectation2 = (f'{self.runtime_artifacts_bucket_url}/'
                     f'gms_core_cache_arm_userdebug_{version_reference}.tar')
-    self.assertEqual(expectation2, variables['ARM_USERDEBUG_GMS_CORE_CACHE'])
+    self.assertEqual({
+        'X86_64_USER_UREADAHEAD_PACK': expectation1,
+        'ARM_USERDEBUG_GMS_CORE_CACHE': expectation2,
+    }, variables)
+
+  def testUpdateDataCollectorArtifactsPin(self):
+    android_version = 100
+    android_pin_version = 50
+    # Mock by default runtime artifacts are not found.
+    self.setupMockRuntimeDataBuild(android_version)
+    self.setupMockRuntimeDataBuild(android_pin_version)
+
+    # Override few as existing.
+    path1 = (f'{self.runtime_artifacts_bucket_url}/ureadahead_pack_x86_64_'
+             f'user_{android_pin_version}.tar')
+    self.gs_mock.AddCmdResult(['stat', '--', path1],
+                              stdout=(self.STAT_OUTPUT) % path1)
+    path2 = (f'{self.runtime_artifacts_bucket_url}/gms_core_cache_arm_'
+             f'userdebug_{android_pin_version}.tar')
+    self.gs_mock.AddCmdResult(['stat', '--', path2],
+                              stdout=(self.STAT_OUTPUT) % path2)
+    pin_path = (f'{self.runtime_artifacts_bucket_url}/'
+                f'{constants.ANDROID_PI_BUILD_BRANCH}_pin_version')
+    self.gs_mock.AddCmdResult(['stat', '--', pin_path],
+                              stdout=(self.STAT_OUTPUT) % pin_path)
+    self.gs_mock.AddCmdResult(['cat', pin_path],
+                              stdout=str(android_pin_version))
+
+    variables = cros_mark_android_as_stable.UpdateDataCollectorArtifacts(
+        android_version,
+        self.runtime_artifacts_bucket_url,
+        constants.ANDROID_PI_BUILD_BRANCH)
+
+    version_reference = '50'
+    expectation1 = (f'{self.runtime_artifacts_bucket_url}/'
+                    f'ureadahead_pack_x86_64_user_{version_reference}.tar')
+    expectation2 = (f'{self.runtime_artifacts_bucket_url}/'
+                    f'gms_core_cache_arm_userdebug_{version_reference}.tar')
+    self.assertEqual({
+        'X86_64_USER_UREADAHEAD_PACK': expectation1,
+        'ARM_USERDEBUG_GMS_CORE_CACHE': expectation2,
+    }, variables)