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

aws_xray_sdk/core/lambda_launcher.py

@@ -2,11 +2,12 @@
 import logging
 import threading
 
+from aws_xray_sdk.sdk_config import SDKConfig
 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__)
 
 
@@ -74,6 +75,11 @@ def put_subsegment(self, subsegment):
             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)
 

aws_xray_sdk/core/patcher.py

@@ -7,6 +7,7 @@
 import sys
 import wrapt
 
+from aws_xray_sdk.sdk_config import SDKConfig
 from .utils.compat import PY2, is_classmethod, is_instance_method
 
 log = logging.getLogger(__name__)
@@ -60,6 +61,9 @@ def _is_valid_import(module):
 
 
 def patch(modules_to_patch, raise_errors=True, ignore_module_patterns=None):
+    enabled = SDKConfig.sdk_enabled()
+    if not enabled:
+        return  # Disable module patching if the SDK is disabled.
     modules = set()
     for module_to_patch in modules_to_patch:
         # boto3 depends on botocore and patching botocore is sufficient

aws_xray_sdk/core/recorder.py

@@ -6,6 +6,7 @@
 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
@@ -18,12 +19,13 @@
 from .daemon_config import DaemonConfig
 from .plugins.utils import get_plugin_modules
 from .lambda_launcher import check_in_lambda
-from .exceptions.exceptions import SegmentNameMissingException
+from .exceptions.exceptions import SegmentNameMissingException, SegmentNotFoundException
 from .utils.compat import string_types
 from .utils import stacktrace
 
 log = logging.getLogger(__name__)
 
+XRAY_ENABLED_KEY = 'AWS_XRAY_ENABLED'
 TRACING_NAME_KEY = 'AWS_XRAY_TRACING_NAME'
 DAEMON_ADDR_KEY = 'AWS_XRAY_DAEMON_ADDRESS'
 CONTEXT_MISSING_KEY = 'AWS_XRAY_CONTEXT_MISSING'
@@ -64,6 +66,7 @@ def __init__(self):
             self._context = Context()
             self._sampler = DefaultSampler()
 
+        self._enabled = True
         self._emitter = UDPEmitter()
         self._sampling = True
         self._max_trace_back = 10
@@ -77,7 +80,7 @@ def __init__(self):
         if type(self.sampler).__name__ == 'DefaultSampler':
             self.sampler.load_settings(DaemonConfig(), self.context)
 
-    def configure(self, sampling=None, plugins=None,
+    def configure(self, enabled=None, sampling=None, plugins=None,
                   context_missing=None, sampling_rules=None,
                   daemon_address=None, service=None,
                   context=None, emitter=None, streaming=None,
@@ -88,7 +91,16 @@ def configure(self, sampling=None, plugins=None,
 
         Configure needs to run before patching thrid party libraries
         to avoid creating dangling subsegment.
-
+        :param bool enabled: If not enabled, the recorder automatically
+            generates DummySegments for every segment creation, whether
+            through patched extensions nor middlewares, and thus
+            not send any Segments out to the Daemon. May be set through an
+            environmental variable, where the environmental variable will
+            always take precedence over hardcoded configurations. The environment
+            variable is set as a case-insensitive string boolean. If the environment
+            variable exists but is an invalid string boolean, this enabled flag
+            will automatically be set to true. If no enabled flag is given and the
+            env variable is not set, then it will also default to being enabled.
         :param bool sampling: If sampling is enabled, every time the recorder
             creates a segment it decides whether to send this segment to
             the X-Ray daemon. This setting is not used if the recorder
@@ -134,10 +146,11 @@ class to have your own implementation of the streaming process.
             by auto-capture. Lower this if a single document becomes too large.
         :param bool stream_sql: Whether SQL query texts should be streamed.
 
-        Environment variables AWS_XRAY_DAEMON_ADDRESS, AWS_XRAY_CONTEXT_MISSING
+        Environment variables AWS_XRAY_ENABLED, AWS_XRAY_DAEMON_ADDRESS, AWS_XRAY_CONTEXT_MISSING
         and AWS_XRAY_TRACING_NAME respectively overrides arguments
-        daemon_address, context_missing and service.
+        enabled, daemon_address, context_missing and service.
         """
+
         if sampling is not None:
             self.sampling = sampling
         if sampler:
@@ -164,6 +177,12 @@ class to have your own implementation of the streaming process.
             self.max_trace_back = max_trace_back
         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)
 
         if plugins:
             plugin_modules = get_plugin_modules(plugins)
@@ -219,6 +238,12 @@ def begin_segment(self, name=None, traceid=None,
         # depending on if centralized or local sampling rule takes effect.
         decision = True
 
+        # 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():
+            sampling = 0
+
         # we respect the input sampling decision
         # regardless of recorder configuration.
         if sampling == 0:
@@ -273,6 +298,7 @@ def begin_subsegment(self, name, namespace='local'):
         :param str name: the name of the subsegment.
         :param str namespace: currently can only be 'local', 'remote', 'aws'.
         """
+
         segment = self.current_segment()
         if not segment:
             log.warning("No segment found, cannot begin subsegment %s." % name)
@@ -396,6 +422,16 @@ def capture(self, name=None):
     def record_subsegment(self, wrapped, instance, args, kwargs, name,
                           namespace, meta_processor):
 
+        # 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():
+            try:
+                self.current_segment()
+            except SegmentNotFoundException:
+                segment = DummySegment(name)
+                self.context.put_segment(segment)
+
         subsegment = self.begin_subsegment(name, namespace)
 
         exception = None
@@ -473,6 +509,14 @@ def _is_subsegment(self, entity):
 
         return (hasattr(entity, 'type') and entity.type == 'subsegment')
 
+    @property
+    def enabled(self):
+        return self._enabled
+
+    @enabled.setter
+    def enabled(self, value):
+        self._enabled = value
+
     @property
     def sampling(self):
         return self._sampling

aws_xray_sdk/sdk_config.py

@@ -0,0 +1,35 @@
+import os
+
+
+class SDKConfig(object):
+    """
+    Global Configuration Class that defines SDK-level configuration properties.
+    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'
+    __SDK_ENABLED = str(os.getenv(XRAY_ENABLED_KEY, 'true')).lower() != 'false'
+
+    @staticmethod
+    def sdk_enabled():
+        """
+        Returns whether the SDK is enabled or not.
+        """
+        return SDKConfig.__SDK_ENABLED
+
+    @staticmethod
+    def set_sdk_enabled(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
+        env variable is an invalid string boolean, it will default to true.
+
+        :param bool value: Flag to set whether the SDK is enabled or disabled.
+
+        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'
+        else:
+            SDKConfig.__SDK_ENABLED = value

tests/ext/aiohttp/test_middleware.py

@@ -283,3 +283,21 @@ async def get_delay():
     # Ensure all ID's are different
     ids = [item.id for item in recorder.emitter.local]
     assert len(ids) == len(set(ids))
+
+
+async def test_disabled_sdk(test_client, loop, recorder):
+    """
+    Test a normal response when the SDK is disabled.
+
+    :param test_client: AioHttp test client fixture
+    :param loop: Eventloop fixture
+    :param recorder: X-Ray recorder fixture
+    """
+    recorder.configure(enabled=False)
+    client = await test_client(ServerTest.app(loop=loop))
+
+    resp = await client.get('/')
+    assert resp.status == 200
+
+    segment = recorder.emitter.pop()
+    assert not segment

tests/ext/django/test_middleware.py

@@ -102,3 +102,10 @@ def test_response_header(self):
 
         assert 'Sampled=1' in trace_header
         assert segment.trace_id in trace_header
+
+    def test_disabled_sdk(self):
+        xray_recorder.configure(enabled=False)
+        url = reverse('200ok')
+        self.client.get(url)
+        segment = xray_recorder.emitter.pop()
+        assert not segment

tests/ext/flask/test_flask.py

@@ -143,3 +143,11 @@ def test_sampled_response_header():
     resp_header = resp.headers[http.XRAY_HEADER]
     assert segment.trace_id in resp_header
     assert 'Sampled=1' in resp_header
+
+
+def test_disabled_sdk():
+    recorder.configure(enabled=False)
+    path = '/ok'
+    app.get(path)
+    segment = recorder.emitter.pop()
+    assert not segment

tests/test_lambda_context.py

@@ -1,7 +1,9 @@
 import os
 
+from aws_xray_sdk.sdk_config import SDKConfig
 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
 
 
 TRACE_ID = '1-5759e988-bd862e3fe1be46a994272793'
@@ -41,3 +43,11 @@ def test_put_subsegment():
 
     context.end_subsegment()
     assert context.get_trace_entity().id == segment.id
+
+
+def test_disable():
+    SDKConfig.set_sdk_enabled(False)
+    segment = context.get_trace_entity()
+    subsegment = Subsegment('name', 'local', segment)
+    context.put_subsegment(subsegment)
+    assert type(context.get_trace_entity()) is DummySubsegment

tests/test_patcher.py

@@ -13,6 +13,7 @@
         # Python versions < 3 have reload built-in
         pass
 
+from aws_xray_sdk.sdk_config import SDKConfig
 from aws_xray_sdk.core import patcher, xray_recorder
 from aws_xray_sdk.core.context import Context
 
@@ -172,3 +173,11 @@ def test_external_submodules_ignores_module():
     assert xray_recorder.current_segment().subsegments[0].name == 'mock_init'
     assert xray_recorder.current_segment().subsegments[1].name == 'mock_func'
     assert xray_recorder.current_segment().subsegments[2].name == 'mock_no_doublepatch'  # It is patched with decorator
+
+
+def test_disable_sdk_disables_patching():
+    SDKConfig.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
+    assert len(xray_recorder.current_segment().subsegments) == 0

tests/test_recorder.py

@@ -5,7 +5,14 @@
 from aws_xray_sdk.version import VERSION
 from .util import get_new_stubbed_recorder
 
+import os
+from aws_xray_sdk.sdk_config import SDKConfig
+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
 
 
 @pytest.fixture(autouse=True)
@@ -166,3 +173,45 @@ def test_in_segment_exception():
                     raise Exception('test exception')
 
     assert len(subsegment.cause['exceptions']) == 1
+
+
+def test_default_enabled():
+    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
+
+
+def test_disable_is_dummy():
+    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_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