Add ability to enable/disable the SDK (#26)

Revision 2

Changes:
    - SDKConfig now built on top of recorder. SDK-level configuration should all be done
      through SDKConfig.
    - Renamed the SDK Enable environment variable to AWS_XRAY_SDK_ENABLED
    - Disabling the SDK now causes LambdaContext to set sampling decision of facade segment to False.
      Previously, it would force all subsegments to be generated as dummy subsegments.
    - SDKConfig.set_sdk_enabled(value) method defaults to True and throws an exception
      if value is not of type boolean.

Author: chanchiem

aws_xray_sdk/__init__.py

@@ -0,0 +1,3 @@
+from .sdk_config import SDKConfig
+
+global_sdk_config = SDKConfig()

aws_xray_sdk/core/lambda_launcher.py

@@ -2,10 +2,9 @@
 import logging
 import threading
 
-from aws_xray_sdk.sdk_config import SDKConfig
+import aws_xray_sdk
 from .models.facade_segment import FacadeSegment
 from .models.trace_header import TraceHeader
-from .models.dummy_entities import DummySubsegment
 from .context import Context
 
 log = logging.getLogger(__name__)
@@ -72,14 +71,10 @@ def put_subsegment(self, subsegment):
         current_entity = self.get_trace_entity()
 
         if not self._is_subsegment(current_entity) and current_entity.initializing:
-            log.warning("Subsegment %s discarded due to Lambda worker still initializing" % subsegment.name)
+            if sdk_config_module.sdk_enabled():
+                log.warning("Subsegment %s discarded due to Lambda worker still initializing" % subsegment.name)
             return
 
-        enabled = SDKConfig.sdk_enabled()
-        if not enabled:
-            # For lambda, if the SDK is not enabled, we force the subsegment to be a dummy segment.
-            subsegment = DummySubsegment(current_entity)
-
         current_entity.add_subsegment(subsegment)
         self._local.entities.append(subsegment)
 
@@ -99,6 +94,9 @@ def _refresh_context(self):
         """
         header_str = os.getenv(LAMBDA_TRACE_HEADER_KEY)
         trace_header = TraceHeader.from_header_str(header_str)
+        if not aws_xray_sdk.global_sdk_config.sdk_enabled():
+            trace_header._sampled = False
+
         segment = getattr(self._local, 'segment', None)
 
         if segment:
@@ -130,7 +128,10 @@ def _initialize_context(self, trace_header):
         set by AWS Lambda and initialize storage for subsegments.
         """
         sampled = None
-        if trace_header.sampled == 0:
+        if not aws_xray_sdk.global_sdk_config.sdk_enabled():
+            # Force subsequent subsegments to be disabled and turned into DummySegments.
+            sampled = False
+        elif trace_header.sampled == 0:
             sampled = False
         elif trace_header.sampled == 1:
             sampled = True

aws_xray_sdk/core/patcher.py

@@ -7,7 +7,7 @@
 import sys
 import wrapt
 
-from aws_xray_sdk.sdk_config import SDKConfig
+import aws_xray_sdk
 from .utils.compat import PY2, is_classmethod, is_instance_method
 
 log = logging.getLogger(__name__)
@@ -61,7 +61,7 @@ def _is_valid_import(module):
 
 
 def patch(modules_to_patch, raise_errors=True, ignore_module_patterns=None):
-    enabled = SDKConfig.sdk_enabled()
+    enabled = aws_xray_sdk.global_sdk_config.sdk_enabled()
     if not enabled:
         return  # Disable module patching if the SDK is disabled.
     modules = set()

aws_xray_sdk/core/recorder.py

@@ -6,7 +6,6 @@
 import time
 
 from aws_xray_sdk.version import VERSION
-from aws_xray_sdk.sdk_config import SDKConfig
 from .models.segment import Segment, SegmentContextManager
 from .models.subsegment import Subsegment, SubsegmentContextManager
 from .models.default_dynamic_naming import DefaultDynamicNaming
@@ -178,11 +177,7 @@ class to have your own implementation of the streaming process.
         if stream_sql is not None:
             self.stream_sql = stream_sql
         if enabled is not None:
-            SDKConfig.set_sdk_enabled(enabled)
-        else:
-            # By default we enable if no enable parameter is given. Prevents unit tests from breaking
-            # if setup doesn't explicitly set enabled while other tests set enabled to false.
-            SDKConfig.set_sdk_enabled(True)
+            self.enabled = enabled
 
         if plugins:
             plugin_modules = get_plugin_modules(plugins)
@@ -241,7 +236,7 @@ def begin_segment(self, name=None, traceid=None,
         # To disable the recorder, we set the sampling decision to always be false.
         # This way, when segments are generated, they become dummy segments and are ultimately never sent.
         # The call to self._sampler.should_trace() is never called either so the poller threads are never started.
-        if not SDKConfig.sdk_enabled():
+        if not self.enabled:
             sampling = 0
 
         # we respect the input sampling decision
@@ -425,7 +420,7 @@ def record_subsegment(self, wrapped, instance, args, kwargs, name,
         # In the case when the SDK is disabled, we ensure that a parent segment exists, because this is usually
         # handled by the middleware. We generate a dummy segment as the parent segment if one doesn't exist.
         # This is to allow potential segment method calls to not throw exceptions in the captured method.
-        if not SDKConfig.sdk_enabled():
+        if not self.enabled:
             try:
                 self.current_segment()
             except SegmentNotFoundException:

aws_xray_sdk/core/sampling/sampler.py

@@ -9,6 +9,7 @@
 from .target_poller import TargetPoller
 from .connector import ServiceConnector
 from .reservoir import ReservoirDecision
+import aws_xray_sdk
 
 log = logging.getLogger(__name__)
 
@@ -37,6 +38,9 @@ def start(self):
         Start rule poller and target poller once X-Ray daemon address
         and context manager is in place.
         """
+        if not aws_xray_sdk.global_sdk_config.sdk_enabled():
+            return
+
         with self._lock:
             if not self._started:
                 self._rule_poller.start()
@@ -51,6 +55,9 @@ def should_trace(self, sampling_req=None):
         All optional arguments are extracted from incoming requests by
         X-Ray middleware to perform path based sampling.
         """
+        if not aws_xray_sdk.global_sdk_config.sdk_enabled():
+            return False
+
         if not self._started:
             self.start() # only front-end that actually uses the sampler spawns poller threads
 

aws_xray_sdk/sdk_config.py

@@ -1,4 +1,12 @@
 import os
+import aws_xray_sdk.core
+
+
+class InvalidParameterTypeException(Exception):
+    """
+    Exception thrown when an invalid parameter is passed into SDKConfig.set_sdk_enabled.
+    """
+    pass
 
 
 class SDKConfig(object):
@@ -7,18 +15,18 @@ class SDKConfig(object):
     It is recommended to only use the recorder to set this configuration's enabled
     flag to maintain thread safety.
     """
-    XRAY_ENABLED_KEY = 'AWS_XRAY_ENABLED'
+    XRAY_ENABLED_KEY = 'AWS_XRAY_SDK_ENABLED'
     __SDK_ENABLED = str(os.getenv(XRAY_ENABLED_KEY, 'true')).lower() != 'false'
 
-    @staticmethod
-    def sdk_enabled():
+    @classmethod
+    def sdk_enabled(cls):
         """
         Returns whether the SDK is enabled or not.
         """
-        return SDKConfig.__SDK_ENABLED
+        return cls.__SDK_ENABLED
 
-    @staticmethod
-    def set_sdk_enabled(value):
+    @classmethod
+    def set_sdk_enabled(cls, value):
         """
         Modifies the enabled flag if the "AWS_XRAY_ENABLED" environment variable is not set,
         otherwise, set the enabled flag to be equal to the environment variable. If the
@@ -29,7 +37,16 @@ def set_sdk_enabled(value):
         Environment variables AWS_XRAY_ENABLED overrides argument value.
         """
         # Environment Variables take precedence over hardcoded configurations.
-        if SDKConfig.XRAY_ENABLED_KEY in os.environ:
-            SDKConfig.__SDK_ENABLED = str(os.getenv(SDKConfig.XRAY_ENABLED_KEY, 'true')).lower() != 'false'
+        if cls.XRAY_ENABLED_KEY in os.environ:
+            cls.__SDK_ENABLED = str(os.getenv(cls.XRAY_ENABLED_KEY, 'true')).lower() != 'false'
         else:
-            SDKConfig.__SDK_ENABLED = value
+            if type(value) == bool:
+                cls.__SDK_ENABLED = value
+            else:
+                cls.__SDK_ENABLED = True
+                raise InvalidParameterTypeException(
+                    "Invalid parameter type passed into set_sdk_enabled(). Defaulting to True..."
+                )
+
+        # Modify all key paths.
+        aws_xray_sdk.core.xray_recorder.configure(enabled=cls.__SDK_ENABLED)

tests/ext/aiohttp/test_middleware.py

@@ -4,6 +4,7 @@
 Expects pytest-aiohttp
 """
 import asyncio
+import aws_xray_sdk.sdk_config_module
 from unittest.mock import patch
 
 from aiohttp import web
@@ -108,7 +109,9 @@ def recorder(loop):
     patcher.start()
 
     xray_recorder.clear_trace_entities()
+    sdk_enabled_state = aws_xray_sdk.global_sdk_config.sdk_enabled()
     yield xray_recorder
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(sdk_enabled_state)
     xray_recorder.clear_trace_entities()
     patcher.stop()
 
@@ -293,7 +296,7 @@ async def test_disabled_sdk(test_client, loop, recorder):
     :param loop: Eventloop fixture
     :param recorder: X-Ray recorder fixture
     """
-    recorder.configure(enabled=False)
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(False)
     client = await test_client(ServerTest.app(loop=loop))
 
     resp = await client.get('/')

tests/ext/django/test_middleware.py

@@ -1,4 +1,5 @@
 import django
+import aws_xray_sdk
 from django.core.urlresolvers import reverse
 from django.test import TestCase
 
@@ -14,6 +15,7 @@ def setUp(self):
         xray_recorder.configure(context=Context(),
                                 context_missing='LOG_ERROR')
         xray_recorder.clear_trace_entities()
+        aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
 
     def tearDown(self):
         xray_recorder.clear_trace_entities()
@@ -104,7 +106,7 @@ def test_response_header(self):
         assert segment.trace_id in trace_header
 
     def test_disabled_sdk(self):
-        xray_recorder.configure(enabled=False)
+        aws_xray_sdk.global_sdk_config.set_sdk_enabled(False)
         url = reverse('200ok')
         self.client.get(url)
         segment = xray_recorder.emitter.pop()

tests/test_lambda_context.py

@@ -1,6 +1,6 @@
 import os
 
-from aws_xray_sdk.sdk_config import SDKConfig
+import aws_xray_sdk
 from aws_xray_sdk.core import lambda_launcher
 from aws_xray_sdk.core.models.subsegment import Subsegment
 from aws_xray_sdk.core.models.dummy_entities import DummySubsegment
@@ -12,6 +12,7 @@
 
 os.environ[lambda_launcher.LAMBDA_TRACE_HEADER_KEY] = HEADER_VAR
 context = lambda_launcher.LambdaContext()
+aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
 
 
 def test_facade_segment_generation():
@@ -46,8 +47,11 @@ def test_put_subsegment():
 
 
 def test_disable():
-    SDKConfig.set_sdk_enabled(False)
+    context.clear_trace_entities()
     segment = context.get_trace_entity()
-    subsegment = Subsegment('name', 'local', segment)
-    context.put_subsegment(subsegment)
-    assert type(context.get_trace_entity()) is DummySubsegment
+    assert segment.sampled
+
+    context.clear_trace_entities()
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(False)
+    segment = context.get_trace_entity()
+    assert not segment.sampled

tests/test_patcher.py

@@ -13,7 +13,7 @@
         # Python versions < 3 have reload built-in
         pass
 
-from aws_xray_sdk.sdk_config import SDKConfig
+import aws_xray_sdk
 from aws_xray_sdk.core import patcher, xray_recorder
 from aws_xray_sdk.core.context import Context
 
@@ -35,12 +35,14 @@ def construct_ctx():
     """
     pre_run_modules = set(module for module in sys.modules.keys())
 
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
     xray_recorder.configure(service='test', sampling=False, context=Context())
     xray_recorder.clear_trace_entities()
     xray_recorder.begin_segment('name')
     yield
     xray_recorder.end_segment()
     xray_recorder.clear_trace_entities()
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(True)
 
     # Reload wrapt.importer references to modules to start off clean
     reload(wrapt)
@@ -176,7 +178,7 @@ def test_external_submodules_ignores_module():
 
 
 def test_disable_sdk_disables_patching():
-    SDKConfig.set_sdk_enabled(False)
+    aws_xray_sdk.global_sdk_config.set_sdk_enabled(False)
     patcher.patch(['tests.mock_module'])
     imported_modules = [module for module in TEST_MODULES if module in sys.modules]
     assert not imported_modules

tests/test_recorder.py

@@ -5,14 +5,13 @@
 from aws_xray_sdk.version import VERSION
 from .util import get_new_stubbed_recorder
 
-import os
-from aws_xray_sdk.sdk_config import SDKConfig
+import aws_xray_sdk
 from aws_xray_sdk.core.models.segment import Segment
 from aws_xray_sdk.core.models.subsegment import Subsegment
 from aws_xray_sdk.core.models.dummy_entities import DummySegment, DummySubsegment
 
 xray_recorder = get_new_stubbed_recorder()
-XRAY_ENABLED_KEY = SDKConfig.XRAY_ENABLED_KEY
+XRAY_ENABLED_KEY = aws_xray_sdk.global_sdk_config.XRAY_ENABLED_KEY
 
 
 @pytest.fixture(autouse=True)
@@ -188,30 +187,3 @@ def test_disable_is_dummy():
     subsegment = xray_recorder.begin_subsegment('name')
     assert type(xray_recorder.current_segment()) is DummySegment
     assert type(xray_recorder.current_subsegment()) is DummySubsegment
-
-
-def test_disable_env_precedence():
-    os.environ[XRAY_ENABLED_KEY] = "False"
-    xray_recorder.configure(enabled=True)
-    segment = xray_recorder.begin_segment('name')
-    subsegment = xray_recorder.begin_subsegment('name')
-    assert type(xray_recorder.current_segment()) is DummySegment
-    assert type(xray_recorder.current_subsegment()) is DummySubsegment
-
-
-def test_disable_env():
-    os.environ[XRAY_ENABLED_KEY] = "False"
-    xray_recorder.configure(enabled=False)
-    segment = xray_recorder.begin_segment('name')
-    subsegment = xray_recorder.begin_subsegment('name')
-    assert type(xray_recorder.current_segment()) is DummySegment
-    assert type(xray_recorder.current_subsegment()) is DummySubsegment
-
-
-def test_enable_env():
-    os.environ[XRAY_ENABLED_KEY] = "True"
-    xray_recorder.configure(enabled=True)
-    segment = xray_recorder.begin_segment('name')
-    subsegment = xray_recorder.begin_subsegment('name')
-    assert type(xray_recorder.current_segment()) is Segment
-    assert type(xray_recorder.current_subsegment()) is Subsegment