Refactor all gRPC proxy code into a single class.

client/utils/grpc_proxy.py

+# Copyright 2017 The LUCI Authors. All rights reserved.
+# Use of this source code is governed under the Apache License, Version 2.0
+# that can be found in the LICENSE file.
+
+"""Common gRPC implementation for Swarming and Isolate"""
+
+import logging
+import os
+import re
+import urlparse
+from utils import net
+
+# gRPC may not be installed on the worker machine. This is fine, as long as
+# the bot doesn't attempt to use gRPC (checked in IsolateServerGrpc.__init__).
+# Full external requirements are: grpcio, certifi.
+try:
+  import grpc
+  from google import auth as google_auth
+  from google.auth.transport import grpc as google_auth_transport_grpc
+  from google.auth.transport import requests as google_auth_transport_requests
+except ImportError as err:
+  grpc = None
+
+
+# If gRPC was successfully imported, try to import certifi as well.  This is not
+# actually used anywhere in this module, but if certifi is missing,
+# google.auth.transport will fail (see
+# https://stackoverflow.com/questions/24973326). So checking it here allows us
+# to print out a somewhat-sane error message.
+certifi = None
+if grpc is not None:
+  try:
+    import certifi
+  except ImportError:
+    # Will print out error messages later (ie when we have a logger)
+    pass
+
+
+# How many times to retry a gRPC call
+MAX_GRPC_ATTEMPTS = 30
+
+
+# Longest time to sleep between gRPC calls
+MAX_GRPC_SLEEP = 10.
+
+
+# Start the timeout at three minutes.
+GRPC_TIMEOUT_SEC = 3 * 60
+
+
+def available():
+  """Returns true if gRPC can be used on this host."""
+  return grpc != None
+
+
+class Proxy(object):
+  """Represents a gRPC proxy.
+
+  If the proxy begins with 'https', the returned channel will be secure and
+  authorized using default application credentials - see
+  https://developers.google.com/identity/protocols/application-default-credentials.
+  Currently, we're using Cloud Container Builder scopes for testing; this may
+  change in the future to allow different scopes to be passed in for different
+  channels.
+
+  To use the returned channel to call methods directly, say:
+
+    proxy = grpc_proxy.Proxy('https://grpc.luci.org/resource/prefix',
+                             myapi_pb2.MyApiStub)
+
+  To make a unary call with retries (recommended):
+
+    proto_output = proxy.call_unary('MyMethod', proto_input)
+
+  To make a unary call without retries, or to pass in a client side stream
+  (proto_input can be an iterator here):
+
+    proto_output = proxy.call_no_retries('MyMethod', proto_input)
+
+  You can also call the stub directly (not recommended, since no errors will be
+  caught or logged):
+
+    proto_output = proxy.stub.MyMethod(proto_input)
+
+  To make a call to a server-side streaming call (these are not retried):
+
+    for response in proxy.get_stream('MyStreaminingMethod', proto_input):
+      <process response>
+
+  To retrieve the prefix:
+
+    prefix = proxy.prefix # returns "prefix/for/resource/names"
+
+  All exceptions are logged using "logging.warning."
+  """
+
+  def __init__(self, proxy, stub_class):
+    url = urlparse.urlparse(proxy)
+    if url.scheme is 'http':
+      self._secure = False
+    elif url.scheme is 'https':
+      self._secure = True
+    else:
+      raise ValueError('gRPC proxy must use http or https: %s' % proxy)
+    if url.netloc is '':
+      raise ValueError('gRPC proxy is missing hostname: %s' % proxy)
+    self._host = url.netloc
+    self._prefix = url.path
+    if not self._prefix.endswith('/'):
+      self._prefix += '/'
+    if url.params is not '' or url.fragment is not '':
+      raise ValueError('gRPC proxy may not contain params or fragments: %s' %
+                       proxy)
+    self._debug_info = ['full proxy name: ' + proxy]

[M-A Ruel, 2017/08/04 02:05:37]

instead of _debug_info, you can create a local logger, e.g.

log = logging.getLogger('gprc.proxy')

then use log.info()...

[aludwin, 2017/08/04 14:53:10]

The idea here is to *save* these bits of information unless/until something goes
wrong, and only then print them out. From what I've read, logging.Logger isn't
really a great fit for that - it seems pretty heavyweight for what I'm really
trying to do (delayed output).

+    self._channel = self._create_channel()
+    self._stub = stub_class(self._channel)
+    logging.info('%s: initialized', self.name)
+
+  @property

[M-A Ruel, 2017/08/04 02:05:37]

Most of these are not needed, you can add only once needed.

[aludwin, 2017/08/04 14:53:10]

Done.

+  def host(self):
+    return self._host
+
+  @property
+  def prefix(self):
+    return self._prefix
+
+  @property
+  def secure(self):
+    return self._secure
+
+  @property
+  def channel(self):
+    return self._channel
+
+  @property
+  def stub(self):
+    return self._stub
+
+  @property
+  def name(self):
+    security = 'insecure'
+    if self._secure:
+      security = 'secure'
+    return 'gRPC %s proxy %s/%s' % (
+        security, self.host, self._stub.__class__.__name__)
+
+  def call_unary(self, name, request):
+    """Calls a method, waiting if the service is not available.
+
+    Usage: proto_output = proxy.call_unary('MyMethod', proto_input)
+    """
+    grpc_error = None
+    for attempt in range(1, MAX_GRPC_ATTEMPTS+1):
+      try:
+        return self.call_no_retries(name, request)
+      except grpc.RpcError as g:
+        if g.code() is not grpc.StatusCode.UNAVAILABLE:
+          raise
+        logging.warning('%s: call_grpc - proxy is unavailable (attempt %d/%d)',
+                        self.name, attempt, MAX_GRPC_ATTEMPTS)
+        # Save the error in case we need to return it
+        grpc_error = g
+        time.sleep(net.calculate_sleep_before_retry(attempt, MAX_GRPC_SLEEP))
+    # If we get here, it must be because we got (and saved) an error
+    assert grpc_error is not None
+    raise grpc_error
+
+  def get_stream(self, name, request):
+    """Calls a server-side streaming method, returning an iterator.
+
+    Usage: for resp in proxy.get_stream('MyMethod', proto_input'):
+    """
+    def wrap_grpc_iterator(stream, name):
+      while True:
+        # The lambda "next(stream, 1)" will return a protobuf on success, or the
+        # integer 1 if the stream has ended. This allows us to avoid attempting
+        # to catch StopIteration, which gets logged by _wrap_grpc_operation.
+        response = self._wrap_grpc_operation(name + ' pull from stream',
+                                             lambda: next(stream, 1))
+        if isinstance(response, int):
+          # Iteration is finished
+          return
+        yield response
+    stream = self.call_no_retries(name, request)
+    return wrap_grpc_iterator(stream, name)

[M-A Ruel, 2017/08/04 02:05:37]

There's still no value in wrap_grpc_iterator, you can rewrite the whole function
as a generator:

def get_stream(self, name, request):
  stream = self.call_no_retries(name, request)
  while True:
    response = self._wrap_grpc_operation(name + ' pull from stream', lambda:
next(stream, 1))
    if isinstance(response, int):
      return
    yield response


that's it! (plus comments)

[aludwin, 2017/08/04 14:53:10]

Oh, right :P Thanks!

+
+  def call_no_retries(self, name, request):
+    """Calls a method without any retries.
+
+    Recommended for client-side streaming or nonidempotent unary calls.
+    """
+    method = getattr(self._stub, name)
+    if method is None:
+      raise NameError('%s: "%s" is not a valid method name', self.name, name)
+    return self._wrap_grpc_operation(
+        name, lambda: method(request, timeout=GRPC_TIMEOUT_SEC))
+
+  def _wrap_grpc_operation(self, name, fn):
+    """Wraps a gRPC operation (call or iterator increment) for logging."""
+    try:
+      return fn()
+    except grpc.RpcError as g:
+      logging.warning('%s/%s - gRPC error: %s', self.name, name, g)
+      self._dump_proxy_info()
+      raise g
+    except Exception as e:
+      logging.warning('caught exception %s; none is %s', repr(e), e is None)
+      logging.warning('%s/%s - generic error: %s', self.name, name, e)
+      self._dump_proxy_info()
+      raise e
+
+  def _dump_proxy_info(self):
+    logging.warning('%s: DETAILED PROXY INFO', self.name)
+    logging.warning('%s: prefix = %s', self.name, self.prefix)
+    logging.warning('%s: debug info:\n\t%s', self.name,
+                    '\n\t'.join(self._debug_info))
+
+  def _create_channel(self):
+    # Make sure grpc was successfully imported
+    assert available()
+
+    if not self._secure:
+      return grpc.insecure_channel(self._host)
+
+    # Authenticate the host.
+    #
+    # You're allowed to override the root certs and server if necessary. For
+    # example, if you're running your proxy on localhost, you'll need to set
+    # GRPC_PROXY_TLS_ROOTS to the "roots.crt" file specifying the certificate for

[M-A Ruel, 2017/08/04 02:05:37]

wrap all

[aludwin, 2017/08/04 14:53:10]

Done.

+    # the root CA that your localhost server has used to certify itself, and the
+    # GRPC_PROXY_TLS_OVERRIDE to the name that your server is using to identify
+    # itself. For example, the ROOTS env var might be "/path/to/roots.crt" while
+    # the OVERRIDE env var might be "test_server," if this is what's used by the
+    # server you're running.
+    #
+    # If you're connecting to a real server with real SSL, none of this should be
+    # used.
+    if not certifi:
+      self._debug_info.append('CERTIFI IS NOT PRESENT; gRPC HTTPS CONNECTIONS MAY FAIL')
+    root_certs = None
+    roots = os.environ.get('GRPC_PROXY_TLS_ROOTS')
+    if roots is not None:

[M-A Ruel, 2017/08/04 02:05:37]

if roots:

[aludwin, 2017/08/04 14:53:10]

Done.

+      self._debug_info.append('Overridden root CA: %s' % roots)
+      with open(roots) as f:
+        root_certs = f.read()
+    else:
+      self._debug_info.append('Using default root CAs from certifi')
+    overd = os.environ.get('GRPC_PROXY_TLS_OVERRIDE')
+    options = ()
+    if overd is not None:

[M-A Ruel, 2017/08/04 02:05:37]

if overd:

[aludwin, 2017/08/04 14:53:10]

Done.

+      options=(('grpc.ssl_target_name_override', overd),)
+    ssl_creds = grpc.ssl_channel_credentials(root_certificates=root_certs)
+
+    # Authenticate the user.
+    scopes = ('https://www.googleapis.com/auth/cloud-source-tools',)
+    self._debug_info.append('Scopes are: %s' % repr(scopes))

[M-A Ruel, 2017/08/04 02:05:37]

%r instead of repr

[aludwin, 2017/08/04 14:53:10]

Done.

+    user_creds, _ = google_auth.default(scopes=scopes)
+
+    # Create the channel.
+    request = google_auth_transport_requests.Request()
+    self._debug_info.append('Options are: %s' % repr(options))
+    return google_auth_transport_grpc.secure_authorized_channel(
+        user_creds, request, self._host, ssl_creds, options=options)