def lambda_wrapper(*args, **kwargs):
if _is_context_already_wrapped(*args):
return func(*args, **kwargs)
_add_wrap_flag_to_context(*args)
executed = False
ret_val = None
try:
SpansContainer.create_span(*args, is_new_invocation=True)
with lumigo_safe_execute("auto tag"):
AutoTagEvent.auto_tag_event(args[0])
SpansContainer.get_span().start(*args)
try:
executed = True
ret_val = func(*args, **kwargs)
except Exception as e:
with lumigo_safe_execute("Customer's exception"):
SpansContainer.get_span().add_exception_event(
e, inspect.trace())
raise
finally:
with lumigo_safe_execute("end"):
SpansContainer.get_span().end(ret_val, *args)
return ret_val
except Exception:
# The case where our wrapping raised an exception
if not executed:
TimeoutMechanism.stop()
get_logger().exception("exception in the wrapper",
exc_info=True)
return func(*args, **kwargs)
else:
raise
def started(self, event):
with lumigo_safe_execute("pymongo started"):
span_id = str(uuid.uuid4())
LumigoMongoMonitoring.request_to_span_id[
event.request_id] = span_id
SpansContainer.get_span().add_span({
"id":
span_id,
"type":
self.MONGO_SPAN,
"started":
get_current_ms_time(),
"databaseName":
event.database_name,
"commandName":
event.command_name,
"request":
lumigo_dumps(event.command),
"mongoRequestId":
event.request_id,
"mongoOperationId":
event.operation_id,
"mongoConnectionId":
event.connection_id,
})
def _read_stream_wrapper_generator(stream_generator, instance):
for partial_response in stream_generator:
with lumigo_safe_execute("parse response.read_chunked"):
update_event_response(None, instance.status,
dict(instance.headers.items()),
partial_response)
yield partial_response
def execute_wrapper(func, instance, args, kwargs):
result = func(*args, **kwargs)
with lumigo_safe_execute("sqlalchemy: listen to engine"):
listen(result, "before_cursor_execute", _before_cursor_execute)
listen(result, "after_cursor_execute", _after_cursor_execute)
listen(result, "handle_error", _handle_error)
return result
def _after_cursor_execute(conn, cursor, statement, parameters, context, executemany):
with lumigo_safe_execute("handle sqlalchemy after execute"):
span = SpansContainer.get_span().get_last_span()
if not span:
get_logger().warning("Redis span ended without a record on its start")
return
span.update({"ended": get_current_ms_time(), "response": ""})
def wrap_sqlalchemy():
with lumigo_safe_execute("wrap sqlalchemy"):
if importlib.util.find_spec("sqlalchemy") and listen:
get_logger().debug("wrapping sqlalchemy")
wrap_function_wrapper(
"sqlalchemy.engine.strategies", "DefaultEngineStrategy.create", execute_wrapper
)
def parse_triggered_by(event: dict):
"""
This function parses the event and build the dictionary that describes the given event.
The current possible values are:
* {triggeredBy: unknown}
* {triggeredBy: apigw, api: <host>, resource: <>, httpMethod: <>, stage: <>, identity: <>, referer: <>}
"""
with lumigo_safe_execute("triggered by"):
if not isinstance(event, dict):
if _is_step_function(event):
return _parse_step_function(event)
return None
if _is_supported_http_method(event):
return _parse_http_method(event)
if _is_load_balancer_method(event):
return _parse_load_balancer_method(event)
elif _is_supported_sns(event):
return _parse_sns(event)
elif _is_supported_streams(event):
return _parse_streams(event)
elif _is_supported_cw(event):
return _parse_cw(event)
elif _is_step_function(event):
return _parse_step_function(event)
elif _is_event_bridge(event):
return _parse_event_bridge(event)
elif _is_appsync(event):
return _parse_appsync(event)
return _parse_unknown(event)
def aiohttp_session_init_wrapper(func, instance, args, kwargs):
with lumigo_safe_execute("aiohttp aiohttp_session_init_wrapper"):
traces = kwargs.get("trace_configs") or []
if isinstance(traces, list):
traces.append(trace_config)
kwargs.update({"trace_configs": traces})
return func(*args, **kwargs)
def wrap_redis():
with lumigo_safe_execute("wrap redis"):
if importlib.util.find_spec("redis"):
get_logger().debug("wrapping redis")
wrap_function_wrapper("redis.client", "Redis.execute_command",
execute_command_wrapper)
wrap_function_wrapper("redis.client", "Pipeline.execute",
execute_wrapper)
async def on_response_chunk_received(session, trace_config_ctx, params):
with lumigo_safe_execute("aiohttp on_response_chunk_received"):
span_id = getattr(trace_config_ctx, LUMIGO_SPAN_ID_KEY)
span = SpansContainer.get_span().get_span_by_id(span_id)
http_info = span.get("info", {}).get("httpInfo", {})
http_info["response"]["body"] = concat_old_body_to_new(
http_info.get("response", {}).get("body"), params.chunk
)
def _add_wrap_flag_to_context(*args):
"""
This function is here in order to validate that we didn't already wrap this invocation
(using the sls plugin / auto instrumentation / etc.).
We are adding lumigo's flag to the context, and check it's value in _is_context_already_wrapped.
"""
if len(args) >= 2:
with lumigo_safe_execute("wrap context"):
setattr(args[1], CONTEXT_WRAPPED_BY_LUMIGO_KEY, True)
def _read_stream_wrapper_generator(stream_generator, instance):
for partial_response in stream_generator:
with lumigo_safe_execute("parse response.read_chunked"):
span_id = HttpState.response_id_to_span_id.get(
get_lumigo_connection_id(instance._original_response))
update_event_response(span_id, None, instance.status,
dict(instance.headers.items()),
partial_response)
yield partial_response
def start_new_invocation(self, event: Dict[str, str]):
with lumigo_safe_execute("Extension: start_new_invocation"):
current_bandwidth = get_current_bandwidth()
if self.request_id:
self._finish_previous_invocation(current_bandwidth)
self.sampler.start_sampling()
self.lambda_service.ready_for_next_event()
self.request_id = event.get("requestId")
self.start_time = datetime.now()
self.bandwidth = current_bandwidth
def start_extension_loop(lambda_service: LambdaService):
with lumigo_safe_execute("Extension main initialization"):
get_extension_logger().debug(
f"Extension started running with extension id: {lambda_service.extension_id}"
)
config()
extension = LumigoExtension(lambda_service)
for event in lambda_service.events_generator():
get_extension_logger().debug(f"Extension got event: {event}")
if event.get("eventType") == "INVOKE":
with lumigo_safe_execute("Extension: start new invocation"):
extension.start_new_invocation(event)
elif event.get("eventType") == "SHUTDOWN":
with lumigo_safe_execute("Extension: shutdown"):
extension.shutdown()
break
else:
get_extension_logger().error(f"Extension got unknown event: {event}")
get_extension_logger().debug("Extension finished running")
def _read_wrapper(func, instance, args, kwargs):
"""
This is the wrapper of the function that can be called only after `getresponse` was called.
"""
ret_val = func(*args, **kwargs)
if ret_val:
with lumigo_safe_execute("parse response.read"):
update_event_response(None, instance.code,
dict(instance.headers.items()), ret_val)
return ret_val
def command_finished(ret_val: Dict):
with lumigo_safe_execute("redis command finished"):
span = SpansContainer.get_span().get_last_span()
if not span:
get_logger().warning(
"Redis span ended without a record on its start")
return
span.update({
"ended": get_current_ms_time(),
"response": lumigo_dumps(copy.deepcopy(ret_val))
})
def _response_wrapper(func, instance, args, kwargs):
"""
This is the wrapper of the function that can be called only after that the http request was sent.
Note that we don't examine the response data because it may change the original behaviour (ret_val.peek()).
"""
ret_val = func(*args, **kwargs)
with lumigo_safe_execute("parse response"):
headers = dict(ret_val.headers.items())
status_code = ret_val.code
update_event_response(instance.host, status_code, headers, b"")
return ret_val
def _requests_wrapper(func, instance, args, kwargs):
"""
This is the wrapper of the function `requests.request`.
This function is being wrapped specifically because it initializes the connection by itself and parses the response,
which creates a gap from the traditional http.client wrapping.
"""
start_time = datetime.now()
ret_val = func(*args, **kwargs)
with lumigo_safe_execute("requests wrapper time updates"):
SpansContainer.get_span().update_event_times(start_time=start_time)
return ret_val
def command_failed(exception: Exception):
with lumigo_safe_execute("redis command failed"):
span = SpansContainer.get_span().get_last_span()
if not span:
get_logger().warning(
"Redis span ended without a record on its start")
return
span.update({
"ended": get_current_ms_time(),
"error": exception.args[0] if exception.args else None
})
def _requests_wrapper(func, instance, args, kwargs):
"""
This is the wrapper of the function `requests.request`.
This function is being wrapped specifically because it initializes the connection by itself and parses the response,
which creates a gap from the traditional http.client wrapping.
Moreover, these "extra" steps may raise exceptions. We should attach the error to the http span.
"""
start_time = datetime.now()
try:
ret_val = func(*args, **kwargs)
except Exception as exception:
with lumigo_safe_execute("requests wrapper exception occurred"):
method = safe_get_list(args, 0, kwargs.get("method", "")).upper()
url = safe_get_list(args, 1, kwargs.get("url"))
if Configuration.is_sync_tracer:
if HttpState.previous_request:
span = SpansContainer.get_span().get_span_by_id(
HttpState.previous_span_id)
else:
span = add_request_event(
None,
HttpRequest(
host=url,
method=method,
uri=url,
body=kwargs.get("data"),
headers=kwargs.get("headers"),
instance_id=id(instance),
),
)
span_id = span["id"]
HttpState.request_id_to_span_id[get_lumigo_connection_id(
instance)] = span_id
SpansContainer.add_exception_to_span(span, exception, [])
raise
with lumigo_safe_execute("requests wrapper time updates"):
span_id = HttpState.response_id_to_span_id.get(
get_lumigo_connection_id(ret_val.raw._original_response))
SpansContainer.get_span().update_event_times(span_id,
start_time=start_time)
return ret_val
def _headers_reminder_wrapper(func, instance, args, kwargs):
"""
This is the wrapper of the function `http.client.HTTPConnection.request` that gets the headers.
Remember the headers helps us to improve performances on requests that use this flow.
"""
with lumigo_safe_execute("add hooked data"):
setattr(
instance,
LUMIGO_HEADERS_HOOK_KEY,
HookedData(headers=kwargs.get("headers"), path=args[1]),
)
return func(*args, **kwargs)
def get_current_cpu_time() -> Optional[int]:
"""
:return: the total number of milliseconds that being used by the CPU.
"""
with lumigo_safe_execute("Extension: get cpu time"):
total = 0
with open("/proc/stat", "r") as stats:
for line in stats.readlines():
if line.startswith("cpu "):
parts = line.split()
total += (int(parts[1]) + int(parts[3])) * 10
return total
async def on_request_start(session, trace_config_ctx, params):
with lumigo_safe_execute("aiohttp on_request_start"):
span = add_request_event(
span_id=None,
parse_params=HttpRequest(
host=params.url.host,
method=params.method,
uri=str(params.url),
headers=dict(params.headers),
body=b"",
),
)
setattr(trace_config_ctx, LUMIGO_SPAN_ID_KEY, span["id"])
def execute_command_wrapper(func, instance, args, kwargs):
with lumigo_safe_execute("redis start"):
command = args[0] if args else None
request_args = args[1:] if args and len(args) > 1 else None
connection_options = instance.connection_pool.connection_kwargs
command_started(command, request_args, connection_options)
try:
ret_val = func(*args, **kwargs)
command_finished(ret_val)
return ret_val
except Exception as e:
command_failed(e)
raise
def get_current_memory() -> Optional[int]:
with lumigo_safe_execute("Extension: get meminfo"):
with open("/proc/meminfo", "r") as meminfo:
res = 0
meminfo_content = meminfo.read()
found_mem_available = re.search(MEM_AVAILABLE_PATTERN,
meminfo_content)
found_mem_total = re.search(MEM_TOTAL_PATTERN, meminfo_content)
if found_mem_total and found_mem_available:
mem_total = float(found_mem_total.group(2))
mem_available = float(found_mem_available.group(2))
res = int(100 * (mem_available / mem_total))
return res
def wrap_aiohttp():
with lumigo_safe_execute("wrap http calls"):
get_logger().debug("wrapping http requests")
if aiohttp:
trace_config = aiohttp.TraceConfig()
trace_config.on_request_start.append(on_request_start)
trace_config.on_request_chunk_sent.append(on_request_chunk_sent)
trace_config.on_request_end.append(on_request_end)
trace_config.on_response_chunk_received.append(on_response_chunk_received)
trace_config.on_request_exception.append(on_request_exception)
wrap_function_wrapper(
"aiohttp.client",
"ClientSession.__init__",
aiohttp_trace_configs_wrapper(trace_config),
)
def execute_wrapper(func, instance, args, kwargs):
with lumigo_safe_execute("redis start"):
commands = instance.command_stack
command = [cmd[0] for cmd in commands if cmd] or None
request_args = [cmd[1:] for cmd in commands if cmd and len(cmd) > 1]
connection_options = instance.connection_pool.connection_kwargs
command_started(lumigo_dumps(command), request_args,
connection_options)
try:
ret_val = func(*args, **kwargs)
command_finished(ret_val)
return ret_val
except Exception as e:
command_failed(e)
raise
def failed(self, event):
with lumigo_safe_execute("pymongo failed"):
if event.request_id not in LumigoMongoMonitoring.request_to_span_id:
get_logger().warning(
"Mongo span ended without a record on its start")
return
span_id = LumigoMongoMonitoring.request_to_span_id.pop(
event.request_id)
span = SpansContainer.get_span().get_span_by_id(span_id)
span.update({
"ended":
span["started"] + (event.duration_micros / 1000),
"error":
lumigo_dumps(event.failure),
})
def _handle_error(context):
with lumigo_safe_execute("handle sqlalchemy error"):
span = SpansContainer.get_span().get_span_by_id(_last_span_id)
if not span:
get_logger().warning(
"Redis span ended without a record on its start")
return
span.update({
"ended":
get_current_ms_time(),
"error":
lumigo_dumps({
"type": context.original_exception.__class__.__name__,
"args": context.original_exception.args,
}),
})
def _before_cursor_execute(conn, cursor, statement, parameters, context, executemany):
with lumigo_safe_execute("handle sqlalchemy before execute"):
SpansContainer.get_span().add_span(
{
"id": str(uuid.uuid4()),
"type": SQL_SPAN,
"started": get_current_ms_time(),
"connectionParameters": {
"host": conn.engine.url.host or conn.engine.url.database,
"port": conn.engine.url.port,
"database": conn.engine.url.database,
"user": conn.engine.url.username,
},
"query": lumigo_dumps(statement),
"values": lumigo_dumps(parameters),
}
)
