Refactor tcp_keepalive code (#1578)

* Move it into a separate function
* set keepalive on the REST API socket

The function will be also used in #1162

patroni/api.py

@@ -10,14 +10,15 @@ import os
 import six
 import socket
 
-from patroni.exceptions import PostgresConnectionException, PostgresException
-from patroni.postgresql.misc import postgres_version_to_int
-from patroni.utils import deep_compare, parse_bool, patch_config, Retry, \
-    RetryFailedError, parse_int, split_host_port, tzutc, uri, cluster_as_json
 from six.moves.BaseHTTPServer import BaseHTTPRequestHandler, HTTPServer
 from six.moves.socketserver import ThreadingMixIn
 from threading import Thread
 
+from .exceptions import PostgresConnectionException, PostgresException
+from .postgresql.misc import postgres_version_to_int
+from .utils import deep_compare, enable_keepalive, parse_bool, patch_config, Retry, \
+    RetryFailedError, parse_int, split_host_port, tzutc, uri, cluster_as_json
+
 logger = logging.getLogger(__name__)
 
 
@@ -626,6 +627,7 @@ class RestApiServer(ThreadingMixIn, HTTPServer, Thread):
     def get_request(self):
         sock = self.socket
         newsock, addr = socket.socket.accept(sock)
+        enable_keepalive(newsock, 10, 3)
         if hasattr(sock, 'context'):  # SSLSocket, we want to do the deferred handshake from a thread
             newsock = (sock, newsock)
         return newsock, addr

patroni/dcs/kubernetes.py

@@ -15,7 +15,7 @@ from threading import Condition, Lock, Thread
 
 from . import AbstractDCS, Cluster, ClusterConfig, Failover, Leader, Member, SyncState, TimelineHistory
 from ..exceptions import DCSError
-from ..utils import deep_compare, Retry, RetryFailedError, tzutc, USER_AGENT
+from ..utils import deep_compare, keepalive_socket_options, Retry, RetryFailedError, tzutc, USER_AGENT
 
 logger = logging.getLogger(__name__)
 
@@ -53,16 +53,8 @@ class CoreV1ApiProxy(object):
         # If we didn't received anything after the loop_wait + retry_timeout it is a time
         # to start worrying (send keepalive messages). Finally, the connection should be
         # considered as dead if we received nothing from the socket after the ttl seconds.
-        cnt = 3
-        idle = int(loop_wait + retry_timeout)
-        intvl = max(1, int(float(ttl - idle) / cnt))
-        self._api.api_client.rest_client.pool_manager.connection_pool_kw['socket_options'] = [
-            (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),
-            (socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, idle),
-            (socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, intvl),
-            (socket.IPPROTO_TCP, socket.TCP_KEEPCNT, cnt),
-            (socket.IPPROTO_TCP, 18, int(ttl * 1000))  # TCP_USER_TIMEOUT
-        ]
+        self._api.api_client.rest_client.pool_manager.connection_pool_kw['socket_options'] = \
+                list(keepalive_socket_options(ttl, int(loop_wait + retry_timeout)))
         self._request_timeout = (1, retry_timeout / 3.0)
 
     def __getattr__(self, func):

patroni/utils.py

@@ -3,6 +3,8 @@ import os
 import platform
 import random
 import re
+import socket
+import sys
 import tempfile
 import time
 
@@ -453,3 +455,38 @@ def data_directory_is_empty(data_dir):
     if not os.path.exists(data_dir):
         return True
     return all(os.name != 'nt' and (n.startswith('.') or n == 'lost+found') for n in os.listdir(data_dir))
+
+
+def keepalive_intvl(timeout, idle, cnt=3):
+    return max(1, int(float(timeout - idle) / cnt))
+
+
+def keepalive_socket_options(timeout, idle, cnt=3):
+    yield (socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
+
+    if sys.platform.startswith('linux'):
+        yield (socket.SOL_TCP, 18, int(timeout * 1000))  # TCP_USER_TIMEOUT
+        TCP_KEEPIDLE = socket.TCP_KEEPIDLE
+        TCP_KEEPINTVL = socket.TCP_KEEPINTVL
+        TCP_KEEPCNT = socket.TCP_KEEPCNT
+    elif sys.platform.startswith('darwin'):
+        TCP_KEEPIDLE = 0x10  # (named "TCP_KEEPALIVE" in C)
+        TCP_KEEPINTVL = 0x101
+        TCP_KEEPCNT = 0x102
+    else:
+        return
+
+    intvl = keepalive_intvl(timeout, idle, cnt)
+    yield (socket.IPPROTO_TCP, TCP_KEEPIDLE, idle)
+    yield (socket.IPPROTO_TCP, TCP_KEEPINTVL, intvl)
+    yield (socket.IPPROTO_TCP, TCP_KEEPCNT, cnt)
+
+
+def enable_keepalive(sock, timeout, idle, cnt=3):
+    SIO_KEEPALIVE_VALS = getattr(socket, 'SIO_KEEPALIVE_VALS', None)
+    if SIO_KEEPALIVE_VALS is not None:  # Windows
+        intvl = keepalive_intvl(timeout, idle, cnt)
+        return sock.ioctl(SIO_KEEPALIVE_VALS, (1, idle * 1000, intvl * 1000))
+
+    for opt in keepalive_socket_options(timeout, idle, cnt):
+        sock.setsockopt(*opt)

tests/test_api.py

@@ -465,7 +465,9 @@ class TestRestApiServer(unittest.TestCase):
         mock_socket.context.wrap_socket.side_effect = socket.error
         self.srv.process_request_thread((mock_socket, 1), '2')
 
-    @patch.object(socket.socket, 'accept', Mock(return_value=(1, '2')))
-    def test_get_request(self):
+    @patch.object(socket.socket, 'accept')
+    def test_get_request(self, mock_accept):
+        newsock = Mock()
+        mock_accept.return_value = (newsock, '2')
         self.srv.socket = Mock()
-        self.assertEqual(self.srv.get_request(), ((self.srv.socket, 1), '2'))
+        self.assertEqual(self.srv.get_request(), ((self.srv.socket, newsock), '2'))

tests/test_utils.py

@@ -2,7 +2,7 @@ import unittest
 
 from mock import Mock, patch
 from patroni.exceptions import PatroniException
-from patroni.utils import Retry, RetryFailedError, polling_loop, validate_directory
+from patroni.utils import Retry, RetryFailedError, enable_keepalive, polling_loop, validate_directory
 
 
 class TestUtils(unittest.TestCase):
@@ -33,6 +33,13 @@ class TestUtils(unittest.TestCase):
     def test_validate_directory_is_not_a_directory(self):
         self.assertRaises(PatroniException, validate_directory, "/tmp")
 
+    def test_enable_keepalive(self):
+        with patch('socket.SIO_KEEPALIVE_VALS', 1, create=True):
+            self.assertIsNotNone(enable_keepalive(Mock(), 10, 5))
+        for platform in ('linux2', 'darwin', 'other'):
+            with patch('sys.platform', platform):
+                self.assertIsNone(enable_keepalive(Mock(), 10, 5))
+
 
 @patch('time.sleep', Mock())
 class TestRetrySleeper(unittest.TestCase):