22# Copyright (c) 2020-2021 Pinecone Systems Inc. All right reserved.
33#
44import logging
5- from typing import NamedTuple
5+ import sys
6+ from typing import NamedTuple , List
67import os
78
89import certifi
910import requests
1011import configparser
12+ import socket
13+
14+ from urllib3 .connection import HTTPConnection
1115
1216from pinecone .core .client .exceptions import ApiKeyError
1317from pinecone .core .api_action import ActionAPI , WhoAmIResponse
1418from pinecone .core .utils import warn_deprecated
15- from pinecone .core .utils .constants import CLIENT_VERSION , PARENT_LOGGER_NAME , DEFAULT_PARENT_LOGGER_LEVEL
19+ from pinecone .core .utils .constants import CLIENT_VERSION , PARENT_LOGGER_NAME , DEFAULT_PARENT_LOGGER_LEVEL , \
20+ TCP_KEEPIDLE , TCP_KEEPINTVL , TCP_KEEPCNT
1621from pinecone .core .client .configuration import Configuration as OpenApiConfiguration
1722
1823__all__ = [
@@ -37,7 +42,7 @@ class _CONFIG:
3742
3843 Order of configs to load:
3944
40- - configs specified explictly in reset
45+ - configs specified explicitly in reset
4146 - environment variables
4247 - configs specified in the INI file
4348 - default configs
@@ -109,6 +114,8 @@ def reset(self, config_file=None, **kwargs):
109114 or default_openapi_config
110115 )
111116
117+ openapi_config .socket_options = self ._get_socket_options ()
118+
112119 config = config ._replace (openapi_config = openapi_config )
113120 self ._config = config
114121
@@ -144,6 +151,54 @@ def _load_config_file(self, config_file: str) -> dict:
144151 config_obj = {** parser ["default" ]}
145152 return config_obj
146153
154+ @staticmethod
155+ def _get_socket_options (do_keep_alive : bool = True ,
156+ keep_alive_idle_sec : int = TCP_KEEPIDLE ,
157+ keep_alive_interval_sec : int = TCP_KEEPINTVL ,
158+ keep_alive_tries : int = TCP_KEEPCNT
159+ ) -> List [tuple ]:
160+ """
161+ Returns the socket options to pass to OpenAPI's Rest client
162+ Args:
163+ do_keep_alive: Whether to enable TCP keep alive mechanism
164+ keep_alive_idle_sec: Time in seconds of connection idleness before starting to send keep alive probes
165+ keep_alive_interval_sec: Interval time in seconds between keep alive probe messages
166+ keep_alive_tries: Number of failed keep alive tries (unanswered KA messages) before terminating the connection
167+
168+ Returns:
169+ A list of socket options for the Rest client's connection pool
170+ """
171+ # Source: https://www.finbourne.com/blog/the-mysterious-hanging-client-tcp-keep-alives
172+
173+ socket_params = HTTPConnection .default_socket_options
174+ if not do_keep_alive :
175+ return socket_params
176+
177+ socket_params += [(socket .SOL_SOCKET , socket .SO_KEEPALIVE , 1 )]
178+
179+ # TCP Keep Alive Probes for different platforms
180+ platform = sys .platform
181+ # TCP Keep Alive Probes for Linux
182+ if platform == 'linux' and hasattr (socket , "TCP_KEEPIDLE" ) and hasattr (socket , "TCP_KEEPINTVL" ) \
183+ and hasattr (socket , "TCP_KEEPCNT" ):
184+ socket_params += [(socket .IPPROTO_TCP , socket .TCP_KEEPIDLE , keep_alive_idle_sec )]
185+ socket_params += [(socket .IPPROTO_TCP , socket .TCP_KEEPINTVL , keep_alive_interval_sec )]
186+ socket_params += [(socket .IPPROTO_TCP , socket .TCP_KEEPCNT , keep_alive_tries )]
187+
188+ # TCP Keep Alive Probes for Windows OS
189+ # NOTE: Changing TCP KA params on windows is done via a different mechanism which OpenAPI's Rest client doesn't expose.
190+ # Since the default values work well, it seems setting `(socket.SO_KEEPALIVE, 1)` is sufficient.
191+ # Leaving this code here for future reference.
192+ # elif platform == 'win32' and hasattr(socket, "SIO_KEEPALIVE_VALS"):
193+ # socket.ioctl((socket.SIO_KEEPALIVE_VALS, (1, keep_alive_idle_sec * 1000, keep_alive_interval_sec * 1000)))
194+
195+ # TCP Keep Alive Probes for Mac OS
196+ elif platform == 'darwin' :
197+ TCP_KEEPALIVE = 0x10
198+ socket_params += [(socket .IPPROTO_TCP , TCP_KEEPALIVE , keep_alive_interval_sec )]
199+
200+ return socket_params
201+
147202 @property
148203 def ENVIRONMENT (self ):
149204 return self ._config .environment
0 commit comments