|
| 1 | +#!/usr/bin/env python |
| 2 | +"""Demonstrate the streaming socket/fd leak from docker/docker-py#2766. |
| 3 | +
|
| 4 | +The streaming helpers in ``docker.api.client`` hand the caller a generator that |
| 5 | +reads from a long-lived socket. Before this change, abandoning that generator |
| 6 | +early (``break``, an exception, or simply dropping the reference) never closed |
| 7 | +the underlying response, so the socket/fd leaked. |
| 8 | +
|
| 9 | +This script reproduces the leak without a Docker daemon, at the raw socket |
| 10 | +level so the result does not depend on connection pooling or garbage-collection |
| 11 | +timing. It serves an endless chunked HTTP response from a local thread, then |
| 12 | +repeatedly: |
| 13 | +
|
| 14 | + 1. opens a streaming connection, |
| 15 | + 2. reads a single chunk, |
| 16 | + 3. stops the iterator early (``generator.close()``), |
| 17 | +
|
| 18 | +while holding every connection open for the whole run. |
| 19 | +
|
| 20 | +``leaky`` is a generator with no cleanup (the pre-fix behaviour). ``fixed`` |
| 21 | +wraps the same loop in ``try/finally: connection.close()`` -- the analogue of |
| 22 | +``response.close()`` that ``APIClient._stream_raw_result`` now performs. Each |
| 23 | +``connection.close()`` here closes exactly one socket, the same way |
| 24 | +``requests.Response.close()`` releases the docker daemon socket. |
| 25 | +
|
| 26 | +Usage: |
| 27 | + python benchmarks/stream_leak.py [--iterations N] |
| 28 | +""" |
| 29 | +import argparse |
| 30 | +import http.client |
| 31 | +import threading |
| 32 | +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| 33 | + |
| 34 | +import psutil |
| 35 | + |
| 36 | + |
| 37 | +class _StreamingHandler(BaseHTTPRequestHandler): |
| 38 | + protocol_version = 'HTTP/1.1' # keep-alive: socket stays open until closed |
| 39 | + |
| 40 | + # Stream chunks forever so the connection stays open until the client |
| 41 | + # closes it -- the follow=True case. |
| 42 | + def do_GET(self): |
| 43 | + self.send_response(200) |
| 44 | + self.send_header('Content-Type', 'application/octet-stream') |
| 45 | + self.send_header('Transfer-Encoding', 'chunked') |
| 46 | + self.end_headers() |
| 47 | + try: |
| 48 | + while True: |
| 49 | + payload = b'log line\n' |
| 50 | + self.wfile.write( |
| 51 | + f'{len(payload):x}\r\n'.encode() + payload + b'\r\n') |
| 52 | + self.wfile.flush() |
| 53 | + except (BrokenPipeError, ConnectionResetError, OSError): |
| 54 | + pass # client went away -- the whole point of the benchmark |
| 55 | + |
| 56 | + def handle(self): |
| 57 | + try: |
| 58 | + super().handle() |
| 59 | + except (ConnectionError, OSError): |
| 60 | + pass # client closed mid-stream -- expected here |
| 61 | + |
| 62 | + def log_message(self, *args): |
| 63 | + pass # keep the benchmark output clean |
| 64 | + |
| 65 | + |
| 66 | +def _start_server(): |
| 67 | + server = ThreadingHTTPServer(('127.0.0.1', 0), _StreamingHandler) |
| 68 | + thread = threading.Thread(target=server.serve_forever, daemon=True) |
| 69 | + thread.start() |
| 70 | + return server |
| 71 | + |
| 72 | + |
| 73 | +class Stream: |
| 74 | + """Minimal stand-in for a streaming docker response over a raw socket. |
| 75 | +
|
| 76 | + ``close()`` releases the socket, mirroring ``requests.Response.close()``. |
| 77 | + """ |
| 78 | + |
| 79 | + def __init__(self, host, port): |
| 80 | + self._conn = http.client.HTTPConnection(host, port) |
| 81 | + self._conn.request('GET', '/') |
| 82 | + self._resp = self._conn.getresponse() |
| 83 | + |
| 84 | + def read(self, n=16): |
| 85 | + return self._resp.read(n) |
| 86 | + |
| 87 | + def close(self): |
| 88 | + # Closing the response releases the socket -- the same call |
| 89 | + # APIClient now makes in its streaming generators' finally block. |
| 90 | + self._resp.close() |
| 91 | + self._conn.close() |
| 92 | + |
| 93 | + @property |
| 94 | + def open(self): |
| 95 | + # http.client moves the socket into the response object, so check the |
| 96 | + # response rather than conn.sock. close() sets isclosed() True. |
| 97 | + return not self._resp.isclosed() |
| 98 | + |
| 99 | + |
| 100 | +def leaky(stream): |
| 101 | + """Pre-fix behaviour: yields chunks, never closes the socket.""" |
| 102 | + while True: |
| 103 | + data = stream.read() |
| 104 | + if not data: |
| 105 | + break |
| 106 | + yield data |
| 107 | + |
| 108 | + |
| 109 | +def fixed(stream): |
| 110 | + """Post-fix behaviour, mirroring APIClient._stream_raw_result.""" |
| 111 | + try: |
| 112 | + while True: |
| 113 | + data = stream.read() |
| 114 | + if not data: |
| 115 | + break |
| 116 | + yield data |
| 117 | + finally: |
| 118 | + stream.close() |
| 119 | + |
| 120 | + |
| 121 | +def established_to(proc, port): |
| 122 | + try: |
| 123 | + return sum( |
| 124 | + 1 for c in proc.net_connections(kind='tcp') |
| 125 | + if c.raddr and c.raddr.port == port and c.status == 'ESTABLISHED' |
| 126 | + ) |
| 127 | + except (psutil.AccessDenied, NotImplementedError): |
| 128 | + return -1 |
| 129 | + |
| 130 | + |
| 131 | +def run(make_stream, host, port, iterations, proc): |
| 132 | + streams = [] |
| 133 | + generators = [] |
| 134 | + for _ in range(iterations): |
| 135 | + stream = Stream(host, port) |
| 136 | + gen = make_stream(stream) |
| 137 | + next(gen) # read a single chunk |
| 138 | + gen.close() # consumer stops early -> GeneratorExit |
| 139 | + streams.append(stream) |
| 140 | + generators.append(gen) |
| 141 | + |
| 142 | + leaked = sum(1 for s in streams if s.open) |
| 143 | + established = established_to(proc, port) |
| 144 | + |
| 145 | + for s in streams: # tidy up before the next run |
| 146 | + s.close() |
| 147 | + return leaked, established |
| 148 | + |
| 149 | + |
| 150 | +def main(): |
| 151 | + parser = argparse.ArgumentParser(description=__doc__) |
| 152 | + parser.add_argument('--iterations', type=int, default=200, |
| 153 | + help='streams opened and abandoned per run') |
| 154 | + args = parser.parse_args() |
| 155 | + |
| 156 | + server = _start_server() |
| 157 | + host, port = server.server_address |
| 158 | + proc = psutil.Process() |
| 159 | + |
| 160 | + print(f'opening {args.iterations} streams, reading one chunk, then ' |
| 161 | + f'stopping each early\n') |
| 162 | + header = (f'{"impl":<8}{"streams":>10}{"sockets leaked":>16}' |
| 163 | + f'{"ESTABLISHED conns":>20}') |
| 164 | + print(header) |
| 165 | + print('-' * len(header)) |
| 166 | + |
| 167 | + results = {} |
| 168 | + for name, fn in (('old', leaky), ('fixed', fixed)): |
| 169 | + leaked, established = run(fn, host, port, args.iterations, proc) |
| 170 | + results[name] = leaked |
| 171 | + print(f'{name:<8}{args.iterations:>10}{leaked:>16}{established:>20}') |
| 172 | + |
| 173 | + server.shutdown() |
| 174 | + print() |
| 175 | + if results['old'] == args.iterations and results['fixed'] == 0: |
| 176 | + print(f'PASS: old leaks all {args.iterations} sockets on early stop; ' |
| 177 | + f'fixed closes every one.') |
| 178 | + else: |
| 179 | + print('NOTE: compare the "sockets leaked" column for the two impls.') |
| 180 | + |
| 181 | + |
| 182 | +if __name__ == '__main__': |
| 183 | + main() |
0 commit comments