Merge pull request #158 from luzfcb/fix-155

Create DummySegment if SDK is disabled

Author: srprash

aws_xray_sdk/core/context.py

@@ -3,6 +3,9 @@
 import os
 
 from .exceptions.exceptions import SegmentNotFoundException
+from .models.dummy_entities import DummySegment
+from aws_xray_sdk import global_sdk_config
+
 
 log = logging.getLogger(__name__)
 
@@ -88,8 +91,11 @@ def get_trace_entity(self):
         """
         Return the current trace entity(segment/subsegment). If there is none,
         it behaves based on pre-defined ``context_missing`` strategy.
+        If the SDK is disabled, returns a DummySegment
         """
         if not getattr(self._local, 'entities', None):
+            if not global_sdk_config.sdk_enabled():
+                return DummySegment()
             return self.handle_context_missing()
 
         return self._local.entities[-1]

aws_xray_sdk/core/recorder.py

@@ -211,6 +211,10 @@ def begin_segment(self, name=None, traceid=None,
         :param str traceid: trace id of the segment
         :param int sampling: 0 means not sampled, 1 means sampled
         """
+        # Disable the recorder; return a generated dummy segment.
+        if not global_sdk_config.sdk_enabled():
+            return DummySegment(global_sdk_config.DISABLED_ENTITY_NAME)
+
         seg_name = name or self.service
         if not seg_name:
             raise SegmentNameMissingException("Segment name is required.")
@@ -220,12 +224,6 @@ 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 global_sdk_config.sdk_enabled():
-            sampling = 0
-
         # we respect the input sampling decision
         # regardless of recorder configuration.
         if sampling == 0:
@@ -251,8 +249,12 @@ def end_segment(self, end_time=None):
         if it is ready to send. Ready means segment and
         all its subsegments are closed.
 
-        :param float end_time: segment compeletion in unix epoch in seconds.
+        :param float end_time: segment completion in unix epoch in seconds.
         """
+        # When the SDK is disabled we return
+        if not global_sdk_config.sdk_enabled():
+            return
+
         self.context.end_segment(end_time)
         segment = self.current_segment()
         if segment and segment.ready_to_send():
@@ -264,6 +266,7 @@ def current_segment(self):
         this will make sure the segment returned is the one created by the
         same thread.
         """
+
         entity = self.get_trace_entity()
         if self._is_subsegment(entity):
             return entity.parent_segment
@@ -280,6 +283,11 @@ 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'.
         """
+        # Generating the parent dummy segment is necessary.
+        # We don't need to store anything in context. Assumption here
+        # is that we only work with recorder-level APIs.
+        if not global_sdk_config.sdk_enabled():
+            return DummySubsegment(DummySegment(global_sdk_config.DISABLED_ENTITY_NAME))
 
         segment = self.current_segment()
         if not segment:
@@ -301,6 +309,9 @@ def current_subsegment(self):
         this will make sure the subsegment returned is one created
         by the same thread.
         """
+        if not global_sdk_config.sdk_enabled():
+            return DummySubsegment(DummySegment(global_sdk_config.DISABLED_ENTITY_NAME))
+
         entity = self.get_trace_entity()
         if self._is_subsegment(entity):
             return entity
@@ -314,6 +325,9 @@ def end_subsegment(self, end_time=None):
 
         :param float end_time: subsegment compeletion in unix epoch in seconds.
         """
+        if not global_sdk_config.sdk_enabled():
+            return
+
         if not self.context.end_subsegment(end_time):
             return
 
@@ -333,6 +347,8 @@ def put_annotation(self, key, value):
         :param object value: annotation value. Any type other than
             string/number/bool will be dropped
         """
+        if not global_sdk_config.sdk_enabled():
+            return
         entity = self.get_trace_entity()
         if entity and entity.sampled:
             entity.put_annotation(key, value)
@@ -348,6 +364,8 @@ def put_metadata(self, key, value, namespace='default'):
         :param str key: metadata key under specified namespace
         :param object value: any object that can be serialized into JSON string
         """
+        if not global_sdk_config.sdk_enabled():
+            return
         entity = self.get_trace_entity()
         if entity and entity.sampled:
             entity.put_metadata(key, value, namespace)
@@ -357,6 +375,9 @@ def is_sampled(self):
         Check if the current trace entity is sampled or not.
         Return `False` if no active entity found.
         """
+        if not global_sdk_config.sdk_enabled():
+            # Disabled SDK is never sampled
+            return False
         entity = self.get_trace_entity()
         if entity:
             return entity.sampled
@@ -404,16 +425,6 @@ 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 global_sdk_config.sdk_enabled():
-            try:
-                self.current_segment()
-            except SegmentNotFoundException:
-                segment = DummySegment(name)
-                self.context.put_segment(segment)
-
         subsegment = self.begin_subsegment(name, namespace)
 
         exception = None

aws_xray_sdk/sdk_config.py

@@ -1,5 +1,6 @@
 import os
 import logging
+from distutils.util import strtobool
 
 log = logging.getLogger(__name__)
 
@@ -27,13 +28,32 @@ class SDKConfig(object):
             to occur.
     """
     XRAY_ENABLED_KEY = 'AWS_XRAY_SDK_ENABLED'
-    __SDK_ENABLED = str(os.getenv(XRAY_ENABLED_KEY, 'true')).lower() != 'false'
+    DISABLED_ENTITY_NAME = 'dummy'
+
+    __SDK_ENABLED = None
+
+    @classmethod
+    def __get_enabled_from_env(cls):
+        """
+        Searches for the environment variable to see if the SDK should be disabled.
+        If no environment variable is found, it returns True by default.
+
+        :return: bool - True if it is enabled, False otherwise.
+        """
+        env_var_str = os.getenv(cls.XRAY_ENABLED_KEY, 'true')
+        try:
+            return bool(strtobool(env_var_str))
+        except ValueError:
+            log.warning("Invalid literal passed into environment variable `AWS_XRAY_SDK_ENABLED`. Defaulting to True...")
+            return True  # If an invalid parameter is passed in, we return True.
 
     @classmethod
     def sdk_enabled(cls):
         """
         Returns whether the SDK is enabled or not.
         """
+        if cls.__SDK_ENABLED is None:
+            cls.__SDK_ENABLED = cls.__get_enabled_from_env()
         return cls.__SDK_ENABLED
 
     @classmethod
@@ -49,7 +69,7 @@ def set_sdk_enabled(cls, value):
         """
         # Environment Variables take precedence over hardcoded configurations.
         if cls.XRAY_ENABLED_KEY in os.environ:
-            cls.__SDK_ENABLED = str(os.getenv(cls.XRAY_ENABLED_KEY, 'true')).lower() != 'false'
+            cls.__SDK_ENABLED = cls.__get_enabled_from_env()
         else:
             if type(value) == bool:
                 cls.__SDK_ENABLED = value

tests/test_recorder.py

@@ -201,6 +201,65 @@ def test_disable_is_dummy():
     assert type(xray_recorder.current_subsegment()) is DummySubsegment
 
 
+def test_disabled_empty_context_current_calls():
+    global_sdk_config.set_sdk_enabled(False)
+    assert type(xray_recorder.current_segment()) is DummySegment
+    assert type(xray_recorder.current_subsegment()) is DummySubsegment
+
+
+def test_disabled_out_of_order_begins():
+    global_sdk_config.set_sdk_enabled(False)
+    xray_recorder.begin_subsegment("Test")
+    xray_recorder.begin_segment("Test")
+    xray_recorder.begin_subsegment("Test1")
+    xray_recorder.begin_subsegment("Test2")
+    assert type(xray_recorder.begin_subsegment("Test3")) is DummySubsegment
+    assert type(xray_recorder.begin_segment("Test4")) is DummySegment
+
+
+def test_disabled_put_methods():
+    global_sdk_config.set_sdk_enabled(False)
+    xray_recorder.put_annotation("Test", "Value")
+    xray_recorder.put_metadata("Test", "Value", "Namespace")
+
+
+# Test for random end segments/subsegments without any entities in context.
+# Should not throw any exceptions
+def test_disabled_ends():
+    global_sdk_config.set_sdk_enabled(False)
+    xray_recorder.end_segment()
+    xray_recorder.end_subsegment()
+    xray_recorder.end_segment()
+    xray_recorder.end_segment()
+    xray_recorder.end_subsegment()
+    xray_recorder.end_subsegment()
+
+
+# Begin subsegment should not fail on its own.
+def test_disabled_begin_subsegment():
+    global_sdk_config.set_sdk_enabled(False)
+    subsegment_entity = xray_recorder.begin_subsegment("Test")
+    assert type(subsegment_entity) is DummySubsegment
+
+
+# When disabled, force sampling should still return dummy entities.
+def test_disabled_force_sampling():
+    global_sdk_config.set_sdk_enabled(False)
+    xray_recorder.configure(sampling=True)
+    segment_entity = xray_recorder.begin_segment("Test1")
+    subsegment_entity = xray_recorder.begin_subsegment("Test2")
+    assert type(segment_entity) is DummySegment
+    assert type(subsegment_entity) is DummySubsegment
+
+
+# When disabled, get_trace_entity should return DummySegment if an entity is not present in the context
+def test_disabled_get_context_entity():
+    global_sdk_config.set_sdk_enabled(False)
+    entity = xray_recorder.get_trace_entity()
+    assert type(entity) is DummySegment
+
+
+
 def test_max_stack_trace_zero():
     xray_recorder.configure(max_trace_back=1)
     with pytest.raises(Exception):

tests/test_sdk_config.py

@@ -56,6 +56,18 @@ def test_env_enable_case():
     global_sdk_config.set_sdk_enabled(True)
     assert global_sdk_config.sdk_enabled() is True
 
+    os.environ[XRAY_ENABLED_KEY] = "1"
+    global_sdk_config.set_sdk_enabled(True)
+    assert global_sdk_config.sdk_enabled() is True
+
+    os.environ[XRAY_ENABLED_KEY] = "y"
+    global_sdk_config.set_sdk_enabled(True)
+    assert global_sdk_config.sdk_enabled() is True
+
+    os.environ[XRAY_ENABLED_KEY] = "t"
+    global_sdk_config.set_sdk_enabled(True)
+    assert global_sdk_config.sdk_enabled() is True
+
     os.environ[XRAY_ENABLED_KEY] = "False"
     global_sdk_config.set_sdk_enabled(True)
     assert global_sdk_config.sdk_enabled() is False
@@ -64,6 +76,10 @@ def test_env_enable_case():
     global_sdk_config.set_sdk_enabled(True)
     assert global_sdk_config.sdk_enabled() is False
 
+    os.environ[XRAY_ENABLED_KEY] = "0"
+    global_sdk_config.set_sdk_enabled(True)
+    assert global_sdk_config.sdk_enabled() is False
+
 
 def test_invalid_env_string():
     os.environ[XRAY_ENABLED_KEY] = "INVALID"
@@ -78,3 +94,7 @@ def test_invalid_env_string():
     os.environ[XRAY_ENABLED_KEY] = "1-.0"
     global_sdk_config.set_sdk_enabled(False)
     assert global_sdk_config.sdk_enabled() is True
+
+    os.environ[XRAY_ENABLED_KEY] = "T RUE"
+    global_sdk_config.set_sdk_enabled(True)
+    assert global_sdk_config.sdk_enabled() is True