Add portend-2.6

lib/portend.py

@@ -0,0 +1,222 @@
+# -*- coding: utf-8 -*-
+
+"""
+A simple library for managing the availability of ports.
+"""
+
+from __future__ import print_function, division
+
+import time
+import socket
+import argparse
+import sys
+import itertools
+import contextlib
+import platform
+
+try:
+    from collections import abc
+except ImportError:
+    import collections as abc
+
+from tempora import timing
+
+
+def client_host(server_host):
+    """Return the host on which a client can connect to the given listener."""
+    if server_host == '0.0.0.0':
+        # 0.0.0.0 is INADDR_ANY, which should answer on localhost.
+        return '127.0.0.1'
+    if server_host in ('::', '::0', '::0.0.0.0'):
+        # :: is IN6ADDR_ANY, which should answer on localhost.
+        # ::0 and ::0.0.0.0 are non-canonical but common
+        # ways to write IN6ADDR_ANY.
+        return '::1'
+    return server_host
+
+
+class Checker(object):
+    def __init__(self, timeout=1.0):
+        self.timeout = timeout
+
+    def assert_free(self, host, port=None):
+        """
+        Assert that the given addr is free
+        in that all attempts to connect fail within the timeout
+        or raise a PortNotFree exception.
+
+        >>> free_port = find_available_local_port()
+
+        >>> Checker().assert_free('localhost', free_port)
+        >>> Checker().assert_free('127.0.0.1', free_port)
+        >>> Checker().assert_free('::1', free_port)
+
+        Also accepts an addr tuple
+
+        >>> addr = '::1', free_port, 0, 0
+        >>> Checker().assert_free(addr)
+
+        Host might refer to a server bind address like '::', which
+        should use localhost to perform the check.
+
+        >>> Checker().assert_free('::', free_port)
+        """
+        if port is None and isinstance(host, abc.Sequence):
+            host, port = host[:2]
+        if platform.system() == 'Windows':
+            host = client_host(host)
+        info = socket.getaddrinfo(host, port, socket.AF_UNSPEC, socket.SOCK_STREAM)
+        list(itertools.starmap(self._connect, info))
+
+    def _connect(self, af, socktype, proto, canonname, sa):
+        s = socket.socket(af, socktype, proto)
+        # fail fast with a small timeout
+        s.settimeout(self.timeout)
+
+        with contextlib.closing(s):
+            try:
+                s.connect(sa)
+            except socket.error:
+                return
+
+        # the connect succeeded, so the port isn't free
+        port, host = sa[:2]
+        tmpl = "Port {port} is in use on {host}."
+        raise PortNotFree(tmpl.format(**locals()))
+
+
+class Timeout(IOError):
+    pass
+
+
+class PortNotFree(IOError):
+    pass
+
+
+def free(host, port, timeout=float('Inf')):
+    """
+    Wait for the specified port to become free (dropping or rejecting
+    requests). Return when the port is free or raise a Timeout if timeout has
+    elapsed.
+
+    Timeout may be specified in seconds or as a timedelta.
+    If timeout is None or ∞, the routine will run indefinitely.
+
+    >>> free('localhost', find_available_local_port())
+    """
+    if not host:
+        raise ValueError("Host values of '' or None are not allowed.")
+
+    timer = timing.Timer(timeout)
+
+    while True:
+        try:
+            # Expect a free port, so use a small timeout
+            Checker(timeout=0.1).assert_free(host, port)
+            return
+        except PortNotFree:
+            if timer.expired():
+                raise Timeout("Port {port} not free on {host}.".format(**locals()))
+            # Politely wait.
+            time.sleep(0.1)
+
+
+wait_for_free_port = free
+
+
+def occupied(host, port, timeout=float('Inf')):
+    """
+    Wait for the specified port to become occupied (accepting requests).
+    Return when the port is occupied or raise a Timeout if timeout has
+    elapsed.
+
+    Timeout may be specified in seconds or as a timedelta.
+    If timeout is None or ∞, the routine will run indefinitely.
+
+    >>> occupied('localhost', find_available_local_port(), .1)
+    Traceback (most recent call last):
+    ...
+    Timeout: Port ... not bound on localhost.
+    """
+    if not host:
+        raise ValueError("Host values of '' or None are not allowed.")
+
+    timer = timing.Timer(timeout)
+
+    while True:
+        try:
+            Checker(timeout=0.5).assert_free(host, port)
+            if timer.expired():
+                raise Timeout("Port {port} not bound on {host}.".format(**locals()))
+            # Politely wait
+            time.sleep(0.1)
+        except PortNotFree:
+            # port is occupied
+            return
+
+
+wait_for_occupied_port = occupied
+
+
+def find_available_local_port():
+    """
+    Find a free port on localhost.
+
+    >>> 0 < find_available_local_port() < 65536
+    True
+    """
+    infos = socket.getaddrinfo(None, 0, socket.AF_UNSPEC, socket.SOCK_STREAM)
+    family, proto, _, _, addr = next(iter(infos))
+    sock = socket.socket(family, proto)
+    sock.bind(addr)
+    addr, port = sock.getsockname()[:2]
+    sock.close()
+    return port
+
+
+class HostPort(str):
+    """
+    A simple representation of a host/port pair as a string
+
+    >>> hp = HostPort('localhost:32768')
+
+    >>> hp.host
+    'localhost'
+
+    >>> hp.port
+    32768
+
+    >>> len(hp)
+    15
+    """
+
+    @property
+    def host(self):
+        host, sep, port = self.partition(':')
+        return host
+
+    @property
+    def port(self):
+        host, sep, port = self.partition(':')
+        return int(port)
+
+
+def _main():
+    parser = argparse.ArgumentParser()
+
+    def global_lookup(key):
+        return globals()[key]
+
+    parser.add_argument('target', metavar='host:port', type=HostPort)
+    parser.add_argument('func', metavar='state', type=global_lookup)
+    parser.add_argument('-t', '--timeout', default=None, type=float)
+    args = parser.parse_args()
+    try:
+        args.func(args.target.host, args.target.port, timeout=args.timeout)
+    except Timeout as timeout:
+        print(timeout, file=sys.stderr)
+        raise SystemExit(1)
+
+
+if __name__ == '__main__':
+    _main()