From 38d668cf1fb1d97008b06b326e033d5ff7ae83ec Mon Sep 17 00:00:00 2001
From: Andrew Weiss <wvvwwvw@gmail.com>
Date: Tue, 5 Mar 2019 13:47:47 -0800
Subject: [PATCH] Allow the request pool to be configured

`requests` have an underlying request pool for making requests. They allow
you to configure things like the max pool size and the number of open
connections. This PR exposes that via configuration of rollbar.

Fixes #274
---
 rollbar/__init__.py      | 18 ++++++++++++++++++
 rollbar/lib/transport.py | 13 ++++++++++++-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/rollbar/__init__.py b/rollbar/__init__.py
index 09410d8f..2b441023 100644
--- a/rollbar/__init__.py
+++ b/rollbar/__init__.py
@@ -265,6 +265,9 @@ def _get_pylons_request():
     'http_proxy_user': None,
     'http_proxy_password': None,
     'include_request_body': False,
+    'request_pool_connections': None,
+    'request_pool_maxsize': None,
+    'request_max_retries': None,
 }
 
 _CURRENT_LAMBDA_CONTEXT = None
@@ -319,6 +322,7 @@ def init(access_token, environment='production', scrub_fields=None, url_fields=N
 
     SETTINGS['access_token'] = access_token
     SETTINGS['environment'] = environment
+    _configure_transport(**SETTINGS)
 
     if SETTINGS.get('allow_logging_basic_config'):
         logging.basicConfig()
@@ -370,6 +374,20 @@ def init(access_token, environment='production', scrub_fields=None, url_fields=N
     _initialized = True
 
 
+def _configure_transport(**kw):
+    configuration = _requests_configuration(**kw)
+    transport.configure_pool(**configuration)
+
+
+def _requests_configuration(**kw):
+    keys = {
+        'request_pool_connections': 'pool_connections',
+        'request_pool_maxsize': 'pool_maxsize',
+        'request_max_retries': 'max_retries',
+    }
+    return {keys[k]: kw.get(k, None) for k in keys}
+
+
 def lambda_function(f):
     """
     Decorator for making error handling on AWS Lambda easier
diff --git a/rollbar/lib/transport.py b/rollbar/lib/transport.py
index cf722a5a..4e8653f4 100644
--- a/rollbar/lib/transport.py
+++ b/rollbar/lib/transport.py
@@ -28,6 +28,17 @@ def _get_proxy_cfg(kw):
         }
 
 
+def configure_pool(**kw):
+    keys = ['pool_connections', 'pool_maxsize', 'max_retries']
+    args = {k: kw[k] for k in keys if kw.get(k, None) is not None}
+    if len(args) == 0:
+        return
+    https_adapter = requests.adapters.HTTPAdapter(**args)
+    http_adapter = requests.adapters.HTTPAdapter(**args)
+    _session().mount('https://', https_adapter)
+    _session().mount('http://', http_adapter)
+
+
 def post(*args, **kw):
     proxies = _get_proxy_cfg(kw)
     return _session().post(*args, proxies=proxies, **kw)
@@ -38,4 +49,4 @@ def get(*args, **kw):
     return _session().get(*args, proxies=proxies, **kw)
 
 
-__all__ = ['post', 'get']
+__all__ = ['post', 'get', 'configure_pool']