Repository: redis/redis-py
Pull request #1797 (2021-12-16 03:41:00)

Socket not closing after server disconnect
### Pull Request check-list

_Please make sure to review and check all of these items:_

- [x] Does `$ tox` pass with this change (including linting)?
- [x] Do the CI tests pass with this change (enable it first in your forked repo and wait for the github action build to finish)?
- [x] Is the new or changed code fully tested?
- [x] Is a documentation update included (if this change modifies existing APIs, or introduces new ones)?
- [x] Is there an example added to the examples folder (if applicable)?

_NOTE: these things are not required to open a PR and can be done
afterwards / while the PR is open._

### Description of change

This is based on @michael-lazar's recent work on fixing a Celery memory leak caused by py-amqp: https://github.com/celery/py-amqp/pull/374

I recently learned that `socket.shutdown()` can throw an `OSError` when the server caused the disconnection. I also learned that `socket.shutdown` and `socket.close` have different purposes: https://stackoverflow.com/a/598759

So, you still need to run `socket.close` after `socket.shutdown` to allow garbage collection to deallocate the socket.

Currently, redis-py's `Connection.disconnect()` won't continue to close the socket if it encounters an `OSError` on shutdown. This prevents garbage collection from cleaning up the socket and will use more memory than necessary (maybe even causing a leak?).

This PR puts a separate try/except around `socket.close` and `socket.shutdown` to allow the socket to be closed after a failed shutdown.

chayim: @pawl Great catch, and a great PR, thank you! I'm particularly appreciative of your adding unit tests for this. Approving for merge, and it'll be in the upcoming release. 

Patch:
diff --git a/redis/connection.py b/redis/connection.py
index 6c4494b1..1bb8eb50 100755
--- a/redis/connection.py
+++ b/redis/connection.py
@@ -717,9 +717,14 @@ class Connection:
         self._parser.on_disconnect()
         if self._sock is None:
             return
-        try:
-            if os.getpid() == self.pid:
+
+        if os.getpid() == self.pid:
+            try:
                 self._sock.shutdown(socket.SHUT_RDWR)
+            except OSError:
+                pass
+
+        try:
             self._sock.close()
         except OSError:
             pass


Test patch:
diff --git a/tests/test_connection.py b/tests/test_connection.py
index 22f1b718..d94a8151 100644
--- a/tests/test_connection.py
+++ b/tests/test_connection.py
@@ -3,6 +3,7 @@ from unittest import mock
 
 import pytest
 
+from redis.connection import Connection
 from redis.exceptions import InvalidResponse
 from redis.utils import HIREDIS_AVAILABLE
 
@@ -40,3 +41,36 @@ def test_loading_external_modules(modclient):
     # mod = j(modclient)
     # mod.set("fookey", ".", d)
     # assert mod.get('fookey') == d
+
+
+class TestConnection:
+    def test_disconnect(self):
+        conn = Connection()
+        mock_sock = mock.Mock()
+        conn._sock = mock_sock
+        conn.disconnect()
+        mock_sock.shutdown.assert_called_once()
+        mock_sock.close.assert_called_once()
+        assert conn._sock is None
+
+    def test_disconnect__shutdown_OSError(self):
+        """An OSError on socket shutdown will still close the socket."""
+        conn = Connection()
+        mock_sock = mock.Mock()
+        conn._sock = mock_sock
+        conn._sock.shutdown.side_effect = OSError
+        conn.disconnect()
+        mock_sock.shutdown.assert_called_once()
+        mock_sock.close.assert_called_once()
+        assert conn._sock is None
+
+    def test_disconnect__close_OSError(self):
+        """An OSError on socket close will still clear out the socket."""
+        conn = Connection()
+        mock_sock = mock.Mock()
+        conn._sock = mock_sock
+        conn._sock.close.side_effect = OSError
+        conn.disconnect()
+        mock_sock.shutdown.assert_called_once()
+        mock_sock.close.assert_called_once()
+        assert conn._sock is None