Merge branch 'master' into async

Author: willarmiros

.github/workflows/UnitTesting.yaml

@@ -36,7 +36,7 @@ jobs:
           python-version: ${{ env[matrix.python-version] }}
 
       - name: Install tox
-        run: pip install -U tox-factor
+        run: pip install "tox<=3.27.1" -U tox-factor
 
       - name: Cache tox environment
         # Preserves .tox directory between runs for faster installs

CHANGELOG.rst

@@ -5,6 +5,20 @@ CHANGELOG
 Unreleased
 ==========
 
+2.11.0
+==========
+* bugfix: Fix TypeError by patching register_default_jsonb from psycopg2 `PR350 https://github.com/aws/aws-xray-sdk-python/pull/350`
+* improvement: Add annotations `PR348 https://github.com/aws/aws-xray-sdk-python/pull/348`
+* bugfix: Use service parameter to match centralized sampling rules `PR 353 https://github.com/aws/aws-xray-sdk-python/pull/353`
+* bugfix: Implement PEP3134 to discover underlying problems with python3 `PR355 https://github.com/aws/aws-xray-sdk-python/pull/355`
+* improvement: Allow list TopicArn for SNS PublishBatch request `PR358 https://github.com/aws/aws-xray-sdk-python/pull/358`
+* bugfix: Version pinning flask-sqlalchemy version to 2.5.1 or less `PR360 https://github.com/aws/aws-xray-sdk-python/pull/360`
+* bugfix: Fix UnboundLocalError when aiohttp server raises a CancelledError `PR356 https://github.com/aws/aws-xray-sdk-python/pull/356`
+* improvement: Instrument httpx >= 0.20 `PR357 https://github.com/aws/aws-xray-sdk-python/pull/357`
+* improvement: [LambdaContext] persist original trace header `PR362 https://github.com/aws/aws-xray-sdk-python/pull/362`
+* bugfix: Run tests against Django 4.x `PR361 https://github.com/aws/aws-xray-sdk-python/pull/361`
+* improvement: Oversampling Mitigation `PR366 https://github.com/aws/aws-xray-sdk-python/pull/366`
+
 2.10.0
 ==========
 * bugfix: Only import future for py2. `PR343 <https://github.com/aws/aws-xray-sdk-python/pull/343>`_.

README.md

@@ -124,6 +124,46 @@ xray_recorder.end_subsegment()
 xray_recorder.end_segment()
 ```
 
+### Oversampling Mitigation
+To modify the sampling decision at the subsegment level, subsegments that inherit the decision of their direct parent (segment or subsegment) can be created using `xray_recorder.begin_subsegment()` and unsampled subsegments can be created using
+`xray_recorder.begin_subsegment_without_sampling()`.
+
+The code snippet below demonstrates creating a sampled or unsampled subsegment based on the sampling decision of each SQS message processed by Lambda.
+
+```python
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.models.subsegment import Subsegment
+from aws_xray_sdk.core.utils.sqs_message_helper import SqsMessageHelper
+
+def lambda_handler(event, context):
+
+    for message in event['Records']:
+        if SqsMessageHelper.isSampled(message):
+            subsegment = xray_recorder.begin_subsegment('sampled_subsegment')
+            print('sampled - processing SQS message')
+
+        else:
+            subsegment = xray_recorder.begin_subsegment_without_sampling('unsampled_subsegment')
+            print('unsampled - processing SQS message')
+    
+    xray_recorder.end_subsegment()   
+```
+
+The code snippet below demonstrates wrapping a downstream AWS SDK request with an unsampled subsegment.
+```python
+from aws_xray_sdk.core import xray_recorder, patch_all
+import boto3
+
+patch_all()
+
+def lambda_handler(event, context):
+    subsegment = xray_recorder.begin_subsegment_without_sampling('unsampled_subsegment')
+    client = boto3.client('sqs')
+    print(client.list_queues())
+    
+    xray_recorder.end_subsegment()
+```
+
 ### Capture
 
 As a decorator:

aws_xray_sdk/core/async_recorder.py

@@ -1,5 +1,4 @@
 import time
-import six
 
 from aws_xray_sdk.core.recorder import AWSXRayRecorder
 from aws_xray_sdk.core.utils import stacktrace
@@ -82,10 +81,10 @@ async def record_subsegment_async(self, wrapped, instance, args, kwargs, name,
         try:
             return_value = await wrapped(*args, **kwargs)
             return return_value
-        except Exception as exc:
-            exception = exc
+        except Exception as e:
+            exception = e
             stack = stacktrace.get_stacktrace(limit=self._max_trace_back)
-            six.raise_from(exc, exc)
+            raise
         finally:
             # No-op if subsegment is `None` due to `LOG_ERROR`.
             if subsegment is not None:

aws_xray_sdk/core/context.py

@@ -27,7 +27,7 @@ class Context(object):
 
     This data structure is thread-safe.
     """
-    def __init__(self, context_missing='RUNTIME_ERROR'):
+    def __init__(self, context_missing='LOG_ERROR'):
 
         self._local = threading.local()
         strategy = os.getenv(CXT_MISSING_STRATEGY_KEY, context_missing)

aws_xray_sdk/core/models/dummy_entities.py

@@ -11,7 +11,7 @@ class DummySegment(Segment):
     the segment based on sampling rules.
     Adding data to a dummy segment becomes a no-op except for
     subsegments. This is to reduce the memory footprint of the SDK.
-    A dummy segment will not be sent to the X-Ray daemon. Manually create
+    A dummy segment will not be sent to the X-Ray daemon. Manually creating
     dummy segments is not recommended.
     """
 

aws_xray_sdk/core/models/entity.py

@@ -81,6 +81,10 @@ def add_subsegment(self, subsegment):
         """
         self._check_ended()
         subsegment.parent_id = self.id
+
+        if not self.sampled and subsegment.sampled:
+            log.warning("This sampled subsegment is being added to an unsampled parent segment/subsegment and will be orphaned.")
+
         self.subsegments.append(subsegment)
 
     def remove_subsegment(self, subsegment):

aws_xray_sdk/core/patcher.py

@@ -6,7 +6,6 @@
 import re
 import sys
 import wrapt
-import six
 
 from aws_xray_sdk import global_sdk_config
 from .utils.compat import PY2, is_classmethod, is_instance_method
@@ -110,9 +109,9 @@ def patch(modules_to_patch, raise_errors=True, ignore_module_patterns=None):
 def _patch_module(module_to_patch, raise_errors=True):
     try:
         _patch(module_to_patch)
-    except Exception as exc:
+    except Exception:
         if raise_errors:
-            six.raise_from(exc, exc)
+            raise
         log.debug('failed to patch module %s', module_to_patch)
 
 

aws_xray_sdk/core/recorder.py

@@ -4,7 +4,6 @@
 import os
 import platform
 import time
-import six
 
 from aws_xray_sdk import global_sdk_config
 from aws_xray_sdk.version import VERSION
@@ -275,16 +274,10 @@ def current_segment(self):
         else:
             return entity
 
-    def begin_subsegment(self, name, namespace='local'):
-        """
-        Begin a new subsegment.
-        If there is open subsegment, the newly created subsegment will be the
-        child of latest opened subsegment.
-        If not, it will be the child of the current open segment.
-
-        :param str name: the name of the subsegment.
-        :param str namespace: currently can only be 'local', 'remote', 'aws'.
-        """
+    def _begin_subsegment_helper(self, name, namespace='local', beginWithoutSampling=False):
+        '''
+        Helper method to begin_subsegment and begin_subsegment_without_sampling
+        '''
         # 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.
@@ -295,16 +288,42 @@ def begin_subsegment(self, name, namespace='local'):
         if not segment:
             log.warning("No segment found, cannot begin subsegment %s." % name)
             return None
-
-        if not segment.sampled:
+        
+        current_entity = self.get_trace_entity()
+        if not current_entity.sampled or beginWithoutSampling:
             subsegment = DummySubsegment(segment, name)
         else:
             subsegment = Subsegment(name, namespace, segment)
 
         self.context.put_subsegment(subsegment)
-
         return subsegment
 
+
+
+    def begin_subsegment(self, name, namespace='local'):
+        """
+        Begin a new subsegment.
+        If there is open subsegment, the newly created subsegment will be the
+        child of latest opened subsegment.
+        If not, it will be the child of the current open segment.
+
+        :param str name: the name of the subsegment.
+        :param str namespace: currently can only be 'local', 'remote', 'aws'.
+        """
+        return self._begin_subsegment_helper(name, namespace)
+
+
+    def begin_subsegment_without_sampling(self, name):
+        """
+        Begin a new unsampled subsegment.
+        If there is open subsegment, the newly created subsegment will be the
+        child of latest opened subsegment.
+        If not, it will be the child of the current open segment.
+
+        :param str name: the name of the subsegment.
+        """
+        return self._begin_subsegment_helper(name, beginWithoutSampling=True)
+
     def current_subsegment(self):
         """
         Return the latest opened subsegment. In a multithreading environment,
@@ -436,10 +455,10 @@ def record_subsegment(self, wrapped, instance, args, kwargs, name,
         try:
             return_value = wrapped(*args, **kwargs)
             return return_value
-        except Exception as exc:
-            exception = exc
+        except Exception as e:
+            exception = e
             stack = stacktrace.get_stacktrace(limit=self.max_trace_back)
-            six.raise_from(exc, exc)
+            raise
         finally:
             # No-op if subsegment is `None` due to `LOG_ERROR`.
             if subsegment is not None:
@@ -487,7 +506,8 @@ def _send_segment(self):
 
     def _stream_subsegment_out(self, subsegment):
         log.debug("streaming subsegments...")
-        self.emitter.send_entity(subsegment)
+        if subsegment.sampled:
+            self.emitter.send_entity(subsegment)
 
     def _load_sampling_rules(self, sampling_rules):
 

aws_xray_sdk/core/utils/sqs_message_helper.py

@@ -0,0 +1,11 @@
+SQS_XRAY_HEADER = "AWSTraceHeader"
+class SqsMessageHelper:
+    
+    @staticmethod 
+    def isSampled(sqs_message):
+        attributes = sqs_message['attributes']
+
+        if SQS_XRAY_HEADER not in attributes:
+            return False
+
+        return 'Sampled=1' in attributes[SQS_XRAY_HEADER]