net: simplify backend transport/protocol interface; fixup test futures

kafka/net/asyncio_backend.py

@@ -329,99 +329,64 @@ async def create_connection(self, protocol, host, port, *, ssl=None,
                 'The asyncio backend does not support proxy_url yet; use the '
                 'default selector backend for SOCKS5/HTTP-CONNECT proxying.')
         server_hostname = host.rstrip('.') if ssl is not None else None
-        adapter = _AsyncioProtocolAdapter()
+        adapter = _AsyncioProtocolAdapter(protocol, host, port, socket_options)
         connect = self._loop.create_connection(
             lambda: adapter, host, port, ssl=ssl, server_hostname=server_hostname)
         try:
             if timeout_at is not None:
                 connect = asyncio.wait_for(connect, max(0.0, timeout_at - time.monotonic()))
-            aio_transport, _ = await connect
+            await connect
+            if adapter.error is not None:
+                raise adapter.error
         except asyncio.TimeoutError:
             raise Errors.KafkaConnectionError('Connection timed out')
         except Errors.KafkaError:
             raise
         except Exception as exc:  # noqa: BLE001  -- surface any connect error uniformly
             raise Errors.KafkaConnectionError('unable to connect to %s:%s: %s' % (host, port, exc))
+
+
+class _AsyncioProtocolAdapter(asyncio.Protocol):
+    """Thin asyncio.Protocol that wires a KafkaConnection to a wrapped transport."""
+
+    def __init__(self, conn, host, port, socket_options=()):
+        self._conn = conn
+        self._host = host
+        self._port = port
+        self._socket_options = socket_options
+        self.error = None
+        self.transport = None       # the _AsyncioTransport wrapper
+
+    def connection_made(self, aio_transport):
         sock = aio_transport.get_extra_info('socket')
         if sock is not None:
-            for option in socket_options:
+            for option in self._socket_options:
                 try:
                     sock.setsockopt(*option)
                 except OSError:
                     pass
-        return _AsyncioTransport(aio_transport, adapter, host, port)
-
-
-class _AsyncioProtocolAdapter(asyncio.Protocol):
-    """Bridges asyncio's Protocol callbacks to a KafkaConnection.
-
-    asyncio calls ``connection_made`` on this adapter during
-    ``create_connection`` -- before the manager wires the KafkaConnection via
-    ``transport.set_protocol(conn)`` (Option A: the caller runs connection_made
-    after its "closed during connect" check). Reading is paused until wired, so
-    no data is delivered early; a buffer/latched-loss is kept as a safety net.
-    """
-
-    def __init__(self):
-        self._kafka = None          # the KafkaConnection, once wired
-        self._transport = None
-        self._buffer = []
-        self._eof = False
-        self._lost = False
-        self._lost_exc = None
-        self._paused_writing = False
-        self._on_read = None        # bumps the wrapper's last_read
-
-    def connection_made(self, transport):
-        self._transport = transport
-        # Hold off delivery until the KafkaConnection is wired + resumes reading.
-        transport.pause_reading()
+        self.transport = _AsyncioTransport(aio_transport, self._host, self._port)
+        try:
+            self._conn.connection_made(self.transport)
+        except Exception as exc:  # noqa: BLE001  -- conn refused (closed mid-connect)
+            self.error = exc
+            aio_transport.abort()
 
     def data_received(self, data):
-        if self._on_read is not None:
-            self._on_read()
-        if self._kafka is None:
-            self._buffer.append(data)
-        else:
-            self._kafka.data_received(data)
+        self.transport._bump_read()
+        self._conn.data_received(data)
 
     def eof_received(self):
-        if self._kafka is not None:
-            return self._kafka.eof_received()
-        self._eof = True
-        return None
+        return self._conn.eof_received()
 
     def connection_lost(self, exc):
-        if self._kafka is not None:
-            self._kafka.connection_lost(exc)
-        else:
-            self._lost = True
-            self._lost_exc = exc
+        self._conn.connection_lost(exc)
 
     def pause_writing(self):
-        if self._kafka is not None:
-            self._kafka.pause_writing()
-        else:
-            self._paused_writing = True
+        self._conn.pause_writing()
 
     def resume_writing(self):
-        if self._kafka is not None:
-            self._kafka.resume_writing()
-        else:
-            self._paused_writing = False
-
-    def _wire(self, kafka_protocol, on_read):
-        self._kafka = kafka_protocol
-        self._on_read = on_read
-        if self._paused_writing:
-            kafka_protocol.pause_writing()
-        for data in self._buffer:
-            kafka_protocol.data_received(data)
-        self._buffer = []
-        if self._eof:
-            kafka_protocol.eof_received()
-        if self._lost:
-            kafka_protocol.connection_lost(self._lost_exc)
+        self._conn.resume_writing()
 
 
 class _AsyncioTransport:
@@ -433,9 +398,8 @@ class _AsyncioTransport:
     protocol, host/host_port/getPeer, and last_activity for idle sweeping.
     """
 
-    def __init__(self, transport, adapter, host, port):
+    def __init__(self, transport, host, port):
         self._t = transport
-        self._adapter = adapter
         self.host = host
         self._port = port
         self._protocol = None
@@ -454,11 +418,11 @@ def get_protocol(self):
 
     def set_protocol(self, protocol):
         self._protocol = protocol
-        self._adapter._wire(protocol, self._bump_read)
 
     def write(self, data):
         self.last_write = time.monotonic()
         self._t.write(data)
+        return len(data)
 
     def close(self):
         self._t.close()

kafka/net/backend.py

@@ -108,10 +108,15 @@ class Transport(Protocol):
     """The transport surface a backend's ``create_connection`` returns.
 
     The subset of the ``asyncio.Transport`` / Twisted ``ITransport`` surface
-    that ``KafkaConnection`` actually drives. The selector's
-    ``KafkaTCPTransport`` and an asyncio-transport adapter both satisfy it.
+    that ``KafkaConnection`` actually drives (plus ``last_activity``, which the
+    manager reads to sweep idle connections). A backend's ``create_connection``
+    builds one and wires it to the conn. The selector's ``KafkaTCPTransport``
+    and an asyncio-transport adapter both satisfy it.
     """
 
+    # Monotonic timestamp of the last read/write; manager idle-sweeping reads it.
+    last_activity: float
+
     def write(self, data: bytes) -> None: ...
     def close(self) -> None: ...
     def abort(self, error: Any = None) -> None: ...
@@ -177,17 +182,20 @@ async def create_connection(
         proxy_url: Optional[str] = None,
         socket_options: Sequence[Any] = (),
         timeout_at: Optional[float] = None,
-    ) -> Transport:
-        """Establish and return a connected :class:`Transport` to ``host:port``.
-
-        The backend owns DNS, connect, TLS and (where supported) proxying. The
-        *caller* wires the transport to the protocol afterwards via
-        ``protocol.connection_made(transport)`` (so manager-level policy such as
-        the "closed during connect" check runs first). ``protocol`` (a
-        ``KafkaConnection``) is passed because backends that own the socket
-        (asyncio/Twisted) need it at connect time to receive transport events;
-        they buffer inbound data until ``connection_made`` is called. Backends
-        without native proxy support raise when ``proxy_url`` is set.
+    ) -> None:
+        """Establish a connected :class:`Transport` to ``host:port`` and wire it.
+
+        The backend owns DNS, connect, TLS and (where supported) proxying, builds
+        a :class:`Transport`, and wires it to ``protocol`` by calling
+        ``protocol.connection_made(transport)`` itself -- mirroring
+        ``asyncio.loop.create_connection`` / Twisted, which own the socket and
+        wire the protocol at connect time. Nothing is returned: the caller drives
+        the connection through ``protocol`` (``conn.transport``), never a
+        transport handle. ``protocol`` (a ``KafkaConnection``) may *refuse* the
+        transport by raising from ``connection_made`` if it closed mid-connect;
+        on that (or any) failure the backend closes the orphaned transport
+        before propagating. Backends without native proxy support raise when
+        ``proxy_url`` is set.
         """
 
     # --- cross-thread bridge ---------------------------------------------

kafka/net/connection.py

@@ -168,7 +168,9 @@ def _send_request(self, request, future=None, timeout_at=None):
         # in_flight_requests (len==1), trip the >= check, pause, and never be
         # written to the transport - hanging forever.
         if not self.paused:
-            self.transport.write(self.parser.send_bytes())
+            total_bytes = self.transport.write(self.parser.send_bytes())
+            if self._sensors:
+                self._sensors.bytes_sent.record(total_bytes)
         if len(self.in_flight_requests) >= self.config['max_in_flight_requests_per_connection']:
             self.pause('max_in_flight')
         return future
@@ -193,6 +195,8 @@ def data_received(self, data):
         if self.closed:
             log.debug('%s: Ignoring %d bytes received by closed connection', self, len(data))
             return
+        if self._sensors:
+            self._sensors.bytes_received.record(len(data))
         responses = self.parser.receive_bytes(data)
 
         # augment responses w/ correlation_id, future, and timestamp
@@ -301,7 +305,9 @@ def unpause(self, v):
             if not self.paused and self.parser and self.transport:
                 to_send = self.parser.send_bytes()
                 if to_send:
-                    self.transport.write(to_send)
+                    total_bytes = self.transport.write(to_send)
+                    if self._sensors:
+                        self._sensors.bytes_sent.record(total_bytes)
 
     def pause_writing(self):
         """ Called when the transport's buffer goes over the high-water mark.

kafka/net/manager.py

@@ -236,41 +236,31 @@ def ssl_enabled(self):
         return self.config['security_protocol'] in ('SSL', 'SASL_SSL')
 
     async def _connect(self, node, conn, reset_backoff_on_connect=True, timeout_at=None):
-        # Tracks ownership of the freshly built transport: while non-None it is
-        # ours to clean up (the connection hasn't taken it over yet), so the
-        # finally clause closes it. Cleared once connection_made() succeeds.
-        transport = None
         try:
-            transport = await self._net.create_connection(
+            await self._net.create_connection(
                 conn, node.host, node.port,
                 ssl=self.ssl_context,
                 proxy_url=self.config['proxy_url'],
                 socket_options=self.config['socket_options'],
                 timeout_at=timeout_at)
-            # The connection (or the whole manager) may have been closed while
-            # we were building the transport. Handing it to connection_made()
-            # would flip the conn back to `initializing` and resurrect a
-            # connection that is already being torn down. Discard
-            # the new transport instead of reviving a dead connection.
-            if conn.closed or self.closed:
-                log.debug('%s: closed during connect; discarding new transport', conn)
-                return
-            conn.connection_made(transport)
-            transport = None  # conn owns cleanup now; skip finally: transport.close()
             # Note: conn.initialize does not currently raise on error;
             # errors are pushed to conn.init_future and raised on await conn
             await conn.initialize(timeout_at=timeout_at)
         except Exception as exc:
+            if conn.closed or self.closed:
+                # A concurrent close() raced the connect (manager / bootstrap
+                # teardown). connection_made() refused to resurrect the conn and
+                # the backend already discarded the transport; don't back off a
+                # connection that is going away.
+                log.debug('%s: closed during connect; discarding', conn)
+                return
             log.error('Connection failed: %s', exc)
             conn.connection_lost(exc)
             self.update_backoff(node.node_id)
             if isinstance(exc, (Errors.SaslAuthenticationFailedError,
                                 Errors.AuthorizationError)):
                 self._auth_failures[node.node_id] = exc
             return
-        finally:
-            if transport is not None:
-                transport.close()
 
         if self._sensors:
             self._sensors.connection_created.record()

kafka/net/selector.py

@@ -490,15 +490,16 @@ def create_future(self):
     async def create_connection(self, protocol, host, port, *, ssl=None,
                                 proxy_url=None, socket_options=(),
                                 timeout_at=None):
-        """Establish and return a connected transport to host:port.
+        """Establish a connected transport to host:port and wire ``protocol``.
 
         The selector owns the raw socket: DNS + non-blocking connect (with
         optional SOCKS5/HTTP-CONNECT proxy via KafkaNetSocket), then wraps it
-        in a TCP or SSL transport and runs the TLS handshake. ``protocol`` (the
-        KafkaConnection) is not used here -- the caller wires it via
-        ``connection_made()`` after its own "closed during connect" check; it
-        is part of the contract because socket-owning backends (asyncio,
-        Twisted) need it at connect time.
+        in a TCP or SSL transport, runs the TLS handshake, and calls
+        ``protocol.connection_made(transport)`` -- mirroring asyncio/Twisted,
+        which own the socket and wire the protocol at connect time. On any
+        failure (handshake error, or a ``protocol`` that refuses the transport
+        because it closed mid-connect) the transport is closed before raising,
+        so the caller never handles a transport instance directly.
         """
         sock = await _inet_create_connection(self, host, port, socket_options,
                                              proxy_url=proxy_url, timeout_at=timeout_at)
@@ -509,8 +510,13 @@ async def create_connection(self, protocol, host, port, *, ssl=None,
         try:
             await transport.handshake()
         except Exception as e:
+            transport.close()
             raise Errors.KafkaConnectionError('Handshake failed: %s' % e)
-        return transport
+        try:
+            protocol.connection_made(transport)
+        except Exception:
+            transport.close()
+            raise
 
     def sleep(self, delay):
         return KernelEvent('_sleep', delay)

kafka/net/transport.py

@@ -95,8 +95,6 @@ async def _read_from_sock(self):
                 return self.abort(error=err)
             log.debug('%s: received %d bytes', self, len(recvd_data))
             self.last_read = time.monotonic()
-            if self._protocol and self._protocol._sensors:
-                self._protocol._sensors.bytes_received.record(len(recvd_data))
             try:
                 self._protocol.data_received(recvd_data)
             except Errors.KafkaProtocolError as e:
@@ -177,6 +175,7 @@ def write(self, data):
         if not self._writing:
             self._writing = True
             self._write_task =  self._net.call_soon(self._write_to_sock)
+        return len(data)
 
     def writelines(self, list_of_data):
         """Write a list (or any iterable) of data bytes to the transport."""
@@ -186,6 +185,7 @@ def writelines(self, list_of_data):
         if not self._writing:
             self._writing = True
             self._write_task = self._net.call_soon(self._write_to_sock)
+        return sum(len(data) for data in list_of_data)
 
     async def _write_to_sock(self):
         try:
@@ -196,8 +196,6 @@ async def _write_to_sock(self):
                     return self.abort(error=err)
                 log.debug('%s: sent %d bytes', self, total_bytes)
                 self.last_write = time.monotonic()
-                if self._protocol and self._protocol._sensors:
-                    self._protocol._sensors.bytes_sent.record(total_bytes)
         finally:
             self._writing = False
         if self._closed:

test/conftest.py

@@ -58,18 +58,22 @@ def client(net, manager, broker):
 
 @pytest.fixture
 def net():
-    return NetworkSelector()
+    backend = NetworkSelector()
+    try:
+        yield backend
+    finally:
+        backend.close()
 
 
 @pytest.fixture
 def manager(net, broker):
+    broker.attach(net)
     manager = KafkaConnectionManager(
         net,
         bootstrap_servers='%s:%d' % (broker.host, broker.port),
         api_version=broker.broker_version,
         request_timeout_ms=5000,
     )
-    broker.attach(manager)
     try:
         yield manager
     finally: