[XRay Sampler] Fix - Ensure all XRay Sampler functionality is under ParentBased logic (#269)

### Issue #, if available:

This PR fixes
aws-observability/aws-otel-js-instrumentation#79,
but in Python
This fixes the bug where:
- While the subcomponents of XRay Sampler uses ParentBased logic to
short-circuit the Parent's Sampling Decision, the XRay Sampling
Statistics recording logic is not skipped when a Parent's Sampling
Decision is found. This causes XRay Sampler to produce Sampling
Statistics based on number of Spans (regardless of Parent's Sampling
Decision), while XRay Sampler should produce Sampling Statistics based
on number of Traces (aka the number of root Spans that makes the
Sampling Decision)

###  Description of changes:
1. Wrap entire XRay Sampler Internal Logic under a single ParentBased
Sampler
a. This will reduce lock contention and not over-count sampling
statistics.
3. Remove use of redundant ParentBased Sampler logic for internal
subcomponents for XRay Sampler.

### Testing:
1. Added unit tests to verify ParentBased sampling decisions
2. Original Sampler functionality tested locally using [Sampler Test
Bed](https://github.com/aws-observability/aws-otel-community/tree/master/centralized-sampling-tests)

By submitting this pull request, I confirm that you can use, modify,
copy, and redistribute this contribution, under the terms of your
choice.

Co-authored-by: Ping Xiang <[email protected]>

Author: jj22ee

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/sampler/_sampling_rule_applier.py

@@ -12,7 +12,7 @@
 from amazon.opentelemetry.distro.sampler._sampling_target import _SamplingTarget
 from opentelemetry.context import Context
 from opentelemetry.sdk.resources import Resource
-from opentelemetry.sdk.trace.sampling import Decision, ParentBased, Sampler, SamplingResult, TraceIdRatioBased
+from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult, TraceIdRatioBased
 from opentelemetry.semconv.resource import CloudPlatformValues, ResourceAttributes
 from opentelemetry.semconv.trace import SpanAttributes
 from opentelemetry.trace import Link, SpanKind
@@ -42,7 +42,7 @@ def __init__(
         self.__borrowing = False
 
         if target is None:
-            self.__fixed_rate_sampler = ParentBased(TraceIdRatioBased(self.sampling_rule.FixedRate))
+            self.__fixed_rate_sampler = TraceIdRatioBased(self.sampling_rule.FixedRate)
             # Until targets are fetched, initialize as borrowing=True if there will be a quota > 0
             if self.sampling_rule.ReservoirSize > 0:
                 self.__reservoir_sampler = self.__create_reservoir_sampler(quota=1)
@@ -55,7 +55,7 @@ def __init__(
             new_quota = target.ReservoirQuota if target.ReservoirQuota is not None else 0
             new_fixed_rate = target.FixedRate if target.FixedRate is not None else 0
             self.__reservoir_sampler = self.__create_reservoir_sampler(quota=new_quota)
-            self.__fixed_rate_sampler = ParentBased(TraceIdRatioBased(new_fixed_rate))
+            self.__fixed_rate_sampler = TraceIdRatioBased(new_fixed_rate)
             if target.ReservoirQuotaTTL is not None:
                 self.__reservoir_expiry = self._clock.from_timestamp(target.ReservoirQuotaTTL)
             else:
@@ -159,7 +159,7 @@ def matches(self, resource: Resource, attributes: Attributes) -> bool:
         )
 
     def __create_reservoir_sampler(self, quota: int) -> Sampler:
-        return ParentBased(_RateLimitingSampler(quota, self._clock))
+        return _RateLimitingSampler(quota, self._clock)
 
     # pylint: disable=no-self-use
     def __get_service_type(self, resource: Resource) -> str:

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/sampler/aws_xray_remote_sampler.py

@@ -24,7 +24,48 @@
 DEFAULT_SAMPLING_PROXY_ENDPOINT = "http://127.0.0.1:2000"
 
 
+# Wrapper class to ensure that all XRay Sampler Functionality in _AwsXRayRemoteSampler
+# uses ParentBased logic to respect the parent span's sampling decision
 class AwsXRayRemoteSampler(Sampler):
+    def __init__(
+        self,
+        resource: Resource,
+        endpoint: str = None,
+        polling_interval: int = None,
+        log_level=None,
+    ):
+        self._root = ParentBased(
+            _AwsXRayRemoteSampler(
+                resource=resource, endpoint=endpoint, polling_interval=polling_interval, log_level=log_level
+            )
+        )
+
+    # pylint: disable=no-self-use
+    @override
+    def should_sample(
+        self,
+        parent_context: Optional[Context],
+        trace_id: int,
+        name: str,
+        kind: SpanKind = None,
+        attributes: Attributes = None,
+        links: Sequence[Link] = None,
+        trace_state: TraceState = None,
+    ) -> SamplingResult:
+        return self._root.should_sample(
+            parent_context, trace_id, name, kind=kind, attributes=attributes, links=links, trace_state=trace_state
+        )
+
+    # pylint: disable=no-self-use
+    @override
+    def get_description(self) -> str:
+        return f"AwsXRayRemoteSampler{{root:{self._root.get_description()}}}"
+
+
+# _AwsXRayRemoteSampler contains all core XRay Sampler Functionality,
+# however it is NOT Parent-based (e.g. Sample logic runs for each span)
+# Not intended for external use, use Parent-based `AwsXRayRemoteSampler` instead.
+class _AwsXRayRemoteSampler(Sampler):
     """
     Remote Sampler for OpenTelemetry that gets sampling configurations from AWS X-Ray
 
@@ -58,7 +99,7 @@ def __init__(
         self.__client_id = self.__generate_client_id()
         self._clock = _Clock()
         self.__xray_client = _AwsXRaySamplingClient(endpoint, log_level=log_level)
-        self.__fallback_sampler = ParentBased(_FallbackSampler(self._clock))
+        self.__fallback_sampler = _FallbackSampler(self._clock)
 
         self.__polling_interval = polling_interval
         self.__target_polling_interval = DEFAULT_TARGET_POLLING_INTERVAL_SECONDS
@@ -114,7 +155,7 @@ def should_sample(
     # pylint: disable=no-self-use
     @override
     def get_description(self) -> str:
-        description = "AwsXRayRemoteSampler{remote sampling with AWS X-Ray}"
+        description = "_AwsXRayRemoteSampler{remote sampling with AWS X-Ray}"
         return description
 
     def __get_and_update_sampling_rules(self) -> None:

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/sampler/test_aws_xray_remote_sampler.py

@@ -10,8 +10,9 @@
 
 from mock_clock import MockClock
 
-from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler
+from amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler import AwsXRayRemoteSampler, _AwsXRayRemoteSampler
 from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import Tracer, TracerProvider
 from opentelemetry.sdk.trace.sampling import Decision
 
 TEST_DIR = os.path.dirname(os.path.realpath(__file__))
@@ -55,27 +56,31 @@ def setUp(self):
     def tearDown(self):
         # Clean up timers
         if self.rs is not None:
-            self.rs._rules_timer.cancel()
-            self.rs._targets_timer.cancel()
+            self.rs._root._root._rules_timer.cancel()
+            self.rs._root._root._targets_timer.cancel()
 
     def test_create_remote_sampler_with_empty_resource(self):
         self.rs = AwsXRayRemoteSampler(resource=Resource.get_empty())
-        self.assertIsNotNone(self.rs._rules_timer)
-        self.assertEqual(self.rs._AwsXRayRemoteSampler__polling_interval, 300)
-        self.assertIsNotNone(self.rs._AwsXRayRemoteSampler__xray_client)
-        self.assertIsNotNone(self.rs._AwsXRayRemoteSampler__resource)
-        self.assertTrue(len(self.rs._AwsXRayRemoteSampler__client_id), 24)
+        self.assertIsNotNone(self.rs._root._root._rules_timer)
+        self.assertEqual(self.rs._root._root._AwsXRayRemoteSampler__polling_interval, 300)
+        self.assertIsNotNone(self.rs._root._root._AwsXRayRemoteSampler__xray_client)
+        self.assertIsNotNone(self.rs._root._root._AwsXRayRemoteSampler__resource)
+        self.assertTrue(len(self.rs._root._root._AwsXRayRemoteSampler__client_id), 24)
 
     def test_create_remote_sampler_with_populated_resource(self):
         self.rs = AwsXRayRemoteSampler(
             resource=Resource.create({"service.name": "test-service-name", "cloud.platform": "test-cloud-platform"})
         )
-        self.assertIsNotNone(self.rs._rules_timer)
-        self.assertEqual(self.rs._AwsXRayRemoteSampler__polling_interval, 300)
-        self.assertIsNotNone(self.rs._AwsXRayRemoteSampler__xray_client)
-        self.assertIsNotNone(self.rs._AwsXRayRemoteSampler__resource)
-        self.assertEqual(self.rs._AwsXRayRemoteSampler__resource.attributes["service.name"], "test-service-name")
-        self.assertEqual(self.rs._AwsXRayRemoteSampler__resource.attributes["cloud.platform"], "test-cloud-platform")
+        self.assertIsNotNone(self.rs._root._root._rules_timer)
+        self.assertEqual(self.rs._root._root._AwsXRayRemoteSampler__polling_interval, 300)
+        self.assertIsNotNone(self.rs._root._root._AwsXRayRemoteSampler__xray_client)
+        self.assertIsNotNone(self.rs._root._root._AwsXRayRemoteSampler__resource)
+        self.assertEqual(
+            self.rs._root._root._AwsXRayRemoteSampler__resource.attributes["service.name"], "test-service-name"
+        )
+        self.assertEqual(
+            self.rs._root._root._AwsXRayRemoteSampler__resource.attributes["cloud.platform"], "test-cloud-platform"
+        )
 
     def test_create_remote_sampler_with_all_fields_populated(self):
         self.rs = AwsXRayRemoteSampler(
@@ -84,35 +89,39 @@ def test_create_remote_sampler_with_all_fields_populated(self):
             polling_interval=120,
             log_level=DEBUG,
         )
-        self.assertIsNotNone(self.rs._rules_timer)
-        self.assertEqual(self.rs._AwsXRayRemoteSampler__polling_interval, 120)
-        self.assertIsNotNone(self.rs._AwsXRayRemoteSampler__xray_client)
-        self.assertIsNotNone(self.rs._AwsXRayRemoteSampler__resource)
+        self.assertIsNotNone(self.rs._root._root._rules_timer)
+        self.assertEqual(self.rs._root._root._AwsXRayRemoteSampler__polling_interval, 120)
+        self.assertIsNotNone(self.rs._root._root._AwsXRayRemoteSampler__xray_client)
+        self.assertIsNotNone(self.rs._root._root._AwsXRayRemoteSampler__resource)
         self.assertEqual(
-            self.rs._AwsXRayRemoteSampler__xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint,
+            self.rs._root._root._AwsXRayRemoteSampler__xray_client._AwsXRaySamplingClient__get_sampling_rules_endpoint,
             "http://abc.com/GetSamplingRules",
         )
-        self.assertEqual(self.rs._AwsXRayRemoteSampler__resource.attributes["service.name"], "test-service-name")
-        self.assertEqual(self.rs._AwsXRayRemoteSampler__resource.attributes["cloud.platform"], "test-cloud-platform")
+        self.assertEqual(
+            self.rs._root._root._AwsXRayRemoteSampler__resource.attributes["service.name"], "test-service-name"
+        )
+        self.assertEqual(
+            self.rs._root._root._AwsXRayRemoteSampler__resource.attributes["cloud.platform"], "test-cloud-platform"
+        )
 
     @patch("requests.Session.post", side_effect=mocked_requests_get)
     @patch("amazon.opentelemetry.distro.sampler.aws_xray_remote_sampler.DEFAULT_TARGET_POLLING_INTERVAL_SECONDS", 2)
     def test_update_sampling_rules_and_targets_with_pollers_and_should_sample(self, mock_post=None):
         self.rs = AwsXRayRemoteSampler(
             resource=Resource.create({"service.name": "test-service-name", "cloud.platform": "test-cloud-platform"})
         )
-        self.assertEqual(self.rs._AwsXRayRemoteSampler__target_polling_interval, 2)
+        self.assertEqual(self.rs._root._root._AwsXRayRemoteSampler__target_polling_interval, 2)
 
         time.sleep(1.0)
         self.assertEqual(
-            self.rs._AwsXRayRemoteSampler__rule_cache._RuleCache__rule_appliers[0].sampling_rule.RuleName,
+            self.rs._root._root._AwsXRayRemoteSampler__rule_cache._RuleCache__rule_appliers[0].sampling_rule.RuleName,
             "test",
         )
         self.assertEqual(self.rs.should_sample(None, 0, "name", attributes={"abc": "1234"}).decision, Decision.DROP)
 
         # wait 2 more seconds since targets polling was patched to 2 seconds (rather than 10s)
         time.sleep(2.0)
-        self.assertEqual(self.rs._AwsXRayRemoteSampler__target_polling_interval, 1000)
+        self.assertEqual(self.rs._root._root._AwsXRayRemoteSampler__target_polling_interval, 1000)
         self.assertEqual(
             self.rs.should_sample(None, 0, "name", attributes={"abc": "1234"}).decision,
             Decision.RECORD_AND_SAMPLE,
@@ -162,9 +171,9 @@ def test_multithreading_with_large_reservoir_with_otel_sdk(self, mock_post=None)
             threads[idx].join()
             sum_sampled += sampled_array[idx]
 
-        test_rule_applier = self.rs._AwsXRayRemoteSampler__rule_cache._RuleCache__rule_appliers[0]
+        test_rule_applier = self.rs._root._root._AwsXRayRemoteSampler__rule_cache._RuleCache__rule_appliers[0]
         self.assertEqual(
-            test_rule_applier._SamplingRuleApplier__reservoir_sampler._root._RateLimitingSampler__reservoir._quota,
+            test_rule_applier._SamplingRuleApplier__reservoir_sampler._RateLimitingSampler__reservoir._quota,
             100000,
         )
         self.assertEqual(sum_sampled, 100000)
@@ -183,19 +192,19 @@ def test_multithreading_with_some_reservoir_with_otel_sdk(self, mock_post=None):
         # which will eat up more than 1 second of reservoir. Using MockClock we can freeze time
         # and pretend all thread jobs start and end at the exact same time,
         # assume and test exactly 1 second of reservoir (100 quota) only
-        mock_clock: MockClock = self.rs._clock
+        mock_clock: MockClock = self.rs._root._root._clock
 
         time.sleep(1.0)
         mock_clock.add_time(1.0)
-        self.assertEqual(mock_clock.now(), self.rs._clock.now())
+        self.assertEqual(mock_clock.now(), self.rs._root._root._clock.now())
         self.assertEqual(
             self.rs.should_sample(None, 0, "name", attributes=attributes).decision, Decision.RECORD_AND_SAMPLE
         )
 
         # wait 2 more seconds since targets polling was patched to 2 seconds (rather than 10s)
         time.sleep(2.0)
         mock_clock.add_time(2.0)
-        self.assertEqual(mock_clock.now(), self.rs._clock.now())
+        self.assertEqual(mock_clock.now(), self.rs._root._root._clock.now())
 
         number_of_spans = 100
         thread_count = 1000
@@ -219,9 +228,79 @@ def test_multithreading_with_some_reservoir_with_otel_sdk(self, mock_post=None):
             threads[idx].join()
             sum_sampled += sampled_array[idx]
 
-        default_rule_applier = self.rs._AwsXRayRemoteSampler__rule_cache._RuleCache__rule_appliers[1]
+        default_rule_applier = self.rs._root._root._AwsXRayRemoteSampler__rule_cache._RuleCache__rule_appliers[1]
         self.assertEqual(
-            default_rule_applier._SamplingRuleApplier__reservoir_sampler._root._RateLimitingSampler__reservoir._quota,
+            default_rule_applier._SamplingRuleApplier__reservoir_sampler._RateLimitingSampler__reservoir._quota,
             100,
         )
         self.assertEqual(sum_sampled, 100)
+
+    def test_get_description(self) -> str:
+        self.rs: AwsXRayRemoteSampler = AwsXRayRemoteSampler(resource=Resource.create({"service.name": "dummy_name"}))
+        self.assertEqual(
+            self.rs.get_description(),
+            "AwsXRayRemoteSampler{root:ParentBased{root:_AwsXRayRemoteSampler{remote sampling with AWS X-Ray},remoteParentSampled:AlwaysOnSampler,remoteParentNotSampled:AlwaysOffSampler,localParentSampled:AlwaysOnSampler,localParentNotSampled:AlwaysOffSampler}}",  # noqa: E501
+        )
+
+    @patch("requests.Session.post", side_effect=mocked_requests_get)
+    def test_parent_based_xray_sampler_updates_statistics_once_for_one_parent_span_with_two_children(
+        self, mock_post=None
+    ):
+        self.rs: AwsXRayRemoteSampler = AwsXRayRemoteSampler(
+            resource=Resource.create({"service.name": "use-default-sample-all-rule"})
+        )
+        time.sleep(1.0)
+
+        provider = TracerProvider(sampler=self.rs)
+        tracer: Tracer = provider.get_tracer("test_tracer_1")
+
+        # child1 and child2 are child spans of root parent0
+        # For AwsXRayRemoteSampler (ParentBased), expect only parent0 to update statistics
+        with tracer.start_as_current_span("parent0") as _:
+            with tracer.start_as_current_span("child1") as _:
+                pass
+            with tracer.start_as_current_span("child2") as _:
+                pass
+        default_rule_applier = self.rs._root._root._AwsXRayRemoteSampler__rule_cache._RuleCache__rule_appliers[1]
+        self.assertEqual(
+            default_rule_applier._SamplingRuleApplier__statistics.RequestCount,
+            1,
+        )
+        self.assertEqual(
+            default_rule_applier._SamplingRuleApplier__statistics.SampleCount,
+            1,
+        )
+
+    @patch("requests.Session.post", side_effect=mocked_requests_get)
+    def test_non_parent_based_xray_sampler_updates_statistics_thrice_for_one_parent_span_with_two_children(
+        self, mock_post=None
+    ):
+        non_parent_based_xray_sampler: _AwsXRayRemoteSampler = _AwsXRayRemoteSampler(
+            resource=Resource.create({"service.name": "use-default-sample-all-rule"})
+        )
+        time.sleep(1.0)
+
+        provider = TracerProvider(sampler=non_parent_based_xray_sampler)
+        tracer: Tracer = provider.get_tracer("test_tracer_2")
+
+        # child1 and child2 are child spans of root parent0
+        # For _AwsXRayRemoteSampler (Non-ParentBased), expect all 3 spans to update statistics
+        with tracer.start_as_current_span("parent0") as _:
+            with tracer.start_as_current_span("child1") as _:
+                pass
+            with tracer.start_as_current_span("child2") as _:
+                pass
+        default_rule_applier = (
+            non_parent_based_xray_sampler._AwsXRayRemoteSampler__rule_cache._RuleCache__rule_appliers[1]
+        )
+        self.assertEqual(
+            default_rule_applier._SamplingRuleApplier__statistics.RequestCount,
+            3,
+        )
+        self.assertEqual(
+            default_rule_applier._SamplingRuleApplier__statistics.SampleCount,
+            3,
+        )
+
+        non_parent_based_xray_sampler._rules_timer.cancel()
+        non_parent_based_xray_sampler._targets_timer.cancel()

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/sampler/test_rule_cache.py

@@ -136,15 +136,15 @@ def test_update_sampling_targets(self):
         # quota should be 1 because of borrowing=true until targets are updated
         rule_applier_0 = rule_cache._RuleCache__rule_appliers[0]
         self.assertEqual(
-            rule_applier_0._SamplingRuleApplier__reservoir_sampler._root._RateLimitingSampler__reservoir._quota, 1
+            rule_applier_0._SamplingRuleApplier__reservoir_sampler._RateLimitingSampler__reservoir._quota, 1
         )
-        self.assertEqual(rule_applier_0._SamplingRuleApplier__fixed_rate_sampler._root._rate, sampling_rule_2.FixedRate)
+        self.assertEqual(rule_applier_0._SamplingRuleApplier__fixed_rate_sampler._rate, sampling_rule_2.FixedRate)
 
         rule_applier_1 = rule_cache._RuleCache__rule_appliers[1]
         self.assertEqual(
-            rule_applier_1._SamplingRuleApplier__reservoir_sampler._root._RateLimitingSampler__reservoir._quota, 1
+            rule_applier_1._SamplingRuleApplier__reservoir_sampler._RateLimitingSampler__reservoir._quota, 1
         )
-        self.assertEqual(rule_applier_1._SamplingRuleApplier__fixed_rate_sampler._root._rate, sampling_rule_1.FixedRate)
+        self.assertEqual(rule_applier_1._SamplingRuleApplier__fixed_rate_sampler._rate, sampling_rule_1.FixedRate)
 
         target_1 = {
             "FixedRate": 0.05,
@@ -179,17 +179,17 @@ def test_update_sampling_targets(self):
         # borrowing=false, use quota from targets
         rule_applier_0 = rule_cache._RuleCache__rule_appliers[0]
         self.assertEqual(
-            rule_applier_0._SamplingRuleApplier__reservoir_sampler._root._RateLimitingSampler__reservoir._quota,
+            rule_applier_0._SamplingRuleApplier__reservoir_sampler._RateLimitingSampler__reservoir._quota,
             target_2["ReservoirQuota"],
         )
-        self.assertEqual(rule_applier_0._SamplingRuleApplier__fixed_rate_sampler._root._rate, target_2["FixedRate"])
+        self.assertEqual(rule_applier_0._SamplingRuleApplier__fixed_rate_sampler._rate, target_2["FixedRate"])
 
         rule_applier_1 = rule_cache._RuleCache__rule_appliers[1]
         self.assertEqual(
-            rule_applier_1._SamplingRuleApplier__reservoir_sampler._root._RateLimitingSampler__reservoir._quota,
+            rule_applier_1._SamplingRuleApplier__reservoir_sampler._RateLimitingSampler__reservoir._quota,
             target_1["ReservoirQuota"],
         )
-        self.assertEqual(rule_applier_1._SamplingRuleApplier__fixed_rate_sampler._root._rate, target_1["FixedRate"])
+        self.assertEqual(rule_applier_1._SamplingRuleApplier__fixed_rate_sampler._rate, target_1["FixedRate"])
 
         # Test target response modified after Rule cache's last modified date
         target_response.LastRuleModification = mock_clock.now().timestamp() + 1