From f4e111e31567de4ea3d7351c45da34b1b97669f4 Mon Sep 17 00:00:00 2001
From: Patrick Ellul <pellul@ipowow.com>
Date: Fri, 11 Dec 2015 14:31:54 +1100
Subject: [PATCH] adding keepalive options to pbc transport connection socket

---
 riak/transports/pbc/connection.py | 8 ++++++++
 riak/transports/pbc/transport.py  | 4 ++++
 2 files changed, 12 insertions(+)

diff --git a/riak/transports/pbc/connection.py b/riak/transports/pbc/connection.py
index 0bc58232..4d15387b 100644
--- a/riak/transports/pbc/connection.py
+++ b/riak/transports/pbc/connection.py
@@ -220,6 +220,12 @@ def _connect(self):
                                                         self._timeout)
             else:
                 self._socket = socket.create_connection(self._address)
+            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()
 
@@ -247,3 +253,5 @@ def _parse_msg(self, code, packet):
     # These are set in the RiakPbcTransport initializer
     _address = None
     _timeout = None
+    _socket_keepalive = None
+    _socket_keepalive_options = None
diff --git a/riak/transports/pbc/transport.py b/riak/transports/pbc/transport.py
index e385c698..14ef2f75 100644
--- a/riak/transports/pbc/transport.py
+++ b/riak/transports/pbc/transport.py
@@ -93,6 +93,8 @@ def __init__(self,
                  node=None,
                  client=None,
                  timeout=None,
+                 socket_keepalive=False,
+                 socket_keepalive_options=None,
                  *unused_options):
         """
         Construct a new RiakPbcTransport object.
@@ -104,6 +106,8 @@ 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
 
     # FeatureDetection API
     def _server_version(self):