diff --git a/patroni/api.py b/patroni/api.py index 23c4e6bf..3d5cf7cc 100644 --- a/patroni/api.py +++ b/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 diff --git a/patroni/dcs/kubernetes.py b/patroni/dcs/kubernetes.py index 44276da8..a4630578 100644 --- a/patroni/dcs/kubernetes.py +++ b/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): diff --git a/patroni/utils.py b/patroni/utils.py index 94b98524..34c70034 100644 --- a/patroni/utils.py +++ b/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) diff --git a/tests/test_api.py b/tests/test_api.py index 32de0869..9d6b18d9 100644 --- a/tests/test_api.py +++ b/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')) diff --git a/tests/test_utils.py b/tests/test_utils.py index ba3405cc..e08891fe 100644 --- a/tests/test_utils.py +++ b/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):