remote sampling - initial classes and rules poller

Author: jj22ee

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

@@ -0,0 +1,79 @@
+# Copyright The OpenTelemetry Authors
+#
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import json
+from logging import getLogger
+from threading import Timer
+from typing import Optional, Sequence
+
+from amazon.opentelemetry.distro.sampler.aws_xray_sampling_client import AwsXRaySamplingClient
+from opentelemetry.context import Context
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.sampling import ALWAYS_OFF, Sampler, SamplingResult
+from opentelemetry.trace import Link, SpanKind
+from opentelemetry.trace.span import TraceState
+from opentelemetry.util.types import Attributes
+
+_logger = getLogger(__name__)
+
+DEFAULT_RULES_POLLING_INTERVAL = 300
+DEFAULT_TARGET_POLLING_INTERVAL = 10
+DEFAULT_SAMPLING_PROXY_ENDPOINT = "http://127.0.0.1:2000"
+
+class AwsXRayRemoteSampler(Sampler):
+    """
+    Remote Sampler for OpenTelemetry that gets sampling configurations from AWS X-Ray
+
+    Args:
+        resource: OpenTelemetry Resource (Optional)
+        endpoint: proxy endpoint for AWS X-Ray Sampling (Optional)
+        polling_interval: Polling interval for getSamplingRules call (Optional)
+        log_level: custom log level configuration for remote sampler (Optional)
+    """
+
+    __resource : Resource
+    __polling_interval : int
+    __xray_client : AwsXRaySamplingClient
+
+    def __init__(self, resource=None, endpoint=DEFAULT_SAMPLING_PROXY_ENDPOINT, polling_interval=DEFAULT_RULES_POLLING_INTERVAL, log_level = None):
+        # Override default log level
+        if log_level is not None:
+            _logger.setLevel(log_level)
+
+        self.__xray_client = AwsXRaySamplingClient(endpoint, log_level=log_level)
+        self.__polling_interval = polling_interval
+        self.__resource = resource
+
+        self.__start_sampling_rule_poller()
+
+    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":
+        # TODO: add sampling functionality
+        return ALWAYS_OFF
+
+    def get_description(self) -> str:
+        description = "AwsXRayRemoteSampler{remote sampling with AWS X-Ray}"
+        return description
+
+    def __get_and_update_sampling_rules(self):
+        sampling_rules = self.__xray_client.get_sampling_rules()
+
+        # TODO: Update sampling rules cache
+        _logger.info(f"Got Sampling Rules: {json.dumps([ob.__dict__ for ob in sampling_rules])}")
+
+    def __start_sampling_rule_poller(self):
+        self.__get_and_update_sampling_rules()
+        # Schedule the next sampling rule poll 
+        self._timer = Timer(self.__polling_interval, self.__start_sampling_rule_poller)
+        self._timer.daemon = True
+        self._timer.start()
+

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

@@ -0,0 +1,45 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+import json
+from logging import getLogger
+
+import requests
+
+from amazon.opentelemetry.distro.sampler.sampling_rule import SamplingRule
+
+_logger = getLogger(__name__)
+
+class AwsXRaySamplingClient:
+    def __init__(self, endpoint=None, log_level=None):
+        # Override default log level
+        if log_level is not None:
+            _logger.setLevel(log_level)
+
+        if endpoint is None:
+            _logger.error("endpoint must be specified")
+        self.__getSamplingRulesEndpoint = endpoint + "/GetSamplingRules"
+
+    def get_sampling_rules(self):
+        sampling_rules = []
+        headers = {'content-type': 'application/json'}
+
+        try:
+            r = requests.post(url=self.__getSamplingRulesEndpoint, headers=headers)
+            if r is None:
+                raise Exception("GetSamplingRules response is None")
+            sampling_rules_response = r.json()
+            if "SamplingRuleRecords" not in sampling_rules_response:
+                raise Exception(f"SamplingRuleRecords is missing in getSamplingRules response:{sampling_rules_response}")
+
+            sampling_rules_records = sampling_rules_response["SamplingRuleRecords"]
+            for record in sampling_rules_records:
+                sampling_rules.append(SamplingRule(**record["SamplingRule"]))
+
+        except requests.exceptions.RequestException as req_err:
+            _logger.exception(f"Request error occurred: {req_err}")
+        except json.JSONDecodeError as json_err:
+            _logger.exception(f"Error in decoding JSON response: {json_err}")
+        except Exception as ex:
+            _logger.exception(f"Exception occurred: {ex}")
+
+        return sampling_rules

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

@@ -0,0 +1,32 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: Apache-2.0
+class SamplingRule:
+    def __init__(
+        self, 
+        Attributes={},
+        FixedRate=None,
+        HTTPMethod=None,
+        Host=None,
+        Priority=None,
+        ReservoirSize=None,
+        ResourceARN=None,
+        RuleARN=None,
+        RuleName=None,
+        ServiceName=None,
+        ServiceType=None,
+        URLPath=None,
+        Version=None
+    ):
+        self.Attributes=Attributes
+        self.FixedRate=FixedRate
+        self.HTTPMethod=HTTPMethod
+        self.Host=Host
+        self.Priority=Priority
+        self.ReservoirSize=ReservoirSize
+        self.ResourceARN=ResourceARN
+        self.RuleARN=RuleARN
+        self.RuleName=RuleName
+        self.ServiceName=ServiceName
+        self.ServiceType=ServiceType
+        self.URLPath=URLPath
+        self.Version=Version