def test_invalid_env_string():
os.environ[XRAY_ENABLED_KEY] = "INVALID"
# Env Variable takes precedence. This is called to activate the internal check
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "1.0"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "1-.0"
global_sdk_config.set_sdk_enabled(False)
assert global_sdk_config.sdk_enabled() is True
def test_env_enable_case():
os.environ[XRAY_ENABLED_KEY] = "TrUE"
# Env Variable takes precedence. This is called to activate the internal check
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "true"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "1"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "y"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "t"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "False"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is False
os.environ[XRAY_ENABLED_KEY] = "falSE"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is False
os.environ[XRAY_ENABLED_KEY] = "0"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is False
def should_trace(self, sampling_req=None):
"""
Return the matched sampling rule name if the sampler finds one
and decide to sample. If no sampling rule matched, it falls back
to the local sampler's ``should_trace`` implementation.
All optional arguments are extracted from incoming requests by
X-Ray middleware to perform path based sampling.
"""
if not global_sdk_config.sdk_enabled():
return False
if not self._started:
self.start(
) # only front-end that actually uses the sampler spawns poller threads
now = int(time.time())
if sampling_req and not sampling_req.get('service_type', None):
sampling_req['service_type'] = self._origin
elif sampling_req is None:
sampling_req = {'service_type': self._origin}
matched_rule = self._cache.get_matched_rule(sampling_req, now)
if matched_rule:
log.debug('Rule %s is selected to make a sampling decision.',
matched_rule.name)
return self._process_matched_rule(matched_rule, now)
else:
log.info(
'No effective centralized sampling rule match. Fallback to local rules.'
)
return self._local_sampler.should_trace(sampling_req)
def create_dummy_segment(self, name):
if not global_sdk_config.sdk_enabled():
try:
self.current_segment()
except SegmentNotFoundException:
segment = DummySegment(name)
self.context.put_segment(segment)
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'.
"""
# 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.
if not global_sdk_config.sdk_enabled():
return DummySubsegment(
DummySegment(global_sdk_config.DISABLED_ENTITY_NAME))
segment = self.current_segment()
if not segment:
log.warning("No segment found, cannot begin subsegment %s." % name)
return None
if not segment.sampled:
subsegment = DummySubsegment(segment, name)
else:
subsegment = Subsegment(name, namespace, segment)
self.context.put_subsegment(subsegment)
return subsegment
def test_env_var_precedence():
os.environ[XRAY_ENABLED_KEY] = "true"
# Env Variable takes precedence. This is called to activate the internal check
global_sdk_config.set_sdk_enabled(False)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "false"
global_sdk_config.set_sdk_enabled(False)
assert global_sdk_config.sdk_enabled() is False
os.environ[XRAY_ENABLED_KEY] = "false"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is False
os.environ[XRAY_ENABLED_KEY] = "true"
global_sdk_config.set_sdk_enabled(True)
assert global_sdk_config.sdk_enabled() is True
os.environ[XRAY_ENABLED_KEY] = "true"
global_sdk_config.set_sdk_enabled(None)
assert global_sdk_config.sdk_enabled() is True
def is_sampled(self):
"""
Check if the current trace entity is sampled or not.
Return `False` if no active entity found.
"""
if not global_sdk_config.sdk_enabled():
# Disabled SDK is never sampled
return False
entity = self.get_trace_entity()
if entity:
return entity.sampled
return False
def get_trace_entity(self):
"""
Return the current trace entity(segment/subsegment). If there is none,
it behaves based on pre-defined ``context_missing`` strategy.
If the SDK is disabled, returns a DummySegment
"""
if not getattr(self._local, 'entities', None):
if not global_sdk_config.sdk_enabled():
return DummySegment()
return self.handle_context_missing()
return self._local.entities[-1]
def start(self):
"""
Start rule poller and target poller once X-Ray daemon address
and context manager is in place.
"""
if not global_sdk_config.sdk_enabled():
return
with self._lock:
if not self._started:
self._rule_poller.start()
self._target_poller.start()
self._started = True
def put_annotation(self, key, value):
"""
Annotate current active trace entity with a key-value pair.
Annotations will be indexed for later search query.
:param str key: annotation key
:param object value: annotation value. Any type other than
string/number/bool will be dropped
"""
if not global_sdk_config.sdk_enabled():
return
entity = self.get_trace_entity()
if entity and entity.sampled:
entity.put_annotation(key, value)
def current_subsegment(self):
"""
Return the latest opened subsegment. In a multithreading environment,
this will make sure the subsegment returned is one created
by the same thread.
"""
if not global_sdk_config.sdk_enabled():
return DummySubsegment(
DummySegment(global_sdk_config.DISABLED_ENTITY_NAME))
entity = self.get_trace_entity()
if self._is_subsegment(entity):
return entity
else:
return None
def begin_segment(self,
name=None,
traceid=None,
parent_id=None,
sampling=None):
"""
Begin a segment on the current thread and return it. The recorder
only keeps one segment at a time. Create the second one without
closing existing one will overwrite it.
:param str name: the name of the segment
:param str traceid: trace id of the segment
:param int sampling: 0 means not sampled, 1 means sampled
"""
seg_name = name or self.service
if not seg_name:
raise SegmentNameMissingException("Segment name is required.")
# 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
# To disable the recorder, we set the sampling decision to always be false.
# This way, when segments are generated, they become dummy segments and are ultimately never sent.
# The call to self._sampler.should_trace() is never called either so the poller threads are never started.
if not global_sdk_config.sdk_enabled():
sampling = 0
# we respect the input sampling decision
# regardless of recorder configuration.
if sampling == 0:
decision = False
elif sampling:
decision = sampling
elif self.sampling:
decision = self._sampler.should_trace()
if not decision:
segment = DummySegment(seg_name)
else:
segment = Segment(name=seg_name,
traceid=traceid,
parent_id=parent_id)
self._populate_runtime_context(segment, decision)
self.context.put_segment(segment)
return segment
def put_metadata(self, key, value, namespace='default'):
"""
Add metadata to the current active trace entity.
Metadata is not indexed but can be later retrieved
by BatchGetTraces API.
:param str namespace: optional. Default namespace is `default`.
It must be a string and prefix `AWS.` is reserved.
:param str key: metadata key under specified namespace
:param object value: any object that can be serialized into JSON string
"""
if not global_sdk_config.sdk_enabled():
return
entity = self.get_trace_entity()
if entity and entity.sampled:
entity.put_metadata(key, value, namespace)
def end_segment(self, end_time=None):
"""
End the current segment and send it to X-Ray daemon
if it is ready to send. Ready means segment and
all its subsegments are closed.
:param float end_time: segment completion in unix epoch in seconds.
"""
# When the SDK is disabled we return
if not global_sdk_config.sdk_enabled():
return
self.context.end_segment(end_time)
segment = self.current_segment()
if segment and segment.ready_to_send():
self._send_segment()
def put_subsegment(self, subsegment):
"""
Refresh the facade segment every time this function is invoked to prevent
a new subsegment from being attached to a leaked segment/subsegment.
"""
current_entity = self.get_trace_entity()
if not self._is_subsegment(
current_entity) and current_entity.initializing:
if global_sdk_config.sdk_enabled():
log.warning(
"Subsegment %s discarded due to Lambda worker still initializing"
% subsegment.name)
return
current_entity.add_subsegment(subsegment)
self._local.entities.append(subsegment)
def begin_segment(self,
name=None,
traceid=None,
parent_id=None,
sampling=None):
"""
Begin a segment on the current thread and return it. The recorder
only keeps one segment at a time. Create the second one without
closing existing one will overwrite it.
:param str name: the name of the segment
:param str traceid: trace id of the segment
:param int sampling: 0 means not sampled, 1 means sampled
"""
# Disable the recorder; return a generated dummy segment.
if not global_sdk_config.sdk_enabled():
return DummySegment(global_sdk_config.DISABLED_ENTITY_NAME)
seg_name = name or self.service
if not seg_name:
raise SegmentNameMissingException("Segment name is required.")
# 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:
decision = False
elif sampling:
decision = sampling
elif self.sampling:
decision = self._sampler.should_trace()
if not decision:
segment = DummySegment(seg_name)
else:
segment = Segment(name=seg_name,
traceid=traceid,
parent_id=parent_id)
self._populate_runtime_context(segment, decision)
self.context.put_segment(segment)
return segment
def record_subsegment(self, wrapped, instance, args, kwargs, name,
namespace, meta_processor):
# In the case when the SDK is disabled, we ensure that a parent segment exists, because this is usually
# handled by the middleware. We generate a dummy segment as the parent segment if one doesn't exist.
# This is to allow potential segment method calls to not throw exceptions in the captured method.
if not global_sdk_config.sdk_enabled():
try:
self.current_segment()
except SegmentNotFoundException:
segment = DummySegment(name)
self.context.put_segment(segment)
subsegment = self.begin_subsegment(name, namespace)
exception = None
stack = None
return_value = None
try:
return_value = wrapped(*args, **kwargs)
return return_value
except Exception as e:
exception = e
stack = stacktrace.get_stacktrace(limit=self.max_trace_back)
raise
finally:
# No-op if subsegment is `None` due to `LOG_ERROR`.
if subsegment is not None:
end_time = time.time()
if callable(meta_processor):
meta_processor(
wrapped=wrapped,
instance=instance,
args=args,
kwargs=kwargs,
return_value=return_value,
exception=exception,
subsegment=subsegment,
stack=stack,
)
elif exception:
subsegment.add_exception(exception, stack)
self.end_subsegment(end_time)
def end_subsegment(self, end_time=None):
"""
End the current active subsegment. If this is the last one open
under its parent segment, the entire segment will be sent.
:param float end_time: subsegment compeletion in unix epoch in seconds.
"""
if not global_sdk_config.sdk_enabled():
return
if not self.context.end_subsegment(end_time):
return
# if segment is already close, we check if we can send entire segment
# otherwise we check if we need to stream some subsegments
if self.current_segment().ready_to_send():
self._send_segment()
else:
self.stream_subsegments()
def patch(modules_to_patch, raise_errors=True, ignore_module_patterns=None):
enabled = global_sdk_config.sdk_enabled()
if not enabled:
log.debug(
"Skipped patching modules %s because the SDK is currently disabled."
% ', '.join(modules_to_patch))
return # Disable module patching if the SDK is disabled.
modules = set()
for module_to_patch in modules_to_patch:
# boto3 depends on botocore and patching botocore is sufficient
if module_to_patch == 'boto3':
modules.add('botocore')
# aioboto3 depends on aiobotocore and patching aiobotocore is sufficient
elif module_to_patch == 'aioboto3':
modules.add('aiobotocore')
# pynamodb requires botocore to be patched as well
elif module_to_patch == 'pynamodb':
modules.add('botocore')
modules.add(module_to_patch)
else:
modules.add(module_to_patch)
unsupported_modules = set(module for module in modules
if module not in SUPPORTED_MODULES)
native_modules = modules - unsupported_modules
external_modules = set(module for module in unsupported_modules
if _is_valid_import(module))
unsupported_modules = unsupported_modules - external_modules
if unsupported_modules:
raise Exception('modules %s are currently not supported for patching' %
', '.join(unsupported_modules))
for m in native_modules:
_patch_module(m, raise_errors)
ignore_module_patterns = [
re.compile(pattern) for pattern in ignore_module_patterns or []
]
for m in external_modules:
_external_module_patch(m, ignore_module_patterns)
def _refresh_context(self):
"""
Get current facade segment. To prevent resource leaking in Lambda worker,
every time there is segment present, we compare its trace id to current
environment variables. If it is different we create a new facade segment
and clean up subsegments stored.
"""
header_str = os.getenv(LAMBDA_TRACE_HEADER_KEY)
trace_header = TraceHeader.from_header_str(header_str)
if not global_sdk_config.sdk_enabled():
trace_header._sampled = False
segment = getattr(self._local, 'segment', None)
if segment:
# Ensure customers don't have leaked subsegments across invocations
if not trace_header.root or trace_header.root == segment.trace_id:
return
else:
self._initialize_context(trace_header)
else:
self._initialize_context(trace_header)
async def __call__(self, wrapped, instance, args, kwargs):
if is_already_recording(wrapped):
# The wrapped function is already decorated, the subsegment will be created later,
# just return the result
return await wrapped(*args, **kwargs)
func_name = self.name
if not func_name:
func_name = wrapped.__name__
if not global_sdk_config.sdk_enabled() or self.instance.app is None:
try:
segment = self.recorder.current_segment()
except SegmentNotFoundException:
segment = DummySegment(func_name)
self.recorder.context.put_segment(segment)
finally:
if segment is None:
error_logger.warning(
CAPTURE_WARNING.format(name=func_name))
elif (hasattr(self.instance.app, "initialized_plugins")
and "incendiary"
not in self.instance.app.initialized_plugins):
error_logger.warning(
CAPTURE_WARNING.format(name=func_name))
try:
return await self.recorder.record_subsegment_async(
wrapped,
instance,
args,
kwargs,
name=func_name,
namespace="local",
meta_processor=None,
)
except exceptions.AlreadyEndedException:
return await wrapped(*args, **kwargs)
def _initialize_context(self, trace_header):
"""
Create a facade segment based on environment variables
set by AWS Lambda and initialize storage for subsegments.
"""
sampled = None
if not global_sdk_config.sdk_enabled():
# Force subsequent subsegments to be disabled and turned into DummySegments.
sampled = False
elif trace_header.sampled == 0:
sampled = False
elif trace_header.sampled == 1:
sampled = True
segment = FacadeSegment(
name='facade',
traceid=trace_header.root,
entityid=trace_header.parent,
sampled=sampled,
)
setattr(self._local, 'segment', segment)
setattr(self._local, 'entities', [])
def test_default_enabled():
assert global_sdk_config.sdk_enabled()
segment = xray_recorder.begin_segment('name')
subsegment = xray_recorder.begin_subsegment('name')
assert type(xray_recorder.current_segment()) is Segment
assert type(xray_recorder.current_subsegment()) is Subsegment
def test_default_enabled():
assert global_sdk_config.sdk_enabled() is True
