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

* Catch brpop before closing connection

* Update unit test

---------

Co-authored-by: mbierma <3448579-mbierma@users.noreply.gitlab.com>
This commit is contained in:
mbierma 2023-08-14 21:16:49 -07:00 committed by GitHub
parent 6f8676630d
commit 492776eb7f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 14 additions and 0 deletions

View File

@ -1079,6 +1079,11 @@ class Channel(virtual.Channel):
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)

View File

@ -1290,6 +1290,15 @@ class test_Redis:
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):