Do not leak the SSL upgrade future when connect() is cancelled early

asyncpg/connect_utils.py

@@ -993,10 +993,26 @@ async def _create_ssl_connection(
     retry: bool = False,
 ) -> typing.Tuple[asyncio.Transport, _ProctolFactoryR]:
 
-    tr, pr = await loop.create_connection(
-        lambda: TLSUpgradeProto(loop, host, port,
-                                ssl_context, ssl_is_advisory),
-        host, port)
+    upgrade_proto: typing.Optional[TLSUpgradeProto] = None
+
+    def upgrade_protocol_factory() -> TLSUpgradeProto:
+        nonlocal upgrade_proto
+        upgrade_proto = TLSUpgradeProto(
+            loop, host, port, ssl_context, ssl_is_advisory)
+        return upgrade_proto
+
+    try:
+        tr, pr = await loop.create_connection(
+            upgrade_protocol_factory, host, port)
+    except (Exception, asyncio.CancelledError):
+        # The protocol can exist before create_connection() returns.  If
+        # that operation fails or is cancelled, nobody will await on_data.
+        if upgrade_proto is not None:
+            if not upgrade_proto.on_data.done():
+                upgrade_proto.on_data.cancel()
+            elif not upgrade_proto.on_data.cancelled():
+                upgrade_proto.on_data.exception()
+        raise
 
     tr.write(struct.pack('!ll', 8, _SSL_REQUEST_CODE))  # SSLRequest message.
 

tests/test_connect.py

@@ -1761,8 +1761,12 @@ async def close():
 class TestConnection(tb.ConnectedTestCase):
 
     async def test_connection_cancelled_during_transport_setup(self):
-        for connection_lost_first in (False, True):
-            with self.subTest(connection_lost_first=connection_lost_first):
+        for sslmode, connection_lost_first in itertools.product(
+            (False, 'prefer'), (False, True),
+        ):
+            with self.subTest(
+                sslmode=sslmode, connection_lost_first=connection_lost_first,
+            ):
                 async def interrupted_connector(factory, *args, **kwargs):
                     proto = factory()
                     if connection_lost_first:
@@ -1775,10 +1779,11 @@ async def interrupted_connector(factory, *args, **kwargs):
                     self.loop, 'create_connection', interrupted_connector,
                 ):
                     with self.assertRaises(asyncio.CancelledError):
-                        await self.connect(host='127.0.0.1', ssl=False)
+                        await self.connect(host='127.0.0.1', ssl=sslmode)
 
-                # Authentication must not report an unobserved failure after
-                # cancellation, even if the transport closes afterwards.
+                # Neither authentication nor the SSL upgrade must report an
+                # unobserved failure after cancellation, even if the
+                # transport closes afterwards.
                 await asyncio.sleep(0)
                 gc.collect()