Implement socket keepalive configuration mechanisms (#332)

2 files changed, 147 insertions, 1 deletions

diff --git a/pymemcache/client/base.py b/pymemcache/client/base.py
index 0999ba7..ad9ed8c 100644
--- a/pymemcache/client/base.py
+++ b/pymemcache/client/base.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import errno
+import platform
 import socket
 import six
 
@@ -38,6 +39,10 @@ VALID_STORE_RESULTS = {
     b'cas':     (b'STORED', b'EXISTS', b'NOT_FOUND'),
 }
 
+SOCKET_KEEPALIVE_SUPPORTED_SYSTEM = {
+    'Linux',
+}
+
 
 # Some of the values returned by the "stats" command
 # need mapping into native Python types
@@ -130,6 +135,44 @@ def normalize_server_spec(server):
     return (host, port)
 
 
+class KeepaliveOpts(object):
+    """
+    A configuration structure to define the socket keepalive.
+
+    This structure must be passed to a client. The client will configure
+    its socket keepalive by using the elements of the structure.
+    """
+    __slots__ = ('idle', 'intvl', 'cnt')
+
+    def __init__(self, idle=1, intvl=1, cnt=5):
+        """
+        Constructor.
+
+        Args:
+          idle: The time (in seconds) the connection needs to remain idle
+            before TCP starts sending keepalive probes. Should be a positive
+            integer most greater than zero.
+          intvl: The time (in seconds) between individual keepalive probes.
+            Should be a positive integer most greater than zero.
+          cnt: The maximum number of keepalive probes TCP should send before
+            dropping the connection. Should be a positive integer most greater
+            than zero.
+        """
+
+        if idle < 1:
+            raise ValueError(
+                "The idle parameter must be greater or equal to 1.")
+        self.idle = idle
+        if intvl < 1:
+            raise ValueError(
+                "The intvl parameter must be greater or equal to 1.")
+        self.intvl = intvl
+        if cnt < 1:
+            raise ValueError(
+                "The cnt parameter must be greater or equal to 1.")
+        self.cnt = cnt
+
+
 class Client(object):
     """
     A client for a single memcached server.
@@ -236,6 +279,7 @@ class Client(object):
                  no_delay=False,
                  ignore_exc=False,
                  socket_module=socket,
+                 socket_keepalive=None,
                  key_prefix=b'',
                  default_noreply=True,
                  allow_unicode_keys=False,
@@ -262,6 +306,9 @@ class Client(object):
             misses. Defaults to False.
           socket_module: socket module to use, e.g. gevent.socket. Defaults to
             the standard library's socket module.
+          socket_keepalive: Activate the socket keepalive feature by passing
+            a KeepaliveOpts structure in this parameter. Disabled by default
+            (None). This feature is only supported on Linux platforms.
           key_prefix: Prefix of key. You can use this as namespace. Defaults
             to b''.
           default_noreply: bool, the default value for 'noreply' as passed to
@@ -281,6 +328,32 @@ class Client(object):
         self.no_delay = no_delay
         self.ignore_exc = ignore_exc
         self.socket_module = socket_module
+        self.socket_keepalive = socket_keepalive
+        user_system = platform.system()
+        if self.socket_keepalive is not None:
+            if user_system not in SOCKET_KEEPALIVE_SUPPORTED_SYSTEM:
+                raise SystemError(
+                    "Pymemcache's socket keepalive mechaniss doesn't "
+                    "support your system ({user_system}). If "
+                    "you see this message it mean that you tried to "
+                    "configure your socket keepalive on an unsupported "
+                    "system. To fix the problem pass `socket_"
+                    "keepalive=False` or use a supported system. "
+                    "Supported systems are: {systems}".format(
+                        user_system=user_system,
+                        systems=", ".join(sorted(
+                            SOCKET_KEEPALIVE_SUPPORTED_SYSTEM))
+                    )
+                )
+            if not isinstance(self.socket_keepalive, KeepaliveOpts):
+                raise ValueError(
+                    "Unsupported keepalive options. If you see this message "
+                    "it means that you passed an unsupported object within "
+                    "the param `socket_keepalive`. To fix it "
+                    "please instantiate and pass to socket_keepalive a "
+                    "KeepaliveOpts object. That's the only supported type "
+                    "of structure."
+                )
         self.sock = None
         if isinstance(key_prefix, six.text_type):
             key_prefix = key_prefix.encode('ascii')
@@ -333,6 +406,14 @@ class Client(object):
 
         try:
             sock.settimeout(self.connect_timeout)
+            if self.socket_keepalive is not None:
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE,
+                                self.socket_keepalive.idle)
+                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL,
+                                self.socket_keepalive.intvl)
+                sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT,
+                                self.socket_keepalive.cnt)
             sock.connect(sockaddr)
             sock.settimeout(self.timeout)
         except Exception:
diff --git a/pymemcache/test/test_client.py b/pymemcache/test/test_client.py
index 7479e48..c09f48a 100644
--- a/pymemcache/test/test_client.py
+++ b/pymemcache/test/test_client.py
@@ -21,13 +21,19 @@ import functools
 import json
 import os
 import mock
+import platform
 import re
 import socket
 import unittest
 
 import pytest
 
-from pymemcache.client.base import PooledClient, Client, normalize_server_spec
+from pymemcache.client.base import (
+    PooledClient,
+    Client,
+    normalize_server_spec,
+    KeepaliveOpts
+)
 from pymemcache.exceptions import (
     MemcacheClientError,
     MemcacheServerError,
@@ -1172,6 +1178,47 @@ class TestClientSocketConnect(unittest.TestCase):
             client._connect()
             assert client.sock.family == socket.AF_UNIX
 
+    @unittest.skipIf('Linux' != platform.system(),
+                     'Socket keepalive only support Linux platforms.')
+    def test_linux_socket_keepalive(self):
+        server = ('::1', 11211)
+        try:
+            client = Client(
+                server,
+                socket_module=MockSocketModule(),
+                socket_keepalive=KeepaliveOpts())
+            client._connect()
+        except SystemError:
+            self.fail("SystemError unexpectedly raised")
+        with self.assertRaises(ValueError):
+            # A KeepaliveOpts object is expected, a ValueError will be raised
+            Client(
+                server,
+                socket_module=MockSocketModule(),
+                socket_keepalive=True)
+
+    @mock.patch('platform.system')
+    def test_osx_socket_keepalive(self, platform_mock):
+        platform_mock.return_value = 'Darwin'
+        server = ('::1', 11211)
+        # For the moment the socket keepalive is only implemented for Linux
+        with self.assertRaises(SystemError):
+            Client(
+                server,
+                socket_module=MockSocketModule(),
+                socket_keepalive=KeepaliveOpts())
+
+    @mock.patch('platform.system')
+    def test_windows_socket_keepalive(self, platform_mock):
+        platform_mock.return_value = 'Windows'
+        server = ('::1', 11211)
+        # For the moment the socket keepalive is only implemented for Linux
+        with self.assertRaises(SystemError):
+            Client(
+                server,
+                socket_module=MockSocketModule(),
+                socket_keepalive=KeepaliveOpts())
+
     def test_socket_connect_closes_on_failure(self):
         server = ("example.com", 11211)
 
@@ -1451,3 +1498,21 @@ class TestNormalizeServerSpec(unittest.TestCase):
         with pytest.raises(ValueError) as excinfo:
             f(12345)
         assert str(excinfo.value) == "Unknown server provided: 12345"
+
+
+@pytest.mark.unit()
+class TestKeepaliveopts(unittest.TestCase):
+    def test_keepalive_opts(self):
+        kao = KeepaliveOpts()
+        assert (kao.idle == 1 and kao.intvl == 1 and kao.cnt == 5)
+        kao = KeepaliveOpts(idle=1, intvl=4, cnt=2)
+        assert (kao.idle == 1 and kao.intvl == 4 and kao.cnt == 2)
+        kao = KeepaliveOpts(idle=8)
+        assert (kao.idle == 8 and kao.intvl == 1 and kao.cnt == 5)
+        kao = KeepaliveOpts(cnt=8)
+        assert (kao.idle == 1 and kao.intvl == 1 and kao.cnt == 8)
+
+        with self.assertRaises(ValueError):
+            KeepaliveOpts(cnt=0)
+        with self.assertRaises(ValueError):
+            KeepaliveOpts(idle=-1)