Tweak code to set socket keepalives and options, and add test

Author: Luke Bakken

riak/tests/test_client.py

@@ -3,6 +3,7 @@
 from six import PY2
 from threading import Thread
 from riak.riak_object import RiakObject
+from riak.transports.tcp import TcpTransport
 from riak.tests import DUMMY_HTTP_PORT, DUMMY_PB_PORT, RUN_POOL
 from riak.tests.base import IntegrationTestBase
 
@@ -13,6 +14,17 @@
 
 
 class ClientTests(IntegrationTestBase, unittest.TestCase):
+    def test_can_set_tcp_keepalive(self):
+        if self.protocol == 'pbc':
+            topts = {'socket_keepalive': True}
+            c = self.create_client(transport_options=topts)
+            for i, r in enumerate(c._tcp_pool.resources):
+                self.assertIsInstance(r, TcpTransport)
+                self.assertTrue(r._socket_keepalive)
+            c.close()
+        else:
+            pass
+
     def test_uses_client_id_if_given(self):
         if self.protocol == 'pbc':
             zero_client_id = "\0\0\0\0"

riak/transports/tcp/connection.py

@@ -195,12 +195,13 @@ def _connect(self):
                                                         self._timeout)
             else:
                 self._socket = socket.create_connection(self._address)
+            if self._socket_tcp_options:
+                ka_opts = self._socket_tcp_options
+                for k, v in ka_opts.iteritems():
+                    self._socket.setsockopt(socket.SOL_TCP, k, v)
             if self._socket_keepalive:
                 self._socket.setsockopt(
                     socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
-                ka_opts = self._socket_keepalive_options or {}
-                for k, v in ka_opts.iteritems():
-                    self._socket.setsockopt(socket.SOL_TCP, k, v)
             if self._client._credentials:
                 self._init_security()
 

riak/transports/tcp/transport.py

@@ -25,8 +25,6 @@ def __init__(self,
                  node=None,
                  client=None,
                  timeout=None,
-                 socket_keepalive=False,
-                 socket_keepalive_options=None,
                  **kwargs):
         super(TcpTransport, self).__init__()
 
@@ -35,11 +33,14 @@ def __init__(self,
         self._address = (node.host, node.pb_port)
         self._timeout = timeout
         self._socket = None
-        self._socket_keepalive = socket_keepalive
-        self._socket_keepalive_options = socket_keepalive_options
         self._pbuf_c = None
         self._ttb_c = None
-        self._use_ttb = kwargs.get('use_ttb', True)
+        self._socket_tcp_options = \
+            kwargs.get('socket_tcp_options', {})
+        self._socket_keepalive = \
+            kwargs.get('socket_keepalive', False)
+        self._use_ttb = \
+            kwargs.get('use_ttb', True)
 
     def _get_pbuf_codec(self):
         if not self._pbuf_c: