comparison mercurial/keepalive.py @ 43076:2372284d9457

formatting: blacken the codebase

This is using my patch to black
(https://github.com/psf/black/pull/826) so we don't un-wrap collection
literals.

Done with:

  hg files 'set:**.py - mercurial/thirdparty/** - "contrib/python-zstandard/**"' | xargs black -S

# skip-blame mass-reformatting only

# no-check-commit reformats foo_bar functions

Differential Revision: https://phab.mercurial-scm.org/D6971

43075:57875cf423c9

    node,
    pycompat,
    urllibcompat,
    util,
)
from .utils import (
    procutil,
)

httplib = util.httplib
urlerr = util.urlerr
urlreq = util.urlreq

DEBUG = None

class ConnectionManager(object):
    """
    The connection manager must be able to:
      * keep track of all existing
      """
    def __init__(self):
        self._lock = threading.Lock()
        self._hostmap = collections.defaultdict(list) # host -> [connection]
        self._connmap = {} # map connections to host
        self._readymap = {} # map connection to ready state

    def add(self, host, connection, ready):
        self._lock.acquire()
        try:
            self._hostmap[host].append(connection)

        if host:
            return list(self._hostmap[host])
        else:
            return dict(self._hostmap)

class KeepAliveHandler(object):
    def __init__(self, timeout=None):
        self._cm = ConnectionManager()
        self._timeout = timeout
        self.requestscount = 0

                h = self._cm.get_ready_conn(host)
            else:
                # no (working) free connections were found.  Create a new one.
                h = http_class(host, timeout=self._timeout)
                if DEBUG:
                    DEBUG.info("creating new connection to %s (%d)",
                               host, id(h))
                self._cm.add(host, h, False)
                self._start_transaction(h, req)
                r = h.getresponse()
        # The string form of BadStatusLine is the status line. Add some context
        # to make the error message slightly more useful.
        except httplib.BadStatusLine as err:
            raise urlerr.urlerror(
                _('bad HTTP status line: %s') % pycompat.sysbytes(err.line))
        except (socket.error, httplib.HTTPException) as err:
            raise urlerr.urlerror(err)

        # If not a persistent connection, don't try to reuse it. Look
        # for this using getattr() since vcr doesn't define this

            r = h.getresponse()
            # note: just because we got something back doesn't mean it
            # worked.  We'll check the version below, too.
        except (socket.error, httplib.HTTPException):
            r = None
        except: # re-raises
            # adding this block just in case we've missed
            # something we will still raise the exception, but
            # lets try and close the connection and remove it
            # first.  We previously got into a nasty loop
            # where an exception was uncaught, and so the
            # connection stayed open.  On the next try, the
            # same exception was raised, etc.  The trade-off is
            # that it's now possible this call will raise
            # a DIFFERENT exception
            if DEBUG:
                DEBUG.error("unexpected exception - closing "
                            "connection to %s (%d)", host, id(h))
            self._cm.remove(h)
            h.close()
            raise

        if r is None or r.version == 9:
            # httplib falls back to assuming HTTP 0.9 if it gets a
            # bad header back.  This is most likely to happen if
            # the socket has been closed by the server since we
            # last used the connection.
            if DEBUG:
                DEBUG.info("failed to re-use connection to %s (%d)",
                           host, id(h))
            r = None
        else:
            if DEBUG:
                DEBUG.info("re-using connection to %s (%d)", host, id(h))

                skipheaders[r'skip_' + n.replace(r'-', r'_')] = 1
        try:
            if urllibcompat.hasdata(req):
                data = urllibcompat.getdata(req)
                h.putrequest(
                    req.get_method(), urllibcompat.getselector(req),
                    **skipheaders)
                if r'content-type' not in headers:
                    h.putheader(r'Content-type',
                                r'application/x-www-form-urlencoded')
                if r'content-length' not in headers:
                    h.putheader(r'Content-length', r'%d' % len(data))
            else:
                h.putrequest(
                    req.get_method(), urllibcompat.getselector(req),
                    **skipheaders)
        except socket.error as err:
            raise urlerr.urlerror(err)
        for k, v in headers.items():
            h.putheader(k, v)
        h.endheaders()

        self.sentbytescount += getattr(h, 'sentbytescount', 0) - oldbytescount

        try:
            self.parent.requestscount += 1
            self.parent.sentbytescount += (
                getattr(h, 'sentbytescount', 0) - oldbytescount)
        except AttributeError:
            pass

class HTTPHandler(KeepAliveHandler, urlreq.httphandler):
    pass

class HTTPResponse(httplib.HTTPResponse):
    # we need to subclass HTTPResponse in order to
    # 1) add readline(), readlines(), and readinto() methods
    # 2) add close_connection() methods

    # the read method wraps the original to accommodate buffering,
    # although read() never adds to the buffer.
    # Both readline and readlines have been stolen with almost no
    # modification from socket.py


    def __init__(self, sock, debuglevel=0, strict=0, method=None):
        extrakw = {}
        if not pycompat.ispy3:
            extrakw[r'strict'] = True
            extrakw[r'buffering'] = True
        httplib.HTTPResponse.__init__(self, sock, debuglevel=debuglevel,
                                      method=method, **extrakw)
        self.fileno = sock.fileno
        self.code = None
        self.receivedbytescount = 0
        self._rbuf = ''
        self._rbufsize = 8096
        self._handler = None # inserted by the handler later
        self._host = None    # (same)
        self._url = None     # (same)
        self._connection = None # (same)

    _raw_read = httplib.HTTPResponse.read
    _raw_readinto = getattr(httplib.HTTPResponse, 'readinto', None)

    # Python 2.7 has a single close() which closes the socket handle.

    def close(self):
        if self.fp:
            self.fp.close()
            self.fp = None
            if self._handler:
                self._handler._request_closed(self, self._host,
                                              self._connection)

    def _close_conn(self):
        self.close()

    def close_connection(self):

        while True:
            if chunk_left is None:
                line = self.fp.readline()
                i = line.find(';')
                if i >= 0:
                    line = line[:i] # strip chunk-extensions
                try:
                    chunk_left = int(line, 16)
                except ValueError:
                    # close the connection as protocol synchronization is
                    # probably lost

            else:
                parts.append(self._safe_read(chunk_left))
                amt -= chunk_left

            # we read the whole chunk, get another
            self._safe_read(2)      # toss the CRLF at the end of the chunk
            chunk_left = None

        # read and discard trailer up to the CRLF terminator
        ### note: we shouldn't have any trailers!
        while True:

    def readinto(self, dest):
        if self._raw_readinto is None:
            res = self.read(len(dest))
            if not res:
                return 0
            dest[0:len(res)] = res
            return len(res)
        total = len(dest)
        have = len(self._rbuf)
        if have >= total:
            dest[0:total] = self._rbuf[:total]

        dest[0:have] = self._rbuf
        got += len(self._rbuf)
        self._rbuf = ''
        return got

def safesend(self, str):
    """Send `str' to the server.

    Shamelessly ripped off from httplib to patch a bad behavior.

        else:
            self.sock.sendall(str)
            self.sentbytescount += len(str)
    except socket.error as v:
        reraise = True
        if v.args[0] == errno.EPIPE:      # Broken pipe
            if self._HTTPConnection__state == httplib._CS_REQ_SENT:
                self._broken_pipe_resp = None
                self._broken_pipe_resp = self.getresponse()
                reraise = False
            self.close()
        if reraise:
            raise

def wrapgetresponse(cls):
    """Wraps getresponse in cls with a broken-pipe sane version.
    """
    def safegetresponse(self):
        # In safesend() we might set the _broken_pipe_resp
        # attribute, in which case the socket has already
        # been closed and we just need to give them the response
        # back. Otherwise, we use the normal response path.
        r = getattr(self, '_broken_pipe_resp', None)
        if r is not None:
            return r
        return cls.getresponse(self)
    safegetresponse.__doc__ = cls.getresponse.__doc__
    return safegetresponse

class HTTPConnection(httplib.HTTPConnection):
    # url.httpsconnection inherits from this. So when adding/removing
    # attributes, be sure to audit httpsconnection() for unintended
    # consequences.

    def __init__(self, *args, **kwargs):
        httplib.HTTPConnection.__init__(self, *args, **kwargs)
        self.sentbytescount = 0
        self.receivedbytescount = 0

#########################################################################
#####   TEST FUNCTIONS
#########################################################################

            break
    fo.close()
    m = md5(foo)
    print(format % ('keepalive readline', node.hex(m.digest())))

def comp(N, url):
    print('  making %i connections to:\n  %s' % (N, url))

    procutil.stdout.write('  first using the normal urllib handlers')
    # first use normal opener

    urlreq.installopener(opener)
    t2 = fetch(N, url)
    print('  TIME: %.3f s' % t2)
    print('  improvement factor: %.2f' % (t1 / t2))

def fetch(N, url, delay=0):
    import time
    lens = []
    starttime = time.time()
    for i in range(N):
        if delay and i > 0:
            time.sleep(delay)

        if not i == lens[0]:
            print("WARNING: inconsistent length on read %i: %i" % (j, i))

    return diff

def test_timeout(url):
    global DEBUG
    dbbackup = DEBUG
    class FakeLogger(object):
        def debug(self, msg, *args):
            print(msg % args)
        info = warning = error = debug
    DEBUG = FakeLogger()
    print("  fetching the file to establish a connection")
    fo = urlreq.urlopen(url)
    data1 = fo.read()
    fo.close()

    comp(N, url)
    print('')
    print("performing dropped-connection check")
    test_timeout(url)

if __name__ == '__main__':
    import time
    try:
        N = int(sys.argv[1])
        url = sys.argv[2]
    except (IndexError, ValueError):
        print("%s <integer> <url>" % sys.argv[0])

43076:2372284d9457

    node,
    pycompat,
    urllibcompat,
    util,
)
from .utils import procutil

httplib = util.httplib
urlerr = util.urlerr
urlreq = util.urlreq

DEBUG = None


class ConnectionManager(object):
    """
    The connection manager must be able to:
      * keep track of all existing
      """

    def __init__(self):
        self._lock = threading.Lock()
        self._hostmap = collections.defaultdict(list)  # host -> [connection]
        self._connmap = {}  # map connections to host
        self._readymap = {}  # map connection to ready state

    def add(self, host, connection, ready):
        self._lock.acquire()
        try:
            self._hostmap[host].append(connection)

        if host:
            return list(self._hostmap[host])
        else:
            return dict(self._hostmap)


class KeepAliveHandler(object):
    def __init__(self, timeout=None):
        self._cm = ConnectionManager()
        self._timeout = timeout
        self.requestscount = 0

                h = self._cm.get_ready_conn(host)
            else:
                # no (working) free connections were found.  Create a new one.
                h = http_class(host, timeout=self._timeout)
                if DEBUG:
                    DEBUG.info(
                        "creating new connection to %s (%d)", host, id(h)
                    )
                self._cm.add(host, h, False)
                self._start_transaction(h, req)
                r = h.getresponse()
        # The string form of BadStatusLine is the status line. Add some context
        # to make the error message slightly more useful.
        except httplib.BadStatusLine as err:
            raise urlerr.urlerror(
                _('bad HTTP status line: %s') % pycompat.sysbytes(err.line)
            )
        except (socket.error, httplib.HTTPException) as err:
            raise urlerr.urlerror(err)

        # If not a persistent connection, don't try to reuse it. Look
        # for this using getattr() since vcr doesn't define this

            r = h.getresponse()
            # note: just because we got something back doesn't mean it
            # worked.  We'll check the version below, too.
        except (socket.error, httplib.HTTPException):
            r = None
        except:  # re-raises
            # adding this block just in case we've missed
            # something we will still raise the exception, but
            # lets try and close the connection and remove it
            # first.  We previously got into a nasty loop
            # where an exception was uncaught, and so the
            # connection stayed open.  On the next try, the
            # same exception was raised, etc.  The trade-off is
            # that it's now possible this call will raise
            # a DIFFERENT exception
            if DEBUG:
                DEBUG.error(
                    "unexpected exception - closing " "connection to %s (%d)",
                    host,
                    id(h),
                )
            self._cm.remove(h)
            h.close()
            raise

        if r is None or r.version == 9:
            # httplib falls back to assuming HTTP 0.9 if it gets a
            # bad header back.  This is most likely to happen if
            # the socket has been closed by the server since we
            # last used the connection.
            if DEBUG:
                DEBUG.info(
                    "failed to re-use connection to %s (%d)", host, id(h)
                )
            r = None
        else:
            if DEBUG:
                DEBUG.info("re-using connection to %s (%d)", host, id(h))

                skipheaders[r'skip_' + n.replace(r'-', r'_')] = 1
        try:
            if urllibcompat.hasdata(req):
                data = urllibcompat.getdata(req)
                h.putrequest(
                    req.get_method(),
                    urllibcompat.getselector(req),
                    **skipheaders
                )
                if r'content-type' not in headers:
                    h.putheader(
                        r'Content-type', r'application/x-www-form-urlencoded'
                    )
                if r'content-length' not in headers:
                    h.putheader(r'Content-length', r'%d' % len(data))
            else:
                h.putrequest(
                    req.get_method(),
                    urllibcompat.getselector(req),
                    **skipheaders
                )
        except socket.error as err:
            raise urlerr.urlerror(err)
        for k, v in headers.items():
            h.putheader(k, v)
        h.endheaders()

        self.sentbytescount += getattr(h, 'sentbytescount', 0) - oldbytescount

        try:
            self.parent.requestscount += 1
            self.parent.sentbytescount += (
                getattr(h, 'sentbytescount', 0) - oldbytescount
            )
        except AttributeError:
            pass


class HTTPHandler(KeepAliveHandler, urlreq.httphandler):
    pass


class HTTPResponse(httplib.HTTPResponse):
    # we need to subclass HTTPResponse in order to
    # 1) add readline(), readlines(), and readinto() methods
    # 2) add close_connection() methods

    # the read method wraps the original to accommodate buffering,
    # although read() never adds to the buffer.
    # Both readline and readlines have been stolen with almost no
    # modification from socket.py

    def __init__(self, sock, debuglevel=0, strict=0, method=None):
        extrakw = {}
        if not pycompat.ispy3:
            extrakw[r'strict'] = True
            extrakw[r'buffering'] = True
        httplib.HTTPResponse.__init__(
            self, sock, debuglevel=debuglevel, method=method, **extrakw
        )
        self.fileno = sock.fileno
        self.code = None
        self.receivedbytescount = 0
        self._rbuf = ''
        self._rbufsize = 8096
        self._handler = None  # inserted by the handler later
        self._host = None  # (same)
        self._url = None  # (same)
        self._connection = None  # (same)

    _raw_read = httplib.HTTPResponse.read
    _raw_readinto = getattr(httplib.HTTPResponse, 'readinto', None)

    # Python 2.7 has a single close() which closes the socket handle.

    def close(self):
        if self.fp:
            self.fp.close()
            self.fp = None
            if self._handler:
                self._handler._request_closed(
                    self, self._host, self._connection
                )

    def _close_conn(self):
        self.close()

    def close_connection(self):

        while True:
            if chunk_left is None:
                line = self.fp.readline()
                i = line.find(';')
                if i >= 0:
                    line = line[:i]  # strip chunk-extensions
                try:
                    chunk_left = int(line, 16)
                except ValueError:
                    # close the connection as protocol synchronization is
                    # probably lost

            else:
                parts.append(self._safe_read(chunk_left))
                amt -= chunk_left

            # we read the whole chunk, get another
            self._safe_read(2)  # toss the CRLF at the end of the chunk
            chunk_left = None

        # read and discard trailer up to the CRLF terminator
        ### note: we shouldn't have any trailers!
        while True:

    def readinto(self, dest):
        if self._raw_readinto is None:
            res = self.read(len(dest))
            if not res:
                return 0
            dest[0 : len(res)] = res
            return len(res)
        total = len(dest)
        have = len(self._rbuf)
        if have >= total:
            dest[0:total] = self._rbuf[:total]

        dest[0:have] = self._rbuf
        got += len(self._rbuf)
        self._rbuf = ''
        return got


def safesend(self, str):
    """Send `str' to the server.

    Shamelessly ripped off from httplib to patch a bad behavior.

        else:
            self.sock.sendall(str)
            self.sentbytescount += len(str)
    except socket.error as v:
        reraise = True
        if v.args[0] == errno.EPIPE:  # Broken pipe
            if self._HTTPConnection__state == httplib._CS_REQ_SENT:
                self._broken_pipe_resp = None
                self._broken_pipe_resp = self.getresponse()
                reraise = False
            self.close()
        if reraise:
            raise


def wrapgetresponse(cls):
    """Wraps getresponse in cls with a broken-pipe sane version.
    """

    def safegetresponse(self):
        # In safesend() we might set the _broken_pipe_resp
        # attribute, in which case the socket has already
        # been closed and we just need to give them the response
        # back. Otherwise, we use the normal response path.
        r = getattr(self, '_broken_pipe_resp', None)
        if r is not None:
            return r
        return cls.getresponse(self)

    safegetresponse.__doc__ = cls.getresponse.__doc__
    return safegetresponse


class HTTPConnection(httplib.HTTPConnection):
    # url.httpsconnection inherits from this. So when adding/removing
    # attributes, be sure to audit httpsconnection() for unintended
    # consequences.

    def __init__(self, *args, **kwargs):
        httplib.HTTPConnection.__init__(self, *args, **kwargs)
        self.sentbytescount = 0
        self.receivedbytescount = 0


#########################################################################
#####   TEST FUNCTIONS
#########################################################################

            break
    fo.close()
    m = md5(foo)
    print(format % ('keepalive readline', node.hex(m.digest())))


def comp(N, url):
    print('  making %i connections to:\n  %s' % (N, url))

    procutil.stdout.write('  first using the normal urllib handlers')
    # first use normal opener

    urlreq.installopener(opener)
    t2 = fetch(N, url)
    print('  TIME: %.3f s' % t2)
    print('  improvement factor: %.2f' % (t1 / t2))


def fetch(N, url, delay=0):
    import time

    lens = []
    starttime = time.time()
    for i in range(N):
        if delay and i > 0:
            time.sleep(delay)

        if not i == lens[0]:
            print("WARNING: inconsistent length on read %i: %i" % (j, i))

    return diff


def test_timeout(url):
    global DEBUG
    dbbackup = DEBUG

    class FakeLogger(object):
        def debug(self, msg, *args):
            print(msg % args)

        info = warning = error = debug

    DEBUG = FakeLogger()
    print("  fetching the file to establish a connection")
    fo = urlreq.urlopen(url)
    data1 = fo.read()
    fo.close()

    comp(N, url)
    print('')
    print("performing dropped-connection check")
    test_timeout(url)


if __name__ == '__main__':
    import time

    try:
        N = int(sys.argv[1])
        url = sys.argv[2]
    except (IndexError, ValueError):
        print("%s <integer> <url>" % sys.argv[0])