Merge pull request #18 from lyft/ramonpetgrave64-max-connection-pool

add kms performance parameters

Author: ramonpetgrave64

README.md

@@ -104,6 +104,18 @@ validator.decrypt_token(username, token)
 Note: 'to', 'from', and 'user_type' keys are not allowed to be set in
 extra_context.
 
+## Performance Tuning
+
+With the [boto defaults](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html), the AWS KMS client used in `KMSTokenValidator` may not be performant under higher loads, due to latency when communicating with AWS KMS. Try tuning these parameters below with the given starting points.
+
+```python
+...
+max_pool_connections=100,
+connect_timeout=1,
+read_timeout=1,
+...
+```
+
 ## Reporting security vulnerabilities
 
 If you've found a vulnerability or a potential vulnerability in kmsauth

kmsauth/__init__.py

@@ -57,6 +57,9 @@ def __init__(
             endpoint_url=None,
             token_cache_size=4096,
             stats=None,
+            max_pool_connections=None,
+            connect_timeout=None,
+            read_timeout=None,
             ):
         """Create a KMSTokenValidator object.
 
@@ -107,13 +110,19 @@ def __init__(
                 aws_access_key_id=self.aws_creds['AccessKeyId'],
                 aws_secret_access_key=self.aws_creds['SecretAccessKey'],
                 aws_session_token=self.aws_creds['SessionToken'],
-                endpoint_url=endpoint_url
+                endpoint_url=endpoint_url,
+                max_pool_connections=max_pool_connections,
+                connect_timeout=connect_timeout,
+                read_timeout=read_timeout,
             )
         else:
             self.kms_client = kmsauth.services.get_boto_client(
                 'kms',
                 region=self.region,
-                endpoint_url=endpoint_url
+                endpoint_url=endpoint_url,
+                max_pool_connections=max_pool_connections,
+                connect_timeout=connect_timeout,
+                read_timeout=read_timeout,
             )
         if extra_context is None:
             self.extra_context = {}

kmsauth/services.py

@@ -1,6 +1,7 @@
 """Module for accessing boto3 clients, resources and sessions."""
 
 import boto3
+import botocore
 import logging
 
 CLIENT_CACHE = {}
@@ -13,7 +14,10 @@ def get_boto_client(
         aws_access_key_id=None,
         aws_secret_access_key=None,
         aws_session_token=None,
-        endpoint_url=None
+        endpoint_url=None,
+        max_pool_connections=None,
+        connect_timeout=None,
+        read_timeout=None,
         ):
     """Get a boto3 client connection."""
     cache_key = '{0}:{1}:{2}:{3}'.format(
@@ -37,7 +41,12 @@ def get_boto_client(
 
     CLIENT_CACHE[cache_key] = session.client(
         client,
-        endpoint_url=endpoint_url
+        endpoint_url=endpoint_url,
+        config=botocore.config.Config(
+            max_pool_connections=max_pool_connections,
+            connect_timeout=connect_timeout,
+            read_timeout=read_timeout,
+        )
     )
     return CLIENT_CACHE[cache_key]
 

setup.py

@@ -13,7 +13,7 @@
 
 from setuptools import setup, find_packages
 
-VERSION = "0.6.1"
+VERSION = "0.6.2"
 
 requirements = [
     # Boto3 is the Amazon Web Services (AWS) Software Development Kit (SDK)