[fix] block in close

richard78917 · Gr1N · commit bb07ff9490f4 · 2020-02-14T19:30:35.000+03:00
In case that close() is called while doing
read operation, the close() call is blocked and whole application
is stucked. Shut-down of socket will resolve this issue and releases
thread blocked at readline().

Regression test for graceful shutdown also included.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -2,6 +2,8 @@
 
 ## 0.7.0 (2020-XX-XX)
 
+- Fixed issue when threads are blocked on close while reading, #11 by @richard78917
+
 ## 0.6.0 (2020-01-08)
 
 - Added Python 3.8.* support, #10
diff --git a/pynats/client.py b/pynats/client.py
@@ -18,7 +18,11 @@
 
 import pkg_resources
 
-from pynats.exceptions import NATSInvalidResponse, NATSUnexpectedResponse
+from pynats.exceptions import (
+    NATSInvalidResponse,
+    NATSSocketError,
+    NATSUnexpectedResponse,
+)
 from pynats.nuid import NUID
 
 __all__ = ("NATSSubscription", "NATSMessage", "NATSClient")
@@ -154,6 +158,7 @@ def connect(self) -> None:
         self._recv(INFO_RE)
 
     def close(self) -> None:
+        self._socket.shutdown(socket.SHUT_RDWR)
         self._socket_file.close()
         self._socket.close()
 
@@ -279,7 +284,11 @@ def _readline(self, *, size: int = None) -> bytes:
         read = io.BytesIO()
 
         while True:
-            line = cast(bytes, self._socket_file.readline())
+            raw_bytes = self._socket_file.readline()
+            if not raw_bytes:
+                raise NATSSocketError(b"unable to read from socket")
+
+            line = cast(bytes, raw_bytes)
             read.write(line)
 
             if size is not None:
diff --git a/pynats/exceptions.py b/pynats/exceptions.py
@@ -15,3 +15,9 @@ class NATSInvalidResponse(NATSError):
     def __init__(self, line: bytes, *args, **kwargs) -> None:
         self.line = line
         super().__init__()
+
+
+class NATSSocketError(NATSError):
+    def __init__(self, line: bytes, *args, **kwargs) -> None:
+        self.line = line
+        super().__init__()
diff --git a/tests/test_client.py b/tests/test_client.py
@@ -7,6 +7,7 @@
 import pytest
 
 from pynats import NATSClient
+from pynats.exceptions import NATSSocketError
 
 
 @pytest.fixture
@@ -179,3 +180,24 @@ def test_request_timeout(nats_url):
     with NATSClient(nats_url, socket_timeout=2) as client:
         with pytest.raises(socket.timeout):
             client.request("test-subject")
+
+
+def test_graceful_shutdown(nats_url):
+    def worker(client, connected_event):
+        client.connect()
+        connected_event.set()
+        try:
+            client.wait()
+        except NATSSocketError:
+            assert True
+        except Exception:
+            raise AssertionError("unexpected Exception raised")
+
+    client = NATSClient(nats_url)
+    connected_event = threading.Event()
+    thread = threading.Thread(target=worker, args=[client, connected_event])
+    thread.start()
+    assert connected_event.wait(5), "unable to connect"
+    client.close()
+    thread.join(5)
+    assert not thread.is_alive(), "thread did not finish"
