1 files changed, 53 insertions, 13 deletions

diff --git a/benchmarks/harness/tunnel_dispatcher.py b/benchmarks/harness/tunnel_dispatcher.py
index 8dfe2862a..c56fd022a 100644
--- a/benchmarks/harness/tunnel_dispatcher.py
+++ b/benchmarks/harness/tunnel_dispatcher.py
@@ -11,7 +11,25 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tunnel handles setting up connections to remote machines."""
+"""Tunnel handles setting up connections to remote machines.
+
+Tunnel dispatcher is a wrapper around the connection from a local UNIX socket
+and a remote UNIX socket via SSH with port forwarding. This is done to
+initialize the pythonic dockerpy client to run containers on the remote host by
+connecting to /var/run/docker.sock (where Docker is listening). Tunnel
+dispatcher sets up the local UNIX socket and calls the `ssh` command as a
+subprocess, and holds a reference to that subprocess. It manages clean-up on
+exit as best it can by killing the ssh subprocess and deleting the local UNIX
+socket,stored in /tmp for easy cleanup in most systems if this fails.
+
+  Typical usage example:
+
+  t = Tunnel(name, **kwargs)
+  t.connect()
+  client = t.get_docker_client() #
+  client.containers.run("ubuntu", "echo hello world")
+
+"""
 
 import os
 import tempfile
@@ -21,31 +39,53 @@ import docker
 import pexpect
 
 SSH_TUNNEL_COMMAND = """ssh
- -o GlobalKnownHostsFile=/dev/null
- -o UserKnownHostsFile=/dev/null
- -o StrictHostKeyChecking=no
- -nNT -L {filename}:/var/run/docker.sock
- -i {key_path}
- {username}@{hostname}"""
+    -o GlobalKnownHostsFile=/dev/null
+    -o UserKnownHostsFile=/dev/null
+    -o StrictHostKeyChecking=no
+    -o IdentitiesOnly=yes
+    -nNT -L {filename}:/var/run/docker.sock
+    -i {key_path}
+    {username}@{hostname}"""
 
 
-class Tunnel:
+class Tunnel(object):
   """The tunnel object represents the tunnel via ssh.
 
   This connects a local unix domain socket with a remote socket.
+
+  Attributes:
+      _filename: a temporary name of the UNIX socket prefixed by the name
+        argument.
+      _hostname: the IP or resolvable hostname of the remote host.
+      _username: the username of the ssh_key used to run ssh.
+      _key_path: path to a valid key.
+      _key_password: optional password to the ssh key in _key_path
+      _process: holds reference to the ssh subprocess created.
+
+    Returns:
+      The new minimum port.
+
+    Raises:
+      ConnectionError: If no available port is found.
   """
 
-  def __init__(self, name, hostname: str, username: str, key_path: str,
+  def __init__(self,
+               name: str,
+               hostname: str,
+               username: str,
+               key_path: str,
+               key_password: str = "",
                **kwargs):
     self._filename = tempfile.NamedTemporaryFile(prefix=name).name
     self._hostname = hostname
     self._username = username
     self._key_path = key_path
+    self._key_password = key_password
     self._kwargs = kwargs
     self._process = None
 
   def connect(self):
-    """Connects the SSH tunnel."""
+    """Connects the SSH tunnel and stores the subprocess reference in _process."""
     cmd = SSH_TUNNEL_COMMAND.format(
         filename=self._filename,
         key_path=self._key_path,
@@ -54,9 +94,9 @@ class Tunnel:
     self._process = pexpect.spawn(cmd, timeout=10)
 
     # If given a password, assume we'll be asked for it.
-    if "key_password" in self._kwargs:
+    if self._key_password:
       self._process.expect(["Enter passphrase for key .*: "])
-      self._process.sendline(self._kwargs["key_password"])
+      self._process.sendline(self._key_password)
 
     while True:
       # Wait for the tunnel to appear.
@@ -71,7 +111,7 @@ class Tunnel:
     return self._filename
 
   def get_docker_client(self):
-    """Returns a docker client for this Tunne0l."""
+    """Returns a docker client for this Tunnel."""
     return docker.DockerClient(base_url="unix:/" + self._filename)
 
   def __del__(self):