fix(consumer): stop rebuilding client on transient get_records errors (#73) (#93)

* fix: stop rebuilding consumer client on transient fetch errors

* test: dedupe metrics sleep patching

Author: hampsterx

kinesis/consumer.py

@@ -518,8 +518,8 @@ async def get_records(self, shard):
                     1,
                     {"stream_name": self.stream_name, "shard_id": shard["ShardId"], "error_type": "connection"},
                 )
-                log.warning("Connection error {}. rebuilding client..".format(e))
-                await self.get_conn()
+                log.warning("Connection error {}. sleeping..".format(e))
+                await asyncio.sleep(3)
             except (TimeoutError, ReadTimeoutError) as e:
                 self.metrics.increment(
                     MetricType.CONSUMER_ERRORS,
@@ -571,8 +571,9 @@ async def get_records(self, shard):
                         shard.pop("ShardIterator", None)
 
                 elif code == "InternalFailure":
-                    log.warning("Received InternalFailure from Kinesis, rebuilding connection.. ")
-                    await self.get_conn()
+                    message = e.response.get("Error", {}).get("Message", str(e))
+                    log.warning("Received InternalFailure from Kinesis: {}. sleeping..".format(message))
+                    await asyncio.sleep(3)
 
                 else:
                     log.warning("ClientError {}. sleeping..".format(code))

tests/test_metrics_unit.py

@@ -28,6 +28,18 @@
 from kinesis.utils import Throttler
 
 
+@pytest.fixture
+def fast_asyncio_sleep(monkeypatch):
+    """Replace asyncio.sleep with a zero-delay awaitable for retry-path tests."""
+    orig_sleep = asyncio.sleep
+
+    async def fast_sleep(_):
+        await orig_sleep(0)
+
+    monkeypatch.setattr(asyncio, "sleep", fast_sleep)
+    return fast_sleep
+
+
 class TestNoOpMetricsCollector:
     def test_all_operations_are_noop(self):
         """Test that NoOpMetricsCollector does nothing."""
@@ -501,7 +513,7 @@ def _shard(shard_id="shard-0"):
         }
 
     @pytest.mark.asyncio
-    async def test_connection_error_emits_connection_label(self, mock_consumer):
+    async def test_connection_error_emits_connection_label(self, mock_consumer, fast_asyncio_sleep):
         collector = InMemoryMetricsCollector()
         consumer = mock_consumer(metrics_collector=collector)
         consumer.client = MagicMock()
@@ -512,9 +524,10 @@ async def test_connection_error_emits_connection_label(self, mock_consumer):
 
         key = f"consumer_errors_total{{error_type=connection,{_shard_labels('shard-0')}}}"
         assert collector.counters[key] == 1
+        consumer.get_conn.assert_not_awaited()
 
     @pytest.mark.asyncio
-    async def test_botocore_connection_closed_emits_connection_label(self, mock_consumer):
+    async def test_botocore_connection_closed_emits_connection_label(self, mock_consumer, fast_asyncio_sleep):
         """Real AWS commonly raises botocore ConnectionClosedError when an idle
         keep-alive connection is dropped mid-response. Prior to #73 fix this fell
         through to the catch-all and got mislabeled as error_type=unknown."""
@@ -530,12 +543,12 @@ async def test_botocore_connection_closed_emits_connection_label(self, mock_cons
 
         key = f"consumer_errors_total{{error_type=connection,{_shard_labels('shard-0')}}}"
         assert collector.counters[key] == 1
-        # Must rebuild client, not just sleep — same remediation as ClientConnectionError
-        consumer.get_conn.assert_awaited_once()
+        consumer.get_conn.assert_not_awaited()
 
     @pytest.mark.asyncio
-    async def test_botocore_endpoint_connection_error_emits_connection_label(self, mock_consumer):
-        """botocore.EndpointConnectionError (DNS/TCP failure before request is sent)."""
+    async def test_botocore_endpoint_connection_error_emits_connection_label(self, mock_consumer, fast_asyncio_sleep):
+        """botocore.EndpointConnectionError (DNS/TCP failure before request is sent)
+        should be counted as a connection error and return without rebuilding."""
         collector = InMemoryMetricsCollector()
         consumer = mock_consumer(metrics_collector=collector)
         consumer.client = MagicMock()
@@ -548,10 +561,10 @@ async def test_botocore_endpoint_connection_error_emits_connection_label(self, m
 
         key = f"consumer_errors_total{{error_type=connection,{_shard_labels('shard-0')}}}"
         assert collector.counters[key] == 1
-        consumer.get_conn.assert_awaited_once()
+        consumer.get_conn.assert_not_awaited()
 
     @pytest.mark.asyncio
-    async def test_botocore_read_timeout_emits_timeout_label(self, mock_consumer):
+    async def test_botocore_read_timeout_emits_timeout_label(self, mock_consumer, fast_asyncio_sleep):
         """botocore.ReadTimeoutError lands in the timeout bucket, not unknown."""
         collector = InMemoryMetricsCollector()
         consumer = mock_consumer(metrics_collector=collector)
@@ -560,40 +573,25 @@ async def test_botocore_read_timeout_emits_timeout_label(self, mock_consumer):
             side_effect=ReadTimeoutError(endpoint_url="https://kinesis.us-east-1.amazonaws.com/")
         )
 
-        orig_sleep = asyncio.sleep
-
-        async def fast_sleep(_):
-            await orig_sleep(0)
-
-        with pytest.MonkeyPatch().context() as mp:
-            mp.setattr(asyncio, "sleep", fast_sleep)
-            await consumer.get_records(self._shard())
+        await consumer.get_records(self._shard())
 
         key = f"consumer_errors_total{{error_type=timeout,{_shard_labels('shard-0')}}}"
         assert collector.counters[key] == 1
 
     @pytest.mark.asyncio
-    async def test_timeout_emits_timeout_label(self, mock_consumer):
+    async def test_timeout_emits_timeout_label(self, mock_consumer, fast_asyncio_sleep):
         collector = InMemoryMetricsCollector()
         consumer = mock_consumer(metrics_collector=collector)
         consumer.client = MagicMock()
         consumer.client.get_records = AsyncMock(side_effect=asyncio.TimeoutError())
 
-        # Monkeypatch asyncio.sleep inside the except block to avoid 3s wait.
-        orig_sleep = asyncio.sleep
-
-        async def fast_sleep(_):
-            await orig_sleep(0)
-
-        with pytest.MonkeyPatch().context() as mp:
-            mp.setattr(asyncio, "sleep", fast_sleep)
-            await consumer.get_records(self._shard())
+        await consumer.get_records(self._shard())
 
         key = f"consumer_errors_total{{error_type=timeout,{_shard_labels('shard-0')}}}"
         assert collector.counters[key] == 1
 
     @pytest.mark.asyncio
-    async def test_client_error_emits_raw_code_label(self, mock_consumer):
+    async def test_client_error_emits_raw_code_label(self, mock_consumer, fast_asyncio_sleep):
         collector = InMemoryMetricsCollector()
         consumer = mock_consumer(metrics_collector=collector)
         err = ClientError(
@@ -603,35 +601,39 @@ async def test_client_error_emits_raw_code_label(self, mock_consumer):
         consumer.client = MagicMock()
         consumer.client.get_records = AsyncMock(side_effect=err)
 
-        orig_sleep = asyncio.sleep
-
-        async def fast_sleep(_):
-            await orig_sleep(0)
-
-        with pytest.MonkeyPatch().context() as mp:
-            mp.setattr(asyncio, "sleep", fast_sleep)
-            await consumer.get_records(self._shard())
+        await consumer.get_records(self._shard())
 
         key = (
             "consumer_errors_total{" f"error_type=ProvisionedThroughputExceededException,{_shard_labels('shard-0')}" "}"
         )
         assert collector.counters[key] == 1
 
     @pytest.mark.asyncio
-    async def test_unknown_exception_emits_unknown_label(self, mock_consumer):
+    async def test_internal_failure_no_longer_rebuilds(self, mock_consumer, fast_asyncio_sleep):
         collector = InMemoryMetricsCollector()
         consumer = mock_consumer(metrics_collector=collector)
+        err = ClientError(
+            {"Error": {"Code": "InternalFailure", "Message": "internal failure"}},
+            "GetRecords",
+        )
         consumer.client = MagicMock()
-        consumer.client.get_records = AsyncMock(side_effect=RuntimeError("surprise"))
+        consumer.client.get_records = AsyncMock(side_effect=err)
+        consumer.get_conn = AsyncMock()
+
+        await consumer.get_records(self._shard())
 
-        orig_sleep = asyncio.sleep
+        key = f"consumer_errors_total{{error_type=InternalFailure,{_shard_labels('shard-0')}}}"
+        assert collector.counters[key] == 1
+        consumer.get_conn.assert_not_awaited()
 
-        async def fast_sleep(_):
-            await orig_sleep(0)
+    @pytest.mark.asyncio
+    async def test_unknown_exception_emits_unknown_label(self, mock_consumer, fast_asyncio_sleep):
+        collector = InMemoryMetricsCollector()
+        consumer = mock_consumer(metrics_collector=collector)
+        consumer.client = MagicMock()
+        consumer.client.get_records = AsyncMock(side_effect=RuntimeError("surprise"))
 
-        with pytest.MonkeyPatch().context() as mp:
-            mp.setattr(asyncio, "sleep", fast_sleep)
-            await consumer.get_records(self._shard())
+        await consumer.get_records(self._shard())
 
         key = f"consumer_errors_total{{error_type=unknown,{_shard_labels('shard-0')}}}"
         assert collector.counters[key] == 1