fix: Prevent redis task loss when closing connection while in poll (#…

…1733) * Catch brpop before closing connection * Update unit test --------- Co-authored-by: mbierma <[email protected]>
celery · Aug 15, 2023 · 492776e · 492776e
1 parent 6f86766
commit 492776e
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/kombu/transport/redis.py b/kombu/transport/redis.py
@@ -1079,6 +1079,11 @@ def _purge(self, queue):
 
     def close(self):
         self._closing = True
+        if self._in_poll:
+            try:
+                self._brpop_read()
+            except Empty:
+                pass
         if not self.closed:
             # remove from channel poller.
             self.connection.cycle.discard(self)

diff --git a/t/unit/transport/test_redis.py b/t/unit/transport/test_redis.py
@@ -1290,6 +1290,15 @@ def test_close_disconnects(self):
         assert conn1.disconnected
         assert conn2.disconnected
 
+    def test_close_in_poll(self):
+        c = Connection(transport=Transport).channel()
+        conn1 = c.client.connection
+        conn1._sock.data = [('BRPOP', ('test_Redis',))]
+        c._in_poll = True
+        c.close()
+        assert conn1.disconnected
+        assert conn1._sock.data == []
+
     def test_get__Empty(self):
         channel = self.connection.channel()
         with pytest.raises(Empty):