python: fix connection persistance

Author: dkropachev

python/alternator_lb.py

@@ -56,7 +56,7 @@ class Config:
     aws_secret_access_key: str = "fake-alternator-lb-secret-access-key"
     update_interval: int = 10
     connect_timeout: int = 3600
-    max_pool_connections: int = 200
+    max_pool_connections: int = 10
 
     def _get_nodes(self) -> List[str]:
         nodes = []
@@ -120,16 +120,39 @@ def _get_connection_pool(self, parsed):
             if pool:
                 return pool
             if self._config.schema == "http":
-                pool = urllib3.HTTPConnectionPool(
-                    host=parsed.hostname,
-                    port=parsed.port,
-                )
+                if self._config.max_pool_connections == 1:
+                    # There is some bug in urllib3, if `size` = 1, pool still grows more than 1
+                    # only way to make it not to grow is to not provide `size` at all
+                    pool = urllib3.HTTPConnectionPool(
+                        host=parsed.hostname,
+                        port=parsed.port,
+                        timeout=self._config.connect_timeout,
+                    )
+                else:
+                    pool = urllib3.HTTPConnectionPool(
+                        host=parsed.hostname,
+                        port=parsed.port,
+                        timeout=self._config.connect_timeout,
+                        size=self._config.max_pool_connections,
+                    )
             else:
-                pool = urllib3.HTTPSConnectionPool(
-                    host=parsed.hostname,
-                    port=parsed.port,
-                    cert_reqs='CERT_NONE',
-                )
+                if self._config.max_pool_connections == 1:
+                    # There is some bug in urllib3, if `size` = 1, pool still grows more than 1
+                    # only way to make it not to grow is to not provide `size` at all
+                    pool = urllib3.HTTPSConnectionPool(
+                        host=parsed.hostname,
+                        port=parsed.port,
+                        cert_reqs='CERT_NONE',
+                        timeout=self._config.connect_timeout,
+                    )
+                else:
+                    pool = urllib3.HTTPSConnectionPool(
+                        host=parsed.hostname,
+                        port=parsed.port,
+                        cert_reqs='CERT_NONE',
+                        timeout=self._config.connect_timeout,
+                        size=self._config.max_pool_connections,
+                    )
             self._conn_pools[parsed.netloc] = pool
             return pool
 
@@ -192,9 +215,11 @@ def _get_nodes(self, uri: str) -> List[str]:
             url = parsed.path
             if parsed.query:
                 url += "?" + parsed.query
-            response = self._get_connection_pool(parsed).request("GET", url)
-            if response.status != 200:
-                return []
+            pool = self._get_connection_pool(parsed)
+            with pool:
+                response = pool.request("GET", url)
+                if response.status != 200:
+                    return []
 
             nodes = json.loads(response.data)
             return [self._host_to_uri(host) for host in nodes if host and self._validate_node(host)]
@@ -241,13 +266,14 @@ def get_known_nodes(self):
 
     def _init_botocore_config(self) -> config.Config:
         config_params = {
-            "tcp_keepalive": True,
+            "tcp_keepalive": bool(self._config.max_pool_connections),
             "connect_timeout": self._config.connect_timeout,
             "max_pool_connections": self._config.max_pool_connections,
         }
         if self._config.client_cert_file:
             if self._config.client_key_file:
-                config_params["client_cert"] = (self._config.client_cert_file, self._config.client_key_file)
+                config_params["client_cert"] = (
+                    self._config.client_cert_file, self._config.client_key_file)
             else:
                 config_params["client_cert"] = self._config.client_cert_file
         return config.Config(**config_params)