handle timeouted h2 connections, refill reply budget

proxy/cf_h2.py

@@ -98,10 +98,11 @@ class _ReplayPacket:
     pending: Optional[asyncio.Task] = None
     originals: Set[asyncio.Task] = field(default_factory=set, repr=False)
     retired: bool = False
+    reply_id: Optional[bytes] = None
 
     @property
     def replayed(self) -> bool:
-        return self.attempts > 0
+        return self.last_attempt_at > 0 or self.attempts > 0
 
 
 async def _read_packet(reader, decryptor, tag: bytes) -> Tuple[bytes, bool]:
@@ -466,6 +467,9 @@ class _HttpChannel:
             reply = await self.lane._post(body, self.channel_id, replay=replay, channel=self)
             self.last_progress = time.monotonic()
             if reply and not self.closed:
+                if packet is not None and reply[:24] != packet.reply_id:
+                    packet.reply_id = reply[:24]
+                    packet.attempts = 0
                 if self.queue.full() or self.reply_bytes + len(reply) > MAX_CHANNEL_BYTES:
                     raise BufferError('H2 client response queue exceeded bound')
                 self.reply_bytes += len(reply)

proxy/h2_transport.py

@@ -49,7 +49,14 @@ class _ResponseStream(httpx.AsyncByteStream):
         try:
             await asyncio.wait_for(self.changed.wait(), self.read_timeout)
         except asyncio.TimeoutError as exc:
-            raise httpx.ReadTimeout('Timed out waiting for H2 stream %d' % self.stream_id) from exc
+            error = httpx.ReadTimeout('Timed out waiting for H2 stream %d' % self.stream_id)
+            # An idle socket can remain established after the path stops
+            # delivering data. Reconnect once its last active request times
+            # out, while preserving unrelated requests and buffered replies.
+            if not any(stream is not self and not stream.ended
+                       for stream in self.connection.streams.values()):
+                self.connection.fail(error)
+            raise error from exc
 
     def check_error(self):
         if self.error is not None:

tests/test_cf_h2.py

@@ -820,6 +820,68 @@ class HttpMultiplexTest(unittest.IsolatedAsyncioTestCase):
                 recovery.cancel()
                 await asyncio.gather(recovery, return_exceptions=True)
 
+    async def test_new_replies_refresh_idle_receiver_budget_without_native_upload(self):
+        replies = [b'samekey!' + bytes([index]) * 32 for index in range(6)]
+        calls = []
+        exhausted = asyncio.Event()
+
+        def handler(request):
+            calls.append(request.content)
+            count = len(calls)
+            if count == len(replies) + 3:
+                exhausted.set()
+            return httpx.Response(200, content=replies[count - 1] if count <= len(replies) else b'',
+                                  extensions={'http_version': b'HTTP/2'})
+
+        self._transport(handler)
+        channel = self._channel()
+        body = b'samekey!' + b'a' * 32
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .002), \
+                patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .01):
+            await channel.send(body, False)
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                # The native client is waiting for a batch of replies and
+                # sends no new packet between them. Three successful polls
+                # must not use up the budget while more responses arrive.
+                for reply in replies:
+                    self.assertEqual(await asyncio.wait_for(channel.receive(), 1), reply)
+                await asyncio.wait_for(exhausted.wait(), 1)
+                await asyncio.sleep(.05)
+                self.assertEqual(calls, [body] * (len(replies) + 3))
+                self.assertFalse(channel.pending)
+                self.assertFalse(channel.replay_pending)
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
+    async def test_duplicate_reply_does_not_refresh_idle_receiver_budget(self):
+        calls = []
+        reply = b'r' * 40
+
+        def handler(request):
+            calls.append(request.content)
+            return httpx.Response(200, content=reply, extensions={'http_version': b'HTTP/2'})
+
+        self._transport(handler)
+        channel = self._channel()
+        with patch('proxy.cf_h2.REPLAY_IDLE_SECONDS', .005), \
+                patch('proxy.cf_h2.REPLAY_CHECK_SECONDS', .002), \
+                patch('proxy.cf_h2.REPLAY_RETRY_SECONDS', .01):
+            await channel.send(b'samekey!' + b'a' * 32, False)
+            recovery = asyncio.create_task(channel._recover())
+            try:
+                for _ in range(4):
+                    self.assertEqual(await asyncio.wait_for(channel.receive(), 1), reply)
+                await asyncio.sleep(.05)
+                self.assertEqual(len(calls), 4)
+                self.assertFalse(channel.pending)
+                self.assertFalse(channel.replay_pending)
+            finally:
+                recovery.cancel()
+                await asyncio.gather(recovery, return_exceptions=True)
+
     async def test_old_pending_request_has_bounded_retries_after_newer_requests(self):
         body = b'samekey!' + b't' * 32
         original_started = asyncio.Event()

tests/test_h2_transport.py

@@ -1030,13 +1030,31 @@ class H2WireTest(unittest.IsolatedAsyncioTestCase):
     async def test_read_timeout_resets_only_its_stream(self):
         self.lane.client.timeout = httpx.Timeout(2, read=.15)
         stalled, request = await self.post()
+        self.lane.client.timeout = httpx.Timeout(2, read=1)
         healthy, healthy_request = await self.post()
-        self.respond(healthy_request)
-        self.assertEqual(await healthy, b'r' * 40)
         with self.assertRaises(httpx.ReadTimeout):
             await stalled
         self.assertEqual(await asyncio.wait_for(self.resets.get(), .5), (request[1], 8))
         self.assertIsNone(self.transport.connection.error)
+        self.respond(healthy_request)
+        self.assertEqual(await healthy, b'r' * 40)
+        self.assertEqual(len(self.connections), 1)
+
+    async def test_last_stream_timeout_reconnects_instead_of_reusing_silent_connection(self):
+        self.lane.client.timeout = httpx.Timeout(2, read=.05)
+        stalled, _ = await self.post()
+        old_connection = self.transport.connection
+        with self.assertRaises(httpx.ReadTimeout):
+            await stalled
+        self.assertIsNotNone(old_connection.error)
+        self.assertTrue(old_connection.writer.is_closing())
+        self.assertFalse(old_connection.streams)
+        self.lane.client.timeout = httpx.Timeout(2, read=1)
+        recovered, request = await self.post()
+        self.respond(request)
+        self.assertEqual(await recovered, b'r' * 40)
+        self.assertIsNot(self.transport.connection, old_connection)
+        self.assertEqual(len(self.connections), 2)
 
     async def test_peer_reset_does_not_break_other_streams(self):
         bad, request = await self.post()