commit f6511c904bcac1dd8c7e24aa3432326738865f1b
author Jun <Junnplus@users.noreply.github.com> Fri Feb 01 12:07:58 2019 +0800
committer James E. King III <jking@apache.org> Thu Jan 31 23:07:58 2019 -0500
tree 4fd5178eab2ecf88282db1d197e2df6baa74f51d
parent 2ec93c8a2da2531755078ab6d5a65a96e26cf4c2

THRIFT-4767: Added TSocket/TSSLSocket support tcp keep-alive in python (#1725)

* THRIFT-4767: Added TSocket support tcp keep-alive in python
* TSSLSocket support tcp keep-alive

lib/py/src/transport/TSSLSocket.py

diff --git a/lib/py/src/transport/TSSLSocket.py b/lib/py/src/transport/TSSLSocket.py
index 00d1f10..c87e81c 100644
--- a/lib/py/src/transport/TSSLSocket.py
+++ b/lib/py/src/transport/TSSLSocket.py
@@ -232,6 +232,7 @@
           ``validate_callback`` (cert, hostname) -> None:
               Called after SSL handshake. Can raise when hostname does not
               match the cert.
+          ``socket_keepalive`` enable TCP keepalive, default off.
         """
         self.is_valid = False
         self.peercert = None
@@ -259,9 +260,11 @@
             kwargs['cert_reqs'] = ssl.CERT_REQUIRED if validate else ssl.CERT_NONE
 
         unix_socket = kwargs.pop('unix_socket', None)
+        socket_keepalive = kwargs.pop('socket_keepalive', False)
         self._validate_callback = kwargs.pop('validate_callback', _match_hostname)
         TSSLBase.__init__(self, False, host, kwargs)
-        TSocket.TSocket.__init__(self, host, port, unix_socket)
+        TSocket.TSocket.__init__(self, host, port, unix_socket,
+                                 socket_keepalive=socket_keepalive)
 
     def close(self):
         try:

lib/py/src/transport/TSocket.py

diff --git a/lib/py/src/transport/TSocket.py b/lib/py/src/transport/TSocket.py
index a7d6617..c8be25a 100644
--- a/lib/py/src/transport/TSocket.py
+++ b/lib/py/src/transport/TSocket.py
@@ -50,7 +50,9 @@
 class TSocket(TSocketBase):
     """Socket implementation of TTransport base."""
 
-    def __init__(self, host='localhost', port=9090, unix_socket=None, socket_family=socket.AF_UNSPEC):
+    def __init__(self, host='localhost', port=9090, unix_socket=None,
+                 socket_family=socket.AF_UNSPEC,
+                 socket_keepalive=False):
         """Initialize a TSocket
 
         @param host(str)  The host to connect to.
@@ -58,6 +60,7 @@
         @param unix_socket(str)  The filename of a unix socket to connect to.
                                  (host and port will be ignored.)
         @param socket_family(int)  The socket family to use with this socket.
+        @param socket_keepalive(bool) enable TCP keepalive, default off.
         """
         self.host = host
         self.port = port
@@ -65,6 +68,7 @@
         self._unix_socket = unix_socket
         self._timeout = None
         self._socket_family = socket_family
+        self._socket_keepalive = socket_keepalive
 
     def setHandle(self, h):
         self.handle = h
@@ -99,6 +103,11 @@
             raise TTransportException(TTransportException.NOT_OPEN, msg)
         for family, socktype, _, _, sockaddr in addrs:
             handle = self._do_open(family, socktype)
+
+            # TCP_KEEPALIVE
+            if self._socket_keepalive:
+                handle.setsockopt(socket.IPPROTO_TCP, socket.SO_KEEPALIVE, 1)
+
             handle.settimeout(self._timeout)
             try:
                 handle.connect(sockaddr)