annotated_run.py: Add LogDog bootstrapping.

scripts/slave/gce.py

Index: scripts/slave/gce.py
diff --git a/scripts/slave/gce.py b/scripts/slave/gce.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6f69abe9e1bc76e6fcd7d8f6bd1e4817aeb282a
--- /dev/null
+++ b/scripts/slave/gce.py
@@ -0,0 +1,101 @@
+# Copyright 2015 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.
+
+"""
+Utilities for interfacing with Google Compute Engine.

[iannucci, 2016/01/15 04:18:18]

is this stripped-down? No reason to have code here that's not used. If we need
to add something, we can add it later. If you want it to be an exact copy, then
let's put it in third_party.

[dnj, 2016/01/15 22:05:50]

Stripped it down a bit.

+"""
+
+import httplib
+import json
+import logging
+import socket
+import time
+import urlparse
+
+
+LOGGER = logging.getLogger('gce')
+TRY_LIMIT = 5
+
+
+class Authenticator(object):
+  """Authenticator implementation that uses GCE metadata service for token.
+  """
+
+  _INFO_URL = 'http://metadata.google.internal'
+  _ACQUIRE_URL = ('http://metadata/computeMetadata/v1/instance/'
+                  'service-accounts/default/token')
+  _ACQUIRE_HEADERS = {"Metadata-Flavor": "Google"}
+
+  _cache_is_gce = None
+  _token_cache = None
+  _token_expiration = None
+
+  @classmethod
+  def is_gce(cls):
+    if cls._cache_is_gce is None:
+      cls._cache_is_gce = cls._test_is_gce()
+    return cls._cache_is_gce
+
+  @classmethod
+  def _test_is_gce(cls):
+    # Based on https://cloud.google.com/compute/docs/metadata#runninggce
+    try:
+      resp = cls._get(cls._INFO_URL)
+    except socket.error:
+      # Could not resolve URL.
+      return False
+    return resp.getheader('Metadata-Flavor', None) == 'Google'
+
+  @staticmethod
+  def _get(url, **kwargs):
+    next_delay_sec = 1
+    for i in xrange(TRY_LIMIT):
+      if i > 0:
+        # Retry server error status codes.
+        LOGGER.info('Encountered server error; retrying after %d second(s).',
+                    next_delay_sec)
+        time.sleep(next_delay_sec)
+        next_delay_sec *= 2
+
+      p = urlparse.urlparse(url)
+      c = GetConnectionClass(protocol=p.scheme)(p.netloc)
+      c.request('GET', url, **kwargs)
+      resp = c.getresponse()
+      LOGGER.debug('GET [%s] #%d/%d (%d)', url, i+1, TRY_LIMIT, resp.status)
+      if resp.status < httplib.INTERNAL_SERVER_ERROR:
+        return resp
+
+
+  @classmethod
+  def _get_token_dict(cls):
+    if cls._token_cache:
+      # If it expires within 25 seconds, refresh.
+      if cls._token_expiration < time.time() - 25:
+        return cls._token_cache
+
+    resp = cls._get(cls._ACQUIRE_URL, headers=cls._ACQUIRE_HEADERS)

[iannucci, 2016/01/15 04:18:18]

may want to log ("refreshing token") or something

[dnj, 2016/01/15 22:05:50]

This was removed as part of the strip.

+    if resp.status != httplib.OK:
+      return None
+    cls._token_cache = json.load(resp)
+    cls._token_expiration = cls._token_cache['expires_in'] + time.time()
+    return cls._token_cache
+
+  def get_auth_header(self, _host):
+    token_dict = self._get_token_dict()
+    if not token_dict:
+      return None
+    return '%(token_type)s %(access_token)s' % token_dict
+
+
+def GetConnectionClass(protocol=None):
+  if protocol is None:
+    protocol = 'https'
+  if protocol == 'https':
+    return httplib.HTTPSConnection
+  elif protocol == 'http':
+    return httplib.HTTPConnection
+  else:
+    raise RuntimeError(
+        "Don't know how to work with protocol '%s'" % protocol)
+