Repository: cloud-custodian/cloud-custodian

Issue:
add aws-xray-sdk so can include in packaged libraries
Require this one line change in order to have policies which include the following line
        packages: [botocore, boto3, urllib3, aws_xray_sdk]
deployed from a docker container running cloud custodian.  Other option is to add to requirements.txt but got feeling it was not there on purpose.


Discussion:
rather than this,  i'm wondering if we just want to install from requirements-dev.txt which will pull in all the optionals deps (psutil, xray sdk, etc).
also with the poetry packaging change #5320, xray sdk will be in requirements.txt
when I switched out to poetry on the main repo, I last minute replaced the docker file with one that uses poetry instead of setuptools, so that still needs adjustment to remove the --no-dev flag to get the xray and psutil tooling.

Resolving patch:
diff --git a/c7n/cli.py b/c7n/cli.py
--- a/c7n/cli.py
+++ b/c7n/cli.py
@@ -251,7 +251,7 @@ def setup_parser():
     run.add_argument(
         "--trace",
         dest="tracer",
-        help=argparse.SUPPRESS,
+        help="Tracing integration",
         default=None, nargs="?", const="default")
 
     schema_desc = ("Browse the available vocabularies (resources, filters, modes, and "
diff --git a/c7n/resources/aws.py b/c7n/resources/aws.py
--- a/c7n/resources/aws.py
+++ b/c7n/resources/aws.py
@@ -26,9 +26,11 @@
 import sys
 import tempfile
 import time
+import threading
 import traceback
 
 import boto3
+import botocore.client
 
 from botocore.validate import ParamValidator
 
@@ -286,6 +288,7 @@ def __repr__(self):
 
 
 class XrayEmitter:
+    # implement https://github.com/aws/aws-xray-sdk-python/issues/51
 
     def __init__(self):
         self.buf = []
@@ -301,18 +304,33 @@ def flush(self):
         self.buf = []
         for segment_set in utils.chunks(buf, 50):
             self.client.put_trace_segments(
-                TraceSegmentDocuments=[
-                    s.serialize() for s in segment_set])
+                TraceSegmentDocuments=[s.serialize() for s in segment_set])
 
 
 class XrayContext(Context):
+    """Specialized XRay Context for Custodian.
+
+    A context is used as a segment storage stack for currently in
+    progress segments.
+
+    We use a customized context for custodian as policy execution
+    commonly uses a concurrent.futures threadpool pattern during
+    execution for api concurrency. Default xray semantics would use
+    thread local storage and treat each of those as separate trace
+    executions. We want to aggregate/associate all thread pool api
+    executions to the custoidan policy execution. XRay sdk supports
+    this via manual code for every thread pool usage, but we don't
+    want to explicitly couple xray integration everywhere across the
+    codebase. Instead we use a context that is aware of custodian
+    usage of threads and associates subsegments therein to the policy
+    execution active subsegment.
+    """
 
     def __init__(self, *args, **kw):
         super(XrayContext, self).__init__(*args, **kw)
-        # We want process global semantics as policy execution
-        # can span threads.
         self._local = Bag()
         self._current_subsegment = None
+        self._main_tid = threading.get_ident()
 
     def handle_context_missing(self):
         """Custodian has a few api calls out of band of policy execution.
@@ -324,6 +342,91 @@ def handle_context_missing(self):
         so default to disabling context missing output.
         """
 
+    # Annotate any segments/subsegments with their thread ids.
+    def put_segment(self, segment):
+        if getattr(segment, 'thread_id', None) is None:
+            segment.thread_id = threading.get_ident()
+        super().put_segment(segment)
+
+    def put_subsegment(self, subsegment):
+        if getattr(subsegment, 'thread_id', None) is None:
+            subsegment.thread_id = threading.get_ident()
+        super().put_subsegment(subsegment)
+
+    # Override since we're not just popping the end of the stack, we're removing
+    # the thread subsegment from the array by identity.
+    def end_subsegment(self, end_time):
+        subsegment = self.get_trace_entity()
+        if self._is_subsegment(subsegment):
+            subsegment.close(end_time)
+            self._local.entities.remove(subsegment)
+            return True
+        else:
+            log.warning("No subsegment to end.")
+            return False
+
+    # Override get trace identity, any worker thread will find its own subsegment
+    # on the stack, else will use the main thread's sub/segment
+    def get_trace_entity(self):
+        tid = threading.get_ident()
+        entities = self._local.get('entities', ())
+        for s in reversed(entities):
+            if s.thread_id == tid:
+                return s
+            # custodian main thread won't advance (create new segment)
+            # with worker threads still doing pool work.
+            elif s.thread_id == self._main_tid:
+                return s
+        return self.handle_context_missing()
+
+
+if HAVE_XRAY:  # pragma: no cover
+    # sigh, work around aws xray sdk issues, via monkey patch
+    # https://github.com/aws/aws-xray-sdk-python/issues/207
+
+    # we can't do a minimal monkey, due to namespace occlusion they
+    # do with function imports masking modules :-(
+
+    # all of this to add a PutTraceSegments to the exempted methods list.
+    import wrapt
+
+    from aws_xray_sdk.ext.boto_utils import inject_header, aws_meta_processor
+    from aws_xray_sdk.ext import botocore as xray_boto3
+
+    def _xray_traced_botocore(wrapped, instance, args, kwargs):
+        service = instance._service_model.metadata["endpointPrefix"]
+        if service == 'xray':
+            # skip tracing for SDK built-in sampling pollers
+            if ('GetSamplingRules' in args or
+                'GetSamplingTargets' in args or
+                    'PutTraceSegments' in args):
+                return wrapped(*args, **kwargs)
+        return xray_recorder.record_subsegment(
+            wrapped, instance, args, kwargs,
+            name=service,
+            namespace='aws',
+            meta_processor=aws_meta_processor,
+        )
+
+    def _xray_patch_botocore():
+        if hasattr(botocore.client, '_xray_enabled'):
+            return
+        setattr(botocore.client, '_xray_enabled', True)
+
+        wrapt.wrap_function_wrapper(
+            'botocore.client',
+            'BaseClient._make_api_call',
+            _xray_traced_botocore,
+        )
+
+        wrapt.wrap_function_wrapper(
+            'botocore.endpoint',
+            'Endpoint.prepare_request',
+            inject_header,
+        )
+
+    xray_boto3.patch = _xray_patch_botocore
+
 
 @tracer_outputs.register('xray', condition=HAVE_XRAY)
 class XrayTracer:
@@ -335,12 +438,13 @@ class XrayTracer:
     service_name = 'custodian'
 
     @classmethod
-    def initialize(cls):
+    def initialize(cls, config):
         context = XrayContext()
+        sampling = config.get('sample', 'true') == 'true' and True or False
         xray_recorder.configure(
             emitter=cls.use_daemon is False and cls.emitter or None,
             context=context,
-            sampling=True,
+            sampling=sampling,
             context_missing='LOG_ERROR')
         patch(['boto3', 'requests'])
         logging.getLogger('aws_xray_sdk.core').setLevel(logging.ERROR)
@@ -391,6 +495,9 @@ def __exit__(self, exc_type=None, exc_value=None, exc_traceback=None):
         xray_recorder.end_segment()
         if not self.use_daemon:
             self.emitter.flush()
+            log.info(
+                ('View XRay Trace https://console.aws.amazon.com/xray/home?region=%s#/'
+                 'traces/%s' % (self.ctx.options.region, self.segment.trace_id)))
         self.metadata.clear()
 
 
@@ -517,7 +624,7 @@ def initialize(self, options):
         _default_region(options)
         _default_account_id(options)
         if options.tracer and options.tracer.startswith('xray') and HAVE_XRAY:
-            XrayTracer.initialize()
+            XrayTracer.initialize(utils.parse_url_config(options.tracer))
 
         return options
 


Test patch:
diff --git a/tests/test_aws.py b/tests/test_aws.py
--- a/tests/test_aws.py
+++ b/tests/test_aws.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import json
+import threading
 
 from mock import Mock
 
@@ -24,6 +25,10 @@
 from .common import BaseTest
 
 
+from aws_xray_sdk.core.models.segment import Segment
+from aws_xray_sdk.core.models.subsegment import Subsegment
+
+
 class TraceDoc(Bag):
 
     def serialize(self):
@@ -139,13 +144,47 @@ def test_validate(self):
 
 class TracerTest(BaseTest):
 
+    def test_context(self):
+        store = aws.XrayContext()
+        self.assertEqual(store.handle_context_missing(), None)
+        x = Segment('foo')
+        y = Segment('foo')
+        a = Subsegment('bar', 'boo', x)
+        b = Subsegment('bar', 'boo', x)
+        b.thread_id = '123'
+        store.put_segment(x)
+        store.put_subsegment(a)
+        store.put_subsegment(b)
+
+        self.assertEqual(store._local.entities, [x, a, b])
+        self.assertEqual(store.get_trace_entity(), a)
+        store.end_subsegment(a)
+        self.assertEqual(store.get_trace_entity(), x)
+        store.put_segment(y)
+        self.assertEqual(store._local.entities, [y])
+        self.assertEqual(store.get_trace_entity(), y)
+        self.assertFalse(store.end_subsegment(42))
+
+    def test_context_worker_thread_main_acquire(self):
+        store = aws.XrayContext()
+        x = Segment('foo')
+        a = Subsegment('bar', 'boo', x)
+        store.put_segment(x)
+        store.put_subsegment(a)
+
+        def get_ident():
+            return 42
+
+        self.patch(threading, 'get_ident', get_ident)
+        self.assertEqual(store.get_trace_entity(), a)
+
     def test_tracer(self):
         session_factory = self.replay_flight_data('output-xray-trace')
         policy = Bag(name='test', resource_type='ec2')
         ctx = Bag(
             policy=policy,
             session_factory=session_factory,
-            options=Bag(account_id='644160558196'))
+            options=Bag(account_id='644160558196', region='us-east-1',))
         ctx.get_metadata = lambda *args: {}
         config = Bag()
         tracer = aws.XrayTracer(ctx, config)