Initial change on centralized sampling support

add poller logic

cs support beta

cs minor bug fixes

minor fixes

minor fixes

minor fixes

update change log

release commit

Author: haotianw465

CHANGELOG.rst

@@ -2,6 +2,19 @@
 CHANGELOG
 =========
 
+2.0.0
+=====
+* **Breaking**: The default sampler now launches background tasks to poll sampling rules from X-Ray backend. See the new default sampling strategy in more details here: 
+* **Breaking**: The `should_trace` function in the sampler now takes a dictionary for sampling rule matching.
+* **Breaking**: The original sampling modules for local defined rules are moved from `models.sampling` to `models.sampling.local`.
+* **Breaking**: The default behavior of `patch_all` changed to selectively patches libraries to avoid double patching. You can use `patch_all(double_patch=True)` to force it to patch ALL supported libraries. See more details on `ISSUE63 <https://github.com/aws/aws-xray-sdk-python/issues/63>`_
+* **Breaking** The latest `botocore` that has new X-Ray service API `GetSamplingRules` and `GetSamplingTargets` are required.
+* feature: Environment variable `AWS_XRAY_DAEMON_ADDRESS` now takes an additional notation in `tcp:127.0.0.1:2000 udp:127.0.0.2:2001` to set TCP and UDP destination separately. By default it assumes a X-Ray daemon listening to both UDP and TCP traffic on `127.0.0.1:2000`.
+* feature: Added MongoDB python client support. `PR65 <https://github.com/aws/aws-xray-sdk-python/pull/65>`_.
+* bugfix: Support binding connection in sqlalchemy as well as engine. `PR78 <https://github.com/aws/aws-xray-sdk-python/pull/78>`_. 
+* bugfix: Flask middleware safe request teardown. `ISSUE75 <https://github.com/aws/aws-xray-sdk-python/issues/75>`_. 
+
+
 1.1.2
 =====
 * bugfix: Fixed an issue on PynamoDB patcher where the capture didn't handle client timeout.

MANIFEST.in

@@ -1,4 +1,4 @@
 include aws_xray_sdk/ext/resources/*.json
-include aws_xray_sdk/core/sampling/*.json
+include aws_xray_sdk/core/sampling/local/*.json
 include README.md
 include LICENSE

aws_xray_sdk/core/daemon_config.py

@@ -0,0 +1,76 @@
+import os
+
+from .exceptions.exceptions import InvalidDaemonAddressException
+
+DAEMON_ADDRESS_KEY = "AWS_XRAY_DAEMON_ADDRESS"
+DEFAULT_ADDRESS = '127.0.0.1:2000'
+
+
+class DaemonConfig(object):
+    """The class that stores X-Ray daemon configuration about
+    the ip address and port for UDP and TCP port. It gets the address
+    string from ``AWS_TRACING_DAEMON_ADDRESS`` and then from recorder's
+    configuration for ``daemon_address``.
+    A notation of '127.0.0.1:2000' or 'tcp:127.0.0.1:2000 udp:127.0.0.2:2001'
+    are both acceptable. The former one means UDP and TCP are running at
+    the same address.
+    By default it assumes a X-Ray daemon running at 127.0.0.1:2000
+    listening to both UDP and TCP traffic.
+    """
+    def __init__(self, daemon_address=DEFAULT_ADDRESS):
+        if daemon_address is None:
+            daemon_address = DEFAULT_ADDRESS
+
+        val = os.getenv(DAEMON_ADDRESS_KEY, daemon_address)
+        configs = val.split(' ')
+        if len(configs) == 1:
+            self._parse_single_form(configs[0])
+        elif len(configs) == 2:
+            self._parse_double_form(configs[0], configs[1], val)
+        else:
+            raise InvalidDaemonAddressException('Invalid daemon address %s specified.' % val)
+
+    def _parse_single_form(self, val):
+        try:
+            configs = val.split(':')
+            self._udp_ip = configs[0]
+            self._udp_port = int(configs[1])
+            self._tcp_ip = configs[0]
+            self._tcp_port = int(configs[1])
+        except Exception:
+            raise InvalidDaemonAddressException('Invalid daemon address %s specified.' % val)
+
+    def _parse_double_form(self, val1, val2, origin):
+        try:
+            configs1 = val1.split(':')
+            configs2 = val2.split(':')
+            mapping = {
+                configs1[0]: configs1,
+                configs2[0]: configs2,
+            }
+
+            tcp_info = mapping.get('tcp')
+            udp_info = mapping.get('udp')
+
+            self._tcp_ip = tcp_info[1]
+            self._tcp_port = int(tcp_info[2])
+            self._udp_ip = udp_info[1]
+            self._udp_port = int(udp_info[2])
+        except Exception:
+            raise InvalidDaemonAddressException('Invalid daemon address %s specified.' % origin)
+
+    @property
+    def udp_ip(self):
+        return self._udp_ip
+
+    @property
+    def udp_port(self):
+        return self._udp_port
+
+    @property
+    def tcp_ip(self):
+        return self._tcp_ip
+
+    @property
+    def tcp_port(self):
+        return self._tcp_port

aws_xray_sdk/core/emitters/udp_emitter.py

@@ -1,15 +1,15 @@
-import os
-import socket
 import logging
+import socket
 
+from aws_xray_sdk.core.daemon_config import DaemonConfig
 from ..exceptions.exceptions import InvalidDaemonAddressException
 
 log = logging.getLogger(__name__)
 
 
-PROTOCOL_HEADER = "{\"format\": \"json\", \"version\": 1}"
+PROTOCOL_HEADER = "{\"format\":\"json\",\"version\":1}"
 PROTOCOL_DELIMITER = '\n'
-DAEMON_ADDRESS_KEY = "AWS_XRAY_DAEMON_ADDRESS"
+DEFAULT_DAEMON_ADDRESS = '127.0.0.1:2000'
 
 
 class UDPEmitter(object):
@@ -19,12 +19,11 @@ class UDPEmitter(object):
     exception on the actual data transfer between the socket and the daemon,
     it logs the exception and continue.
     """
-    def __init__(self, daemon_address='127.0.0.1:2000'):
+    def __init__(self, daemon_address=DEFAULT_DAEMON_ADDRESS):
 
         self._socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
         self._socket.setblocking(0)
-        address = os.getenv(DAEMON_ADDRESS_KEY, daemon_address)
-        self._ip, self._port = self._parse_address(address)
+        self.set_daemon_address(daemon_address)
 
     def send_entity(self, entity):
         """
@@ -42,11 +41,20 @@ def send_entity(self, entity):
 
     def set_daemon_address(self, address):
         """
-        Takes a full address like 127.0.0.1:2000 and parses it into ip address
-        and port. Throws an exception if the address has invalid format.
+        Set up UDP ip and port from the raw daemon address
+        string using ``DaemonConfig`` class utlities.
         """
         if address:
-            self._ip, self._port = self._parse_address(address)
+            daemon_config = DaemonConfig(address)
+            self._ip, self._port = daemon_config.udp_ip, daemon_config.udp_port
+
+    @property
+    def ip(self):
+        return self._ip
+
+    @property
+    def port(self):
+        return self._port
 
     def _send_data(self, data):
 

aws_xray_sdk/core/models/segment.py

@@ -113,6 +113,16 @@ def set_service(self, service_info):
         """
         self.service = service_info
 
+    def set_rule_name(self, rule_name):
+        """
+        Add the matched centralized sampling rule name
+        if a segment is sampled because of that rule.
+        This method should be only used by the recorder.
+        """
+        if not self.aws.get('xray', None):
+            self.aws['xray'] = {}
+        self.aws['xray']['rule_name'] = rule_name
+
     def save_origin_trace_header(self, trace_header):
         """
         Temporarily store additional data fields in trace header

aws_xray_sdk/core/patcher.py

@@ -14,11 +14,24 @@
     'pymongo',
 )
 
+NO_DOUBLE_PATCH = (
+    'aiobotocore',
+    'botocore',
+    'pynamodb',
+    'requests',
+    'sqlite3',
+    'mysql',
+    'pymongo',
+)
+
 _PATCHED_MODULES = set()
 
 
-def patch_all():
-    patch(SUPPORTED_MODULES, raise_errors=False)
+def patch_all(double_patch=False):
+    if double_patch:
+        patch(SUPPORTED_MODULES, raise_errors=False)
+    else:
+        patch(NO_DOUBLE_PATCH, raise_errors=False)
 
 
 def patch(modules_to_patch, raise_errors=True):

aws_xray_sdk/core/recorder.py

@@ -14,9 +14,11 @@
 from .models.default_dynamic_naming import DefaultDynamicNaming
 from .models.dummy_entities import DummySegment, DummySubsegment
 from .emitters.udp_emitter import UDPEmitter
-from .sampling.default_sampler import DefaultSampler
+from .sampling.sampler import DefaultSampler
+from .sampling.local.sampler import LocalSampler
 from .streaming.default_streaming import DefaultStreaming
 from .context import Context
+from .daemon_config import DaemonConfig
 from .plugins.utils import get_plugin_modules
 from .lambda_launcher import check_in_lambda
 from .exceptions.exceptions import SegmentNameMissingException
@@ -59,11 +61,12 @@ def __init__(self):
             # Special handling when running on AWS Lambda.
             self._context = context
             self.streaming_threshold = 0
+            self._sampler = LocalSampler()
         else:
             self._context = Context()
+            self._sampler = DefaultSampler()
 
         self._emitter = UDPEmitter()
-        self._sampler = DefaultSampler()
         self._sampling = True
         self._max_trace_back = 10
         self._plugins = None
@@ -72,12 +75,15 @@ def __init__(self):
         self._aws_metadata = copy.deepcopy(XRAY_META)
         self._origin = None
 
+        if type(self.sampler).__name__ == 'DefaultSampler':
+            self.sampler.load_settings(DaemonConfig(), self.context)
+
     def configure(self, sampling=None, plugins=None,
                   context_missing=None, sampling_rules=None,
                   daemon_address=None, service=None,
                   context=None, emitter=None, streaming=None,
                   dynamic_naming=None, streaming_threshold=None,
-                  max_trace_back=None):
+                  max_trace_back=None, sampler=None):
         """Configure global X-Ray recorder.
 
         Configure needs to run before patching thrid party libraries
@@ -86,10 +92,16 @@ def configure(self, sampling=None, plugins=None,
         :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
-            is running in AWS Lambda.
-        :param sampling_rules: Pass a set of custom sampling rules.
+            is running in AWS Lambda. The recorder always respect the incoming
+            sampling decisions regardless of this setting.
+        :param sampling_rules: Pass a set of local custom sampling rules.
             Can be an absolute path of the sampling rule config json file
-            or a dictionary that defines those rules.
+            or a dictionary that defines those rules. This will also be the
+            fallback rules in case of centralized sampling opted-in while
+            the cetralized sampling rules are not available.
+        :param sampler: The sampler used to make sampling decisions. The SDK
+            provides two built-in samplers. One is centralized rules based and
+            the other is local rules based. The former is the default.
         :param tuple plugins: plugins that add extra metadata to each segment.
             Currently available plugins are EC2Plugin, ECS plugin and
             ElasticBeanstalkPlugin.
@@ -127,6 +139,8 @@ class to have your own implementation of the streaming process.
         """
         if sampling is not None:
             self.sampling = sampling
+        if sampler:
+            self.sampler = sampler
         if service:
             self.service = os.getenv(TRACING_NAME_KEY, service)
         if sampling_rules:
@@ -160,6 +174,10 @@ class to have your own implementation of the streaming process.
             self._aws_metadata = copy.deepcopy(XRAY_META)
             self._origin = None
 
+        if type(self.sampler).__name__ == 'DefaultSampler':
+            self.sampler.load_settings(DaemonConfig(daemon_address),
+                                       self.context, self._origin)
+
     def begin_segment(self, name=None, traceid=None,
                       parent_id=None, sampling=None):
         """
@@ -175,21 +193,26 @@ def begin_segment(self, name=None, traceid=None,
         if not seg_name:
             raise SegmentNameMissingException("Segment name is required.")
 
-        # we respect sampling decision regardless of recorder configuration.
-        dummy = False
+        # Sampling decision is None if not sampled.
+        # In a sampled case it could be either a string or 1
+        # depending on if centralized or local sampling rule takes effect.
+        decision = True
+
+        # we respect the input sampling decision
+        # regardless of recorder configuration.
         if sampling == 0:
-            dummy = True
-        elif sampling == 1:
-            dummy = False
-        elif self.sampling and not self._sampler.should_trace():
-            dummy = True
+            decision = False
+        elif sampling:
+            decision = sampling
+        elif self.sampling:
+            decision = self._sampler.should_trace()
 
-        if dummy:
+        if not decision:
             segment = DummySegment(seg_name)
         else:
             segment = Segment(name=seg_name, traceid=traceid,
                               parent_id=parent_id)
-            self._populate_runtime_context(segment)
+            self._populate_runtime_context(segment, decision)
 
         self.context.put_segment(segment)
         return segment
@@ -397,13 +420,16 @@ def record_subsegment(self, wrapped, instance, args, kwargs, name,
 
                 self.end_subsegment(end_time)
 
-    def _populate_runtime_context(self, segment):
+    def _populate_runtime_context(self, segment, sampling_decision):
         if self._origin:
             setattr(segment, 'origin', self._origin)
 
-        segment.set_aws(self._aws_metadata)
+        segment.set_aws(copy.deepcopy(self._aws_metadata))
         segment.set_service(SERVICE_INFO)
 
+        if isinstance(sampling_decision, string_types):
+            segment.set_rule_name(sampling_decision)
+
     def _send_segment(self):
         """
         Send the current segment to X-Ray daemon if it is present and
@@ -429,10 +455,10 @@ def _load_sampling_rules(self, sampling_rules):
             return
 
         if isinstance(sampling_rules, dict):
-            self.sampler = DefaultSampler(sampling_rules)
+            self.sampler.load_local_rules(sampling_rules)
         else:
             with open(sampling_rules) as f:
-                self.sampler = DefaultSampler(json.load(f))
+                self.sampler.load_local_rules(json.load(f))
 
     def _is_subsegment(self, entity):