gsutil: Parallel-safe, specify target, add clean.

gsutil.py

 #!/usr/bin/env python
 # Copyright 2014 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Run a pinned gsutil."""
 
 
 import argparse
 import base64
 import hashlib
 import json
 import os
 import shutil
 import subprocess
 import sys
+import tempfile
+import time
 import urllib2
 import zipfile
 
 
 GSUTIL_URL = 'https://storage.googleapis.com/pub/'
 API_URL = 'https://www.googleapis.com/storage/v1/b/pub/o/'
 
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
 DEFAULT_BIN_DIR = os.path.join(THIS_DIR, 'external_bin', 'gsutil')
 DEFAULT_FALLBACK_GSUTIL = os.path.join(
     THIS_DIR, 'third_party', 'gsutil', 'gsutil')
 
-
 class InvalidGsutilError(Exception):
   pass
 
 
 def download_gsutil(version, target_dir):
   """Downloads gsutil into the target_dir."""
   filename = 'gsutil_%s.zip' % version
   target_filename = os.path.join(target_dir, filename)
 
   # Check if the target exists already.
@@ skipping 26 matching lines @@
       f.write(buf)
   return target_filename
 
 
 def check_gsutil(gsutil_bin):
   """Run gsutil version and make sure it runs."""
   return subprocess.call(
       [sys.executable, gsutil_bin, 'version'],
       stdout=subprocess.PIPE, stderr=subprocess.STDOUT) == 0
 
-def ensure_gsutil(version, target):
+def ensure_gsutil(version, target, clean):
   bin_dir = os.path.join(target, 'gsutil_%s' % version)
   gsutil_bin = os.path.join(bin_dir, 'gsutil', 'gsutil')
-  if os.path.isfile(gsutil_bin) and check_gsutil(gsutil_bin):
+  if not clean and os.path.isfile(gsutil_bin) and check_gsutil(gsutil_bin):
     # Everything is awesome! we're all done here.
     return gsutil_bin
 
-  if os.path.isdir(bin_dir):
-    # Clean up if we're redownloading a corrupted gsutil.
-    shutil.rmtree(bin_dir)
-  cache_dir = os.path.join(target, '.cache_dir')
-  if not os.path.isdir(cache_dir):
-    os.makedirs(cache_dir)
-  target_zip_filename = download_gsutil(version, cache_dir)
-  with zipfile.ZipFile(target_zip_filename, 'r') as target_zip:
-    target_zip.extractall(bin_dir)
+  # Clean up if we're redownloading a corrupted gsutil.
+  suffix = '%d_%d' % (time.time(), os.getpid())
+  try:
+    cleanup_path = '%s.%s' % (bin_dir, suffix)
+    shutil.move(bin_dir, cleanup_path)
+    shutil.rmtree(cleanup_path)

[David James, 2015/09/17 20:44:31]

did you intend to catch errors from rmtree?

[dnj (Google), 2015/09/17 21:16:39]

Hmm, probably shouldn't.

+  except IOError:

[David James, 2015/09/17 20:44:31]

These can return OSError too I believe.
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python2.7/shutil.py", line 239, in rmtree
    onerror(os.listdir, path, sys.exc_info())
  File "/usr/lib/python2.7/shutil.py", line 237, in rmtree
    names = os.listdir(path)
OSError: [Errno 13] Permission denied: 'f/f'

Consider the case when multiple users (who have different permissions) are all
writing to the same dir in parallel.

[dnj (Google), 2015/09/17 21:16:39]

Multi-user, gross. I don't think that's an intended use case.

Good point though; I'll go ahead and move the rmtree out of the try/except.

+    # Directory was not present.
+    pass
 
-  # Final check that the gsutil bin is okay.  This should never fail.
-  if not check_gsutil(gsutil_bin):
-    raise InvalidGsutilError()
+  download_dir = tempfile.mkdtemp(prefix='gsutil_py')
+  try:
+    target_zip_filename = download_gsutil(version, download_dir)
+    with zipfile.ZipFile(target_zip_filename, 'r') as target_zip:
+      target_zip.extractall(download_dir)
 
+    try:
+      shutil.move(download_dir, bin_dir)
+    except IOError:

[David James, 2015/09/17 20:44:31]

shutil.Error:
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python2.7/shutil.py", line 292, in move
    raise Error, "Destination path '%s' already exists" % real_dst
shutil.Error: Destination path 'e/f' already exists


OSError as well?
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/lib/python2.7/shutil.py", line 300, in move
    rmtree(src)
  File "/usr/lib/python2.7/shutil.py", line 256, in rmtree
    onerror(os.rmdir, path, sys.exc_info())
  File "/usr/lib/python2.7/shutil.py", line 254, in rmtree
    os.rmdir(path)
OSError: [Errno 13] Permission denied: 'g/f'


You can avoid shutil.Error by using os.rename directly. os.rename also provides
better guarantees for parallel-safeness. (Unfortunately, shutil.move has a
crappy implementation and uses copy/delete in the case of an OSError which then
makes it not as parallel-safe.)

[dnj (Google), 2015/09/17 21:16:39]

Hah did that before I read your comment. Done.

+      # Something else did this in parallel.
+      pass
+
+    # Final check that the gsutil bin is okay.  This should never fail.
+    if not check_gsutil(gsutil_bin):
+      raise InvalidGsutilError()
+  finally:
+    if os.path.isdir(download_dir):
+      shutil.rmtree(download_dir)
   return gsutil_bin
 
 
-def run_gsutil(force_version, fallback, target, args):
+def run_gsutil(force_version, fallback, target, args, clean=False):
   if force_version:
-    gsutil_bin = ensure_gsutil(force_version, target)
+    gsutil_bin = ensure_gsutil(force_version, target, clean)
   else:
     gsutil_bin = fallback
   cmd = [sys.executable, gsutil_bin] + args
   return subprocess.call(cmd)
 
 
 def parse_args():
+  bin_dir = os.environ.get('DEPOT_TOOLS_GSUTIL_BIN_DIR', DEFAULT_BIN_DIR)
+
   parser = argparse.ArgumentParser()
   parser.add_argument('--force-version', default='4.13')
+  parser.add_argument('--clean', action='store_true',
+      help='Clear any existing gsutil package, forcing a new download.')
   parser.add_argument('--fallback', default=DEFAULT_FALLBACK_GSUTIL)
-  parser.add_argument('--target', default=DEFAULT_BIN_DIR)
+  parser.add_argument('--target', default=bin_dir,
+      help='The target directory to download/store a gsutil version in. '
+           '(default is %(default)s).')
   parser.add_argument('args', nargs=argparse.REMAINDER)
 
   args, extras = parser.parse_known_args()
   if args.args and args.args[0] == '--':
     args.args.pop(0)
   if extras:
     args.args = extras + args.args
-  return args.force_version, args.fallback, args.target, args.args
+  return args
 
 
 def main():
-  force_version, fallback, target, args = parse_args()
-  return run_gsutil(force_version, fallback, target, args)
+  args = parse_args()
+  return run_gsutil(args.force_version, args.fallback, args.target, args.args,
+                    clean=args.clean)
 
 if __name__ == '__main__':
   sys.exit(main())