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