For Python < 3.5, automatically retry EINTR

2 files changed, 61 insertions, 4 deletions

diff --git a/redis/_compat.py b/redis/_compat.py
index 97e4919..d0a6cf3 100644
--- a/redis/_compat.py
+++ b/redis/_compat.py
@@ -1,6 +1,63 @@
 """Internal module for Python 2 backwards compatibility."""
 import sys
 
+# For Python older than 3.5, retry EINTR.
+if sys.version_info[0] < 3 or sys.version_info[0] == 3 and sys.version_info[1] < 5:
+    # Adapted from https://bugs.python.org/review/23863/patch/14532/54418
+    import socket
+    import time
+    import errno
+
+    # Wrapper for handling interruptable system calls.
+    def _retryable_call(s, func, *args, **kwargs):
+        # Some modules (SSL) use the _fileobject wrapper directly and
+        # implement a smaller portion of the socket interface, thus we
+        # need to let them continue to do so.
+        timeout, deadline = None, 0.0
+        attempted = False
+        try:
+            timeout = s.gettimeout()
+        except AttributeError:
+            pass
+
+        if timeout:
+            deadline = time.time() + timeout
+
+        try:
+            while True:
+                if attempted and timeout:
+                    now = time.time()
+                    if now >= deadline:
+                        raise socket.error(errno.EWOULDBLOCK, "timed out")
+                    else:
+                        # Overwrite the timeout on the socket object
+                        # to take into account elapsed time.
+                        s.settimeout(deadline - now)
+                try:
+                    attempted = True
+                    return func(*args, **kwargs)
+                except socket.error as e:
+                    if e.args[0] == errno.EINTR:
+                        continue
+                    raise
+        finally:
+            # Set the existing timeout back for future
+            # calls.
+            if timeout:
+                s.settimeout(timeout)
+
+    def recv(sock, *args, **kwargs):
+        return _retryable_call(sock, sock.recv, *args, **kwargs)
+
+    def recv_into(sock, *args, **kwargs):
+        return _retryable_call(sock, sock.recv_into, *args, **kwargs)
+
+else: # Python 3.5 and above automatically retry EINTR
+    def recv(sock, *args, **kwargs):
+        return sock.recv(*args, **kwargs)
+
+    def recv_into(sock, *args, **kwargs):
+        return sock.recv_into(*args, **kwargs)
 
 if sys.version_info[0] < 3:
     from urllib import unquote
diff --git a/redis/connection.py b/redis/connection.py
index 004c7a6..a552820 100755
--- a/redis/connection.py
+++ b/redis/connection.py
@@ -17,7 +17,7 @@ except ImportError:
 from redis._compat import (b, xrange, imap, byte_to_chr, unicode, bytes, long,
                            BytesIO, nativestr, basestring, iteritems,
                            LifoQueue, Empty, Full, urlparse, parse_qs,
-                           unquote)
+                           recv, recv_into, unquote)
 from redis.exceptions import (
     RedisError,
     ConnectionError,
@@ -123,7 +123,7 @@ class SocketBuffer(object):
 
         try:
             while True:
-                data = self._sock.recv(socket_read_size)
+                data = recv(self._sock, socket_read_size)
                 # an empty string indicates the server shutdown the socket
                 if isinstance(data, bytes) and len(data) == 0:
                     raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)
@@ -341,11 +341,11 @@ class HiredisParser(BaseParser):
         while response is False:
             try:
                 if HIREDIS_USE_BYTE_BUFFER:
-                    bufflen = self._sock.recv_into(self._buffer)
+                    bufflen = recv_into(self._sock, self._buffer)
                     if bufflen == 0:
                         raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)
                 else:
-                    buffer = self._sock.recv(socket_read_size)
+                    buffer = recv(self._sock, socket_read_size)
                     # an empty string indicates the server shutdown the socket
                     if not isinstance(buffer, bytes) or len(buffer) == 0:
                         raise socket.error(SERVER_CLOSED_CONNECTION_ERROR)