Add CIPD support for bot assets

BUG=skia:
GOLD_TRYBOT_URL= https://gold.skia.org/search?issue=2085473002

Review-Url: https://codereview.chromium.org/2085473002

8 files changed, 320 insertions, 53 deletions

diff --git a/infra/bots/assets/asset_utils.py b/infra/bots/assets/asset_utils.py
index a13d2290a3..27d6b85de7 100644
--- a/infra/bots/assets/asset_utils.py
+++ b/infra/bots/assets/asset_utils.py
@@ -10,36 +10,147 @@
 
 
 import argparse
+import json
 import os
 import shlex
 import shutil
 import subprocess
 import sys
 
-SKIA_DIR = os.path.abspath(os.path.realpath(os.path.join(
-    os.path.dirname(os.path.abspath(__file__)),
-    os.pardir, os.pardir, os.pardir)))
-INFRA_BOTS_DIR = os.path.join(SKIA_DIR, 'infra', 'bots')
+INFRA_BOTS_DIR = os.path.abspath(os.path.realpath(os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), os.pardir)))
 sys.path.insert(0, INFRA_BOTS_DIR)
 import utils
 import zip_utils
 
 
 ASSETS_DIR = os.path.join(INFRA_BOTS_DIR, 'assets')
+SKIA_DIR = os.path.abspath(os.path.join(INFRA_BOTS_DIR, os.pardir, os.pardir))
+
+CIPD_PACKAGE_NAME_TMPL = 'skia/bots/%s'
+DEFAULT_CIPD_SERVICE_URL = 'https://chrome-infra-packages.appspot.com'
+
 DEFAULT_GS_BUCKET = 'skia-buildbots'
 GS_SUBDIR_TMPL = 'gs://%s/assets/%s'
 GS_PATH_TMPL = '%s/%s.zip'
+
+TAG_PROJECT_SKIA = 'project:skia'
+TAG_VERSION_PREFIX = 'version:'
+TAG_VERSION_TMPL = '%s%%s' % TAG_VERSION_PREFIX
+
 VERSION_FILENAME = 'VERSION'
 ZIP_BLACKLIST = ['.git', '.svn', '*.pyc', '.DS_STORE']
 
 
-class _GSWrapper(object):
+class CIPDStore(object):
+  """Wrapper object for CIPD."""
+  def __init__(self, cipd_url=DEFAULT_CIPD_SERVICE_URL):
+    cipd = 'cipd'
+    platform = 'linux64'
+    if sys.platform == 'darwin':
+      platform = 'mac64'
+    elif sys.platform == 'win32':
+      platform = 'win64'
+      cipd = 'cipd.exe'
+    self._cipd_path = os.path.join(INFRA_BOTS_DIR, 'tools', 'luci-go', platform)
+    self._cipd = os.path.join(self._cipd_path, cipd)
+    self._cipd_url = cipd_url
+    self._check_setup()
+
+  def _check_setup(self):
+    """Verify that we have the CIPD binary and that we're authenticated."""
+    try:
+      subprocess.check_call([self._cipd, 'auth-info'])
+    except OSError:
+      cipd_sha1_path = os.path.join(self._cipd_path, 'cipd.sha1')
+      raise Exception('CIPD binary not found in %s. You may need to run:\n\n'
+                      '$ download_from_google_storage -s %s'
+                      ' --bucket chromium-luci' % (self._cipd, cipd_sha1_path))
+    except subprocess.CalledProcessError:
+      raise Exception('CIPD not authenticated. You may need to run:\n\n'
+                      '$ %s auth-login' % self._cipd)
+
+  def _run(self, cmd):
+    """Run the given command."""
+    subprocess.check_call(
+        [self._cipd]
+        + cmd
+        + ['--service-url', self._cipd_url]
+    )
+
+  def _json_output(self, cmd):
+    """Run the given command, return the JSON output."""
+    with utils.tmp_dir():
+      json_output = os.path.join(os.getcwd(), 'output.json')
+      self._run(cmd + ['--json-output', json_output])
+      with open(json_output) as f:
+        parsed = json.load(f)
+    return parsed.get('result', [])
+
+  def _search(self, pkg_name):
+    res = self._json_output(['search', pkg_name, '--tag', TAG_PROJECT_SKIA])
+    return [r['instance_id'] for r in res]
+
+  def _describe(self, pkg_name, instance_id):
+    """Obtain details about the given package and instance ID."""
+    return self._json_output(['describe', pkg_name, '--version', instance_id])
+
+  def get_available_versions(self, name):
+    """List available versions of the asset."""
+    pkg_name = CIPD_PACKAGE_NAME_TMPL % name
+    versions = []
+    for instance_id in self._search(pkg_name):
+      details = self._describe(pkg_name, instance_id)
+      for tag in details.get('tags'):
+        tag_name = tag.get('tag', '')
+        if tag_name.startswith(TAG_VERSION_PREFIX):
+          trimmed = tag_name[len(TAG_VERSION_PREFIX):]
+          try:
+            versions.append(int(trimmed))
+          except ValueError:
+            raise ValueError('Found package instance with invalid version '
+                             'tag: %s' % tag_name)
+    versions.sort()
+    return versions
+
+  def upload(self, name, version, target_dir):
+    """Create a CIPD package."""
+    self._run([
+        'create',
+        '--name', CIPD_PACKAGE_NAME_TMPL % name,
+        '--in', target_dir,
+        '--tag', TAG_PROJECT_SKIA,
+        '--tag', TAG_VERSION_TMPL % version,
+    ])
+
+  def download(self, name, version, target_dir):
+    """Download a CIPD package."""
+    pkg_name = CIPD_PACKAGE_NAME_TMPL % name
+    version_tag = TAG_VERSION_TMPL % version
+    target_dir = os.path.abspath(target_dir)
+    with utils.tmp_dir():
+      infile = os.path.join(os.getcwd(), 'input')
+      with open(infile, 'w') as f:
+        f.write('%s %s' % (pkg_name, version_tag))
+      self._run([
+          'ensure',
+          '--root', target_dir,
+          '--list', infile,
+      ])
+
+  def delete_contents(self, name):
+    """Delete data for the given asset."""
+    self._run(['pkg-delete', CIPD_PACKAGE_NAME_TMPL % name])
+
+
+class GSStore(object):
   """Wrapper object for interacting with Google Storage."""
-  def __init__(self, gsutil):
+  def __init__(self, gsutil=None, bucket=DEFAULT_GS_BUCKET):
     gsutil = os.path.abspath(gsutil) if gsutil else 'gsutil'
     self._gsutil = [gsutil]
     if gsutil.endswith('.py'):
       self._gsutil = ['python', gsutil]
+    self._gs_bucket = bucket
 
   def copy(self, src, dst):
     """Copy src to dst."""
@@ -53,6 +164,68 @@ class _GSWrapper(object):
       # If the prefix does not exist, we'll get an error, which is okay.
       return []
 
+  def get_available_versions(self, name):
+    """Return the existing version numbers for the asset."""
+    files = self.list(GS_SUBDIR_TMPL % (self._gs_bucket, name))
+    bnames = [os.path.basename(f) for f in files]
+    suffix = '.zip'
+    versions = [int(f[:-len(suffix)]) for f in bnames if f.endswith(suffix)]
+    versions.sort()
+    return versions
+
+  def upload(self, name, version, target_dir):
+    """Upload to GS."""
+    target_dir = os.path.abspath(target_dir)
+    with utils.tmp_dir():
+      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
+      zip_utils.zip(target_dir, zip_file, blacklist=ZIP_BLACKLIST)
+      gs_path = GS_PATH_TMPL % (GS_SUBDIR_TMPL % (self._gs_bucket, name),
+                                str(version))
+      self.copy(zip_file, gs_path)
+
+  def download(self, name, version, target_dir):
+    """Download from GS."""
+    gs_path = GS_PATH_TMPL % (GS_SUBDIR_TMPL % (self._gs_bucket, name),
+                              str(version))
+    target_dir = os.path.abspath(target_dir)
+    with utils.tmp_dir():
+      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
+      self.copy(gs_path, zip_file)
+      zip_utils.unzip(zip_file, target_dir)
+
+  def delete_contents(self, name):
+    """Delete data for the given asset."""
+    gs_path = GS_SUBDIR_TMPL % (self._gs_bucket, name)
+    attempt_delete = True
+    try:
+      subprocess.check_call(['gsutil', 'ls', gs_path])
+    except subprocess.CalledProcessError:
+      attempt_delete = False
+    if attempt_delete:
+      subprocess.check_call(['gsutil', 'rm', '-rf', gs_path])
+
+
+class MultiStore(object):
+  """Wrapper object which uses CIPD as the primary store and GS for backup."""
+  def __init__(self, cipd_url=DEFAULT_CIPD_SERVICE_URL,
+               gsutil=None, gs_bucket=DEFAULT_GS_BUCKET):
+    self._cipd = CIPDStore(cipd_url=cipd_url)
+    self._gs = GSStore(gsutil=gsutil, bucket=gs_bucket)
+
+  def get_available_versions(self, name):
+    return self._cipd.get_available_versions(name)
+
+  def upload(self, name, version, target_dir):
+    self._cipd.upload(name, version, target_dir)
+    self._gs.upload(name, version, target_dir)
+
+  def download(self, name, version, target_dir):
+    self._cipd.download(name, version, target_dir)
+
+  def delete_contents(self, name):
+    self._cipd.delete_contents(name)
+    self._gs.delete_contents(name)
+
 
 def _prompt(prompt):
   """Prompt for input, return result."""
@@ -60,9 +233,8 @@ def _prompt(prompt):
 
 
 class Asset(object):
-  def __init__(self, name, gs_bucket=DEFAULT_GS_BUCKET, gsutil=None):
-    self._gs = _GSWrapper(gsutil)
-    self._gs_subdir = GS_SUBDIR_TMPL % (gs_bucket, name)
+  def __init__(self, name, store):
+    self._store = store
     self._name = name
     self._dir = os.path.join(ASSETS_DIR, self._name)
 
@@ -80,12 +252,7 @@ class Asset(object):
 
   def get_available_versions(self):
     """Return the existing version numbers for this asset."""
-    files = self._gs.list(self._gs_subdir)
-    bnames = [os.path.basename(f) for f in files]
-    suffix = '.zip'
-    versions = [int(f[:-len(suffix)]) for f in bnames if f.endswith(suffix)]
-    versions.sort()
-    return versions
+    return self._store.get_available_versions(self._name)
 
   def get_next_version(self):
     """Find the next available version number for the asset."""
@@ -96,12 +263,7 @@ class Asset(object):
 
   def download_version(self, version, target_dir):
     """Download the specified version of the asset."""
-    gs_path = GS_PATH_TMPL % (self._gs_subdir, str(version))
-    target_dir = os.path.abspath(target_dir)
-    with utils.tmp_dir():
-      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
-      self._gs.copy(gs_path, zip_file)
-      zip_utils.unzip(zip_file, target_dir)
+    self._store.download(self._name, version, target_dir)
 
   def download_current_version(self, target_dir):
     """Download the version of the asset specified in its version file."""
@@ -111,12 +273,7 @@ class Asset(object):
   def upload_new_version(self, target_dir, commit=False):
     """Upload a new version and update the version file for the asset."""
     version = self.get_next_version()
-    target_dir = os.path.abspath(target_dir)
-    with utils.tmp_dir():
-      zip_file = os.path.join(os.getcwd(), '%d.zip' % version)
-      zip_utils.zip(target_dir, zip_file, blacklist=ZIP_BLACKLIST)
-      gs_path = GS_PATH_TMPL % (self._gs_subdir, str(version))
-      self._gs.copy(zip_file, gs_path)
+    self._store.upload(self._name, version, target_dir)
 
     def _write_version():
       with open(self.version_file, 'w') as f:
@@ -134,9 +291,9 @@ class Asset(object):
         _write_version()
 
   @classmethod
-  def add(cls, name, gs_bucket=DEFAULT_GS_BUCKET, gsutil=None):
+  def add(cls, name, store):
     """Add an asset."""
-    asset = cls(name, gs_bucket=gs_bucket, gsutil=gsutil)
+    asset = cls(name, store)
     if os.path.isdir(asset._dir):
       raise Exception('Asset %s already exists!' % asset._name)
 
@@ -159,16 +316,17 @@ class Asset(object):
     print 'Successfully created asset %s.' % asset._name
     return asset
 
-  def remove(self):
+  def remove(self, remove_in_store=False):
     """Remove this asset."""
     # Ensure that the asset exists.
     if not os.path.isdir(self._dir):
       raise Exception('Asset %s does not exist!' % self._name)
 
+    # Cleanup the store.
+    if remove_in_store:
+      self._store.delete_contents(self._name)
+
     # Remove the asset.
     subprocess.check_call([utils.GIT, 'rm', '-rf', self._dir])
     if os.path.isdir(self._dir):
       shutil.rmtree(self._dir)
-
-    # We *could* remove all uploaded versions of the asset in Google Storage but
-    # we choose not to be that destructive.
diff --git a/infra/bots/assets/asset_utils_test.py b/infra/bots/assets/asset_utils_test.py
index edfc271f9e..a76042fe77 100644
--- a/infra/bots/assets/asset_utils_test.py
+++ b/infra/bots/assets/asset_utils_test.py
@@ -20,13 +20,13 @@ import uuid
 
 
 FILE_DIR = os.path.dirname(os.path.abspath(__file__))
-INFRA_BOTS_DIR = os.path.realpath(os.path.join(
-    FILE_DIR, os.pardir, 'infra', 'bots'))
+INFRA_BOTS_DIR = os.path.realpath(os.path.join(FILE_DIR, os.pardir))
 sys.path.insert(0, INFRA_BOTS_DIR)
 import test_utils
 import utils
 
 
+CIPD_DEV_SERVICE_URL = 'https://chrome-infra-packages-dev.appspot.com'
 GS_BUCKET = 'skia-infra-testdata'
 
 
@@ -48,16 +48,131 @@ def _write_stuff(target_dir):
   fw.write(os.path.join('subdir', 'd.txt'), 0640)
 
 
-class AssetUtilsTest(unittest.TestCase):
+class _LocalStore(object):
+  """Local store used for testing."""
+  def __init__(self):
+    self.dir = tempfile.mkdtemp()
+
+  def get_available_versions(self, name):
+    target = os.path.join(self.dir, name)
+    if not os.path.isdir(target):
+      return []
+    contents = os.listdir(os.path.join(self.dir, name))
+    return sorted([int(d) for d in contents])
+
+  def upload(self, name, version, target_dir):
+    shutil.copytree(target_dir, os.path.join(self.dir, name, str(version)))
+
+  def download(self, name, version, target_dir):
+    shutil.copytree(os.path.join(self.dir, name, str(version)), target_dir)
+
+  def delete_contents(self, name):
+    try:
+      shutil.rmtree(self.dir)
+    except OSError:
+      if os.path.exists(self.dir):
+        raise
+
+
+class StoreTest(unittest.TestCase):
+  """Superclass used for testing one of the stores."""
+  def setUp(self):
+    self.asset_name = str(uuid.uuid4())
+
+  def tearDown(self):
+    pass
+
+  def _test_upload_download(self, store):
+    with utils.tmp_dir():
+      # Create input files and directories.
+      input_dir = os.path.join(os.getcwd(), 'input')
+      _write_stuff(input_dir)
+
+      # Upload a version, download it again.
+      store.upload(self.asset_name, 0, input_dir)
+      output_dir = os.path.join(os.getcwd(), 'output')
+      store.download(self.asset_name, 0, output_dir)
+
+      # Compare.
+      test_utils.compare_trees(self, input_dir, output_dir)
+
+  def _test_versions(self, store):
+    with utils.tmp_dir():
+      # Create input files and directories.
+      input_dir = os.path.join(os.getcwd(), 'input')
+      _write_stuff(input_dir)
+      self.assertEqual(store.get_available_versions(self.asset_name), [])
+      store.upload(self.asset_name, 0, input_dir)
+      self.assertEqual(store.get_available_versions(self.asset_name), [0])
+      store.upload(self.asset_name, 1, input_dir)
+      self.assertEqual(store.get_available_versions(self.asset_name), [0, 1])
+      store.delete_contents(self.asset_name)
+      self.assertEqual(store.get_available_versions(self.asset_name), [])
+
+
+class LocalStoreTest(StoreTest):
+  """Test the local store."""
+  def setUp(self):
+    super(LocalStoreTest, self).setUp()
+    self._store = _LocalStore()
+
+  def tearDown(self):
+    self._store.delete_contents(self.asset_name)
+    super(LocalStoreTest, self).tearDown()
+
+  def test_upload_download(self):
+    self._test_upload_download(self._store)
+
+  def test_versions(self):
+    self._test_versions(self._store)
+
+
+class CIPDStoreTest(StoreTest):
+  """Test the CIPD store."""
+  def setUp(self):
+    super(CIPDStoreTest, self).setUp()
+    self._store = asset_utils.CIPDStore(cipd_url=CIPD_DEV_SERVICE_URL)
+
+  def tearDown(self):
+    self._store.delete_contents(self.asset_name)
+    super(CIPDStoreTest, self).tearDown()
+
+  def test_upload_download(self):
+    self._test_upload_download(self._store)
+
+  def test_versions(self):
+    self._test_versions(self._store)
+
+
+class GSStoreTest(StoreTest):
+  """Test the GS store."""
+  def setUp(self):
+    super(GSStoreTest, self).setUp()
+    self._store = asset_utils.GSStore(gsutil=None, bucket=GS_BUCKET)
+
+  def tearDown(self):
+    self._store.delete_contents(self.asset_name)
+    super(GSStoreTest, self).tearDown()
+
+  def test_upload_download(self):
+    self._test_upload_download(self._store)
+
+  def test_versions(self):
+    self._test_versions(self._store)
+
+
+class AssetTest(unittest.TestCase):
+  """Test Asset operations using a local store."""
   def setUp(self):
     self.asset_name = str(uuid.uuid4())
     self.old_prompt = asset_utils._prompt
     asset_utils._prompt = _fake_prompt('y')
-    self.a = asset_utils.Asset.add(self.asset_name, gs_bucket=GS_BUCKET)
+    self._store = _LocalStore()
+    self.a = asset_utils.Asset.add(self.asset_name, self._store)
 
   def tearDown(self):
     if self.a:
-      self.a.remove()
+      self.a.remove(remove_in_store=True)
     asset_utils._prompt = self.old_prompt
 
     gs_path = 'gs://%s/assets/%s' % (GS_BUCKET, self.asset_name)
@@ -72,7 +187,7 @@ class AssetUtilsTest(unittest.TestCase):
   def test_add_remove(self):
     # Ensure that we can't create an asset twice.
     with self.assertRaises(Exception):
-      asset_utils.Asset.add(self.asset_name, gs_bucket=GS_BUCKET)
+      asset_utils.Asset.add(self.asset_name, self._store)
 
     # Ensure that the asset dir exists.
     asset_dir = os.path.join(FILE_DIR, self.asset_name)
diff --git a/infra/bots/assets/assets.py b/infra/bots/assets/assets.py
index 538a41b585..b398988338 100755
--- a/infra/bots/assets/assets.py
+++ b/infra/bots/assets/assets.py
@@ -25,23 +25,23 @@ import utils
 
 def add(args):
   """Add a new asset."""
-  asset_utils.Asset.add(args.asset_name)
+  asset_utils.Asset.add(args.asset_name, asset_utils.MultiStore())
 
 
 def remove(args):
   """Remove an asset."""
-  asset_utils.Asset(args.asset_name).remove()
+  asset_utils.Asset(args.asset_name, asset_utils.MultiStore()).remove()
 
 
 def download(args):
   """Download the current version of an asset."""
-  asset = asset_utils.Asset(args.asset_name, gsutil=args.gsutil)
+  asset = asset_utils.Asset(args.asset_name, asset_utils.MultiStore())
   asset.download_current_version(args.target_dir)
 
 
 def upload(args):
   """Upload a new version of the asset."""
-  asset = asset_utils.Asset(args.asset_name, gsutil=args.gsutil)
+  asset = asset_utils.Asset(args.asset_name, asset_utils.MultiStore())
   asset.upload_new_version(args.target_dir, commit=args.commit)
 
 
diff --git a/infra/bots/download_asset.isolate b/infra/bots/download_asset.isolate
deleted file mode 100644
index 0c4d8738bc..0000000000
--- a/infra/bots/download_asset.isolate
+++ /dev/null
@@ -1,10 +0,0 @@
-{
-  'includes': [
-    'infrabots.isolate',
-  ],
-  'variables': {
-    'command': [
-      'python', 'assets/<(ASSET)/download.py', '-t', '${ISOLATED_OUTDIR}', '--gsutil', '<(GSUTIL)',
-    ],
-  },
-}
diff --git a/infra/bots/tools/luci-go/linux64/cipd.sha1 b/infra/bots/tools/luci-go/linux64/cipd.sha1
new file mode 100644
index 0000000000..b770eddb94
--- /dev/null
+++ b/infra/bots/tools/luci-go/linux64/cipd.sha1
@@ -0,0 +1 @@
+ebb43b0bf38a3ab7bbc2eaff38ec386e60fc7d99
diff --git a/infra/bots/tools/luci-go/mac64/cipd.sha1 b/infra/bots/tools/luci-go/mac64/cipd.sha1
new file mode 100644
index 0000000000..8381b7e7bb
--- /dev/null
+++ b/infra/bots/tools/luci-go/mac64/cipd.sha1
@@ -0,0 +1 @@
+2097f9871e58d4d2d9903d0e5e7de5eac96744af
diff --git a/infra/bots/tools/luci-go/win64/cipd.exe.sha1 b/infra/bots/tools/luci-go/win64/cipd.exe.sha1
new file mode 100644
index 0000000000..8e2d52a9ec
--- /dev/null
+++ b/infra/bots/tools/luci-go/win64/cipd.exe.sha1
@@ -0,0 +1 @@
+6ed0882aa8ba415aec5ff69a7cfdeaeaf60be9ed
diff --git a/infra/bots/utils.py b/infra/bots/utils.py
index a60ebe97b8..12b2ad4776 100644
--- a/infra/bots/utils.py
+++ b/infra/bots/utils.py
@@ -19,6 +19,7 @@ import uuid
 
 GCLIENT = 'gclient.bat' if sys.platform == 'win32' else 'gclient'
 GIT = 'git.bat' if sys.platform == 'win32' else 'git'
+WHICH = 'where' if sys.platform == 'win32' else 'which'
 
 
 class print_timings(object):