def inject_trace_header(headers, entity):
"""
Extract trace id, entity id and sampling decision
from the input entity and inject these information
to headers.
:param dict headers: http headers to inject
:param Entity entity: trace entity that the trace header
value generated from.
"""
if not entity:
return
if hasattr(entity, 'type') and entity.type == 'subsegment':
header = entity.parent_segment.get_origin_trace_header()
else:
header = entity.get_origin_trace_header()
data = header.data if header else None
to_insert = TraceHeader(
root=entity.trace_id,
parent=entity.id,
sampled=entity.sampled,
data=data,
)
value = to_insert.to_header_str()
headers[http.XRAY_HEADER] = value
def test_no_sample():
header = TraceHeader(root=TRACE_ID, parent=PARENT_ID)
assert header.sampled is None
assert header.root == TRACE_ID
assert header.parent == PARENT_ID
assert header.to_header_str() == 'Root=%s;Parent=%s' % (TRACE_ID,
PARENT_ID)
def construct_xray_header(headers):
"""
Construct a ``TraceHeader`` object from dictionary headers
of the incoming request. This method should always return
a ``TraceHeader`` object regardless of tracing header's presence
in the incoming request.
"""
header_str = headers.get(http.XRAY_HEADER) or headers.get(http.ALT_XRAY_HEADER)
if header_str:
return TraceHeader.from_header_str(header_str)
else:
return TraceHeader()
def prepare_response_header(origin_header, segment):
"""
Prepare a trace header to be inserted into response
based on original header and the request segment.
"""
if origin_header and origin_header.sampled == '?':
new_header = TraceHeader(root=segment.trace_id,
sampled=segment.sampled)
else:
new_header = TraceHeader(root=segment.trace_id)
return new_header.to_header_str()
def test_arbitrary_fields():
origin_header_str = 'Root=%s;k1=v1;k2=v2' % TRACE_ID
header = TraceHeader.from_header_str(origin_header_str)
header_str = header.to_header_str()
assert 'k1=v1' in header_str
assert 'k2=v2' in header_str
def test_from_str():
# a full header string that has all fields present
header_str1 = 'Root=%s;Parent=%s;Sampled=1' % (TRACE_ID, PARENT_ID)
header1 = TraceHeader.from_header_str(header_str1)
assert header1.root == TRACE_ID
assert header1.parent == PARENT_ID
assert header1.sampled == 1
# missing parent id
header_str2 = 'Root=%s;Sampled=?' % TRACE_ID
header2 = TraceHeader.from_header_str(header_str2)
assert header2.root == TRACE_ID
assert header2.parent is None
assert header2.sampled == '?'
# missing sampled
header_str3 = 'Root=%s;Parent=%s' % (TRACE_ID, PARENT_ID)
header3 = TraceHeader.from_header_str(header_str3)
assert header3.root == TRACE_ID
assert header3.parent == PARENT_ID
assert header3.sampled is None
def inject_trace_header(headers, entity):
"""
Extract trace id, entity id and sampling decision
from the input entity and inject these information
to headers.
:param dict headers: http headers to inject
:param Entity entity: trace entity that the trace header
value generated from.
"""
if not entity:
return
to_insert = TraceHeader(
root=entity.trace_id,
parent=entity.id,
sampled=entity.sampled,
)
value = to_insert.to_header_str()
headers[http.XRAY_HEADER] = value
def _read_sqs_message():
sqs = boto3.resource("sqs", **config.BOTO_RESOURCE_KWARGS)
"""
It is possible that the queue was not created by the time
the worker launches, because the work queue creation (if needed)
and the Job spawn are on separate promises and work asyncrhonously.
This is a performance improvement but it causes the race condition above.
If this is the case, we just return an empty response
as if we didn't receive a message in this time frame.
"""
try:
queue = sqs.get_queue_by_name(QueueName=config.QUEUE_NAME)
except ClientError as e:
if e.response["Error"]["Code"] == "AWS.SimpleQueueService.NonExistentQueue":
return None
else:
raise e
message = queue.receive_messages(
WaitTimeSeconds=20, AttributeNames=["AWSTraceHeader"]
)
if not message:
return None
# Try to parse it as JSON
try:
message = message[0]
trace_header = message.attributes and message.attributes.get(
"AWSTraceHeader", None
)
if trace_header:
xray.global_sdk_config.set_sdk_enabled(True)
header = TraceHeader.from_header_str(trace_header)
trace_id = header.root
sampled = header.sampled
xray_recorder.begin_segment(
f"worker-{config.CLUSTER_ENV}-{config.SANDBOX_ID}",
traceid=trace_id,
sampling=sampled,
parent_id=header.parent,
)
body = json.loads(message.body)
info("Consumed a message from SQS.")
except Exception as e:
xray_recorder.current_segment().add_exception(e, traceback.format_exc())
info("Exception when loading message", message.body)
info("Exception:", e)
return None
finally:
message.delete()
return body
def test_invalid_str():
header = TraceHeader.from_header_str('some invalid string')
assert header.root is None
assert header.parent is None
assert header.sampled is None
def test_no_parent():
header = TraceHeader(root=TRACE_ID, sampled=1)
assert header.parent is None
assert header.to_header_str() == 'Root=%s;Sampled=1' % TRACE_ID