Adding the initial code for Omnibot multi-machine support

testing/legion/client_lib.py

Index: testing/legion/client_lib.py
diff --git a/testing/legion/client_lib.py b/testing/legion/client_lib.py
new file mode 100644
index 0000000000000000000000000000000000000000..2576e4f9d41ac8415f7425367762bab824c977b4
--- /dev/null
+++ b/testing/legion/client_lib.py
@@ -0,0 +1,240 @@
+# 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.
+"""Client library module."""
+
+import argparse
+import logging
+import os
+import socket
+import subprocess
+import tempfile
+import threading
+import uuid
+import xmlrpclib
+
+#pylint: disable=relative-import
+import common_lib
+import client_rpc_server
+
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+SWARMING_DIR = os.path.join(THIS_DIR, '..', '..', 'tools/swarming_client')
+ISOLATE_PY = os.path.join(SWARMING_DIR, 'isolate.py')
+SWARMING_PY = os.path.join(SWARMING_DIR, 'swarming.py')
+# ISOLATE_SERVER = 'omnibot-isolate-server.appspot.com'
+# SWARMING_SERVER = 'https://omnibot-swarming-server.appspot.com/'
+CLIENT_CONNECTION_TIMEOUT = 30 * 60  # 30 minutes
+
+
+class Error(Exception):
+  pass
+
+
+class ConnectionTimeoutError(Error):
+  pass
+
+
+class Client(object):
+  """The main client class.

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

Few issues:
- The class name is misleading, it's not one client. It's more akin a
ClientsManager.
- This first docstring line is information-less. Starts it with a verb about
what it's doing.

[Mike Meade, 2015/02/03 01:18:08]

Done.

+
+  This class is used to create clients, wait for them to connect to the host,
+  connect to their RPC servers, and run RPC commands.
+  """
+
+  _client_count = 0
+
+  def __init__(self, isolate_file, discovery_server, name=None):
+    """ctor.
+
+    Args:
+      isolate_file: The path to the client isolate file.
+      discovery_server: The discovery server to register with
+      name: A name for the client.
+    """
+    self._IncreaseCount()
+    self.name = name or self._CreateName()
+    self.priority = 100
+    self.isolate_file = isolate_file
+    self.isolated_file = isolate_file + 'd'
+    self.rpc_timeout = None
+    self._connect_event = threading.Event()
+    self._args = self._ParseArgs()
+    self._discovery_server = discovery_server
+    self._connected = False
+    self._ip_address = None
+    self._config_vars = []
+    self._dimensions = []

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

this must be a dict, since duplicate keys are not acceptable.

[Mike Meade, 2015/02/03 01:18:08]

Done.

+    self._otp = str(uuid.uuid1())
+    self._rpc = None
+    self._verbose = False
+
+  @property
+  def connected(self):
+    """Return the value of self._connected."""
+    return self._connected
+
+  @property
+  def connect_event(self):
+    return self._connect_event
+
+  @property
+  def rpc(self):
+    return self._rpc
+
+  @classmethod
+  def _IncreaseCount(cls):

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:38]

Why a function if it is used exactly at one place?

[Mike Meade, 2015/02/03 01:18:08]

Changed it to type(self)._client_count += 1 in the __init__ method. Using a
simple "self._client_count += 1" in __init__ results in the object getting the
value rather than the class. This resulted in multiple Client1 objects.

+    """Increase the client_count parameter."""
+    cls._client_count += 1
+
+  @classmethod
+  def _CreateName(cls):
+    """Create a name for this client.
+
+    By default the name is "Client%d" where %d is the number of clients that
+    currently exist.
+    """
+    return 'Client%d' % cls._client_count
+
+  def _ParseArgs(self):

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

Why is this a method?

[Mike Meade, 2015/02/03 01:18:08]

Moved it into __init__

+    """Parse command line args."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--isolate-server')
+    parser.add_argument('--swarming-server')
+    parser.add_argument('--client-connection-timeout',
+                        default=CLIENT_CONNECTION_TIMEOUT)
+    args, _ = parser.parse_known_args()
+    return args
+
+  def AddConfigVars(self, key, value):
+    """Add a set of config vars to isolate.py.
+
+    Args:
+      key: The config vars key.

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

Frankly, I don't see the value of this documentation. It's pretty
self-explanatory.

[Mike Meade, 2015/02/03 01:18:08]

Done.

+      value: The config vars value.
+    """
+    logging.debug('Adding --config-var %s %s to %s', key, value,
+                  self.name)
+    self._config_vars.append((key, value))

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

The fact previous keys are not overridden IS something worth documenting.

[Mike Meade, 2015/02/03 01:18:07]

I added a note, but should this also be a dict?

+
+  def AddDimension(self, key, value):
+    """Add a set of dimensions to swarming.py.
+
+    Args:
+      key: The dimension key.

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

same

[Mike Meade, 2015/02/03 01:18:08]

Done.

+      value: The dimension value.
+    """
+    logging.debug('Adding --dimension %s %s to %s', key, value,
+                  self.name)
+    self._dimensions.append((key, value))
+
+  def Create(self, wait=False, timeout=None):

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:38]

Why options? Will there be code that will use all the combinations?

[Mike Meade, 2015/02/03 01:18:07]

In thinking about it more this is probably over-engineered. It's probably
cleaner just to have the calling code call Create() and
WaitForConnection(timeout).

+    """Create the client machine and wait for it to be created if specified.
+
+    Args:
+      wait: True to block until created, False to return immediately.
+      timeout: The timeout to block before raising a ConnectionTimeoutError.
+    """
+    logging.info('Creating %s', self.name)
+    self._connect_event.clear()
+    self._RegisterOnConnectCallback()
+    self._ExecuteIsolate()
+    self._ExecuteSwarming()
+    if wait:
+      self.WaitForConnection(timeout)
+
+  def WaitForConnection(self, timeout=None):
+    """Connect to the client machine.
+
+    This method waits for the client machine to register itself
+    with the discovery server.
+
+    Args:
+      timeout: The timeout in seconds.
+
+    Raises:
+      TimeoutError if the client doesn't connect in time.
+    """
+    timeout = timeout or self._args.client_connection_timeout
+    msg = ('Waiting for %s to connect with a timeout of %d seconds' %
+           (self.name, timeout))
+    logging.info(msg)
+    self._connect_event.wait(timeout)
+    if not self._connect_event.is_set():
+      raise ConnectionTimeoutError()
+
+  def Release(self):
+    """Quit the client's RPC server so it can release the machine."""
+    if self._rpc is not None and self._connected:
+      logging.info('Releasing %s', self.name)
+      try:
+        self._rpc.Quit()
+      except (socket.error, xmlrpclib.Fault):
+        logging.error('Unable to connect to %s to call Quit', self.name)
+      self._rpc = None
+      self._connected = False
+
+  def _ExecuteIsolate(self):
+    """Execute isolate.py with the given args."""
+    cmd = [
+        'python',
+        ISOLATE_PY,
+        'archive',
+        '--isolate=' + self.isolate_file,

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

Trust me, that's a recipe for bugs. Use instead:
'--isolate', self.isolate_file,
and the like.

[Mike Meade, 2015/02/03 01:18:07]

Done.

+        '--isolated=' + self.isolated_file,
+        ]
+
+    if self._args.isolate_server:
+      cmd += ['--isolate-server', self._args.isolate_server]

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

I generally prefer cmd.extend(...)

[Mike Meade, 2015/02/03 01:18:07]

Done.

+    for key, value in self._config_vars:
+      cmd += ['--config-var', key, value]
+
+    logging.debug('Running %s', ' '.join(cmd))
+    if subprocess.call(cmd, stdout=subprocess.PIPE) != 0:

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

This may hang BTW since you are not reading stdout.

[Mike Meade, 2015/02/03 01:18:08]

Done.

+      raise Error('Error calling isolate.py')
+
+  def _ExecuteSwarming(self):
+    """Execute swarming.py with the vars."""
+    cmd = [
+        'python',
+        SWARMING_PY,
+        'trigger',
+        self.isolated_file,
+        '--priority=' + str(self.priority)
+        ]
+
+    if self._args.isolate_server:
+      cmd += ['--isolate-server', self._args.isolate_server]
+    if self._args.swarming_server:
+      cmd += ['--swarming', self._args.swarming_server]
+    for key, value in self._dimensions:
+      cmd += ['--dimension', key, value]
+
+    cmd += ['--', '--host', str(common_lib.MY_IP), '--otp', self._otp]
+    if self.rpc_timeout:
+      cmd += ['--idle-timeout', str(self.rpc_timeout)]
+    if self._verbose:
+      cmd += ['--verbose']
+
+    logging.debug('Running %s', ' '.join(cmd))
+    if subprocess.call(cmd, stdout=subprocess.PIPE) != 0:
+      raise Error('Error calling swarming.py')
+
+  def _RegisterOnConnectCallback(self):

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

Same thing, it's used exactly at one place, I don't see value in making it a
function, add a comment at the call site instead if needed but I think that:
self._discovery_server.RegisterClientCallback(self._otp, self._OnConnect)
is completely self-explanatory.

[Mike Meade, 2015/02/03 01:18:08]

Done.

+    """Register a callback with the discovery server.
+
+    This callback is used to receive the client's IP address once it starts

[Marc-Antoine Ruel (Google), 2015/01/30 21:58:39]

What the callback does is not relevant documentation to this function. It should
be commented on the callback itself.
What the callback receives should be documented in RegisterClientCallback(). No
need to repeat it elsewhere unless the arguments are illnamed due to historical
reasons or something like that.

[Mike Meade, 2015/02/03 01:18:07]

Acknowledged.

+    and contacts the discovery server.
+    """
+    self._discovery_server.RegisterClientCallback(self._otp, self._OnConnect)
+
+  def _OnConnect(self, ip_address):
+    """The OnConnect callback method.
+
+    This method receives the ip address received by the discovery server from
+    the client and sets the object's connected state to True.
+    """
+    self._ip_address = ip_address
+    self._connected = True
+    self._rpc = client_rpc_server.RPCServer.Connect(self._ip_address)
+    logging.info('%s connected from %s', self.name, ip_address)
+    self._connect_event.set()