Expand ability to ignore some httplib calls.

Author: jonathangreen

aws_xray_sdk/ext/httplib/__init__.py

@@ -1,3 +1,3 @@
-from .patch import patch, unpatch
+from .patch import patch, unpatch, add_ignored, reset_ignored
 
-__all__ = ['patch', 'unpatch']
+__all__ = ['patch', 'unpatch', 'add_ignored', 'reset_ignored']

aws_xray_sdk/ext/httplib/patch.py

@@ -22,8 +22,33 @@
 
 _XRAY_PROP = '_xray_prop'
 _XRay_Data = namedtuple('xray_data', ['method', 'host', 'url'])
+_XRay_Ignore = namedtuple('xray_ignore', ['subclass', 'hostname', 'urls'])
 # A flag indicates whether this module is X-Ray patched or not
 PATCH_FLAG = '__xray_patched'
+# Calls that should be ignored
+_XRAY_IGNORE = set()
+
+
+def add_ignored(subclass=None, hostname=None, urls=None):
+    global _XRAY_IGNORE
+    if subclass is not None or hostname is not None or urls is not None:
+        urls = urls if urls is None else tuple(urls)
+        _XRAY_IGNORE.add(_XRay_Ignore(subclass=subclass, hostname=hostname, urls=urls))
+
+
+def reset_ignored():
+    global _XRAY_IGNORE
+    _XRAY_IGNORE.clear()
+    _ignored_add_default()
+
+
+def _ignored_add_default():
+    # skip httplib tracing for SDK built-in centralized sampling pollers
+    add_ignored(subclass='botocore.awsrequest.AWSHTTPConnection', urls=['/GetSamplingRules', '/SamplingTargets'])
+
+
+# make sure we have the default rules
+_ignored_add_default()
 
 
 def http_response_processor(wrapped, instance, args, kwargs, return_value,
@@ -77,11 +102,21 @@ def http_send_request_processor(wrapped, instance, args, kwargs, return_value,
         subsegment.add_exception(exception, stack)
 
 
+def _ignore_request(subclass, hostname, url):
+    global _XRAY_IGNORE
+    for rule in _XRAY_IGNORE:
+        subclass_match = subclass == rule.subclass if rule.subclass is not None else True
+        host_match = hostname == rule.hostname if rule.hostname is not None else True
+        url_match = url in rule.urls if rule.urls is not None else True
+        if url_match and host_match and subclass_match:
+            return True
+    return False
+
+
 def _send_request(wrapped, instance, args, kwargs):
     def decompose_args(method, url, body, headers, encode_chunked=False):
-        # skip httplib tracing for SDK built-in centralized sampling pollers
-        if (('/GetSamplingRules' in args or '/SamplingTargets' in args) and
-                type(instance).__name__ == 'botocore.awsrequest.AWSHTTPConnection'):
+        # skip any ignored requests
+        if _ignore_request(type(instance).__name__, instance.host, url):
             return wrapped(*args, **kwargs)
 
         # Only injects headers when the subsegment for the outgoing