Truncate buffer to prevent overflow

Author: thatbudakguy

mqlog/__init__.py

@@ -70,10 +70,18 @@ def format(self, record):
     # Named with an underscore to avoid conflict with logging.Handler.flush
     async def _flush(self):
         if self.buffer:
+            # Try to publish, but if we fail, just log the error – don't
+            # want to cause cascading errors. Something else is responsible
+            # for bringing the connection back up.
             try:
                 msg = "\n".join([line for line in self.buffer])
                 await self.client.publish(self.topic, msg, qos=self.qos)
                 self.buffer.clear()
             except Exception as e:
                 self._logger.error("Failed to publish logs", exc_info=e)
             self.will_flush.clear()
+
+        # If we were supposed to flush but couldn't, truncate the buffer
+        # to prevent it from growing indefinitely
+        if self.buffer and len(self.buffer) >= self.capacity:
+            self.buffer = self.buffer[-self.capacity :]

tests/test_handler.py

@@ -22,6 +22,11 @@ def setUp(self):
         self.logger.addHandler(self.handler)
         self.logger.setLevel(logging.INFO)
 
+    def tearDown(self):
+        """Clean up after each test"""
+        self.handler.buffer.clear()
+        self.handler.will_flush.clear()
+
     def test_default_formatter(self):
         """Handler should use default formatter if none is set"""
         self.handler.setFormatter(None)
@@ -103,3 +108,26 @@ async def do_test(handler: logging.Handler, logger: logging.Logger):
                 call("test_topic", "Test message 3\nTest message 4", qos=0),
             ]
         )
+
+    def test_buffer_overflow(self):
+        """Buffer should get truncated if it exceeds capacity"""
+        self.handler.capacity = 3
+        self.client.publish = AsyncMock(side_effect=Exception("Publish failed"))
+
+        async def do_test(handler: logging.Handler, logger: logging.Logger):
+            asyncio.create_task(handler.run())
+            await asyncio.sleep(0.1)
+            logger.info("Test message 1")
+            await asyncio.sleep(0.1)
+            logger.info("Test message 2")
+            await asyncio.sleep(0.1)
+            logger.info("Test message 3")
+            await asyncio.sleep(0.1)
+            logger.info("Test message 4")  # This will fail to flush
+            await asyncio.sleep(0.1)
+
+        asyncio.run(do_test(self.handler, self.logger))
+
+        self.assertEqual(
+            self.handler.buffer, ["Test message 2", "Test message 3", "Test message 4"]
+        )