From 4f4f6eb818b1f973516911358825938ff6243626 Mon Sep 17 00:00:00 2001 From: Alan Potter Date: Fri, 5 Jul 2024 16:57:59 -0400 Subject: [PATCH] fix: fix Flask event ordering problem If the config specifies that flask should be instrumented, http_server_response events get recorded between the call and return events for Flask.finalize_request. These changes hook finalize_request, and ensure that the events get ordered correctly. --- _appmap/event.py | 18 ++++++++-------- _appmap/web_framework.py | 22 +++++++++++++------ appmap/flask.py | 46 +++++++++++++++++++++++++++++++++++----- 3 files changed, 66 insertions(+), 20 deletions(-) diff --git a/_appmap/event.py b/_appmap/event.py index a333fc78..87a0f310 100644 --- a/_appmap/event.py +++ b/_appmap/event.py @@ -479,18 +479,18 @@ class HttpResponseEvent(ReturnEvent): def __init__(self, status_code, headers=None, **kwargs): super().__init__(**kwargs) + self.response = {} + self.update(status_code, headers) - response = {"status_code": status_code} + def update(self, status_code, headers): + if status_code is not None: + self.response.update({"status_code": status_code}) if headers is not None: - response.update( - { - "mime_type": headers.get("Content-Type"), - "headers": none_if_empty(dict(headers)), - } - ) - - self.response = compact_dict(response) + self.response.update({"mime_type": headers.get("Content-Type")}) + updated_headers = dict(headers) + if len(updated_headers) > 0: + self.response.update({"headers": updated_headers}) # pylint: disable=too-few-public-methods diff --git a/_appmap/web_framework.py b/_appmap/web_framework.py index aebf08a7..ddfccaf5 100644 --- a/_appmap/web_framework.py +++ b/_appmap/web_framework.py @@ -143,8 +143,9 @@ def before_request_main(self, rec, req: Any) -> Tuple[float, int]: raise NotImplementedError # pylint: disable=too-many-arguments - def after_request_main(self, rec, status, headers, start, call_event_id) -> None: - + def after_request_main( + self, rec, status, headers, start, call_event_id + ) -> HttpServerResponseEvent: duration = time.monotonic() - start return_event = HttpServerResponseEvent( parent_id=call_event_id, @@ -153,6 +154,7 @@ def after_request_main(self, rec, status, headers, start, call_event_id) -> None headers=headers, ) rec.add_event(return_event) + return return_event def __init__(self, framework_name): self.record_url = "/_appmap/record" @@ -191,17 +193,21 @@ def after_request_hook( headers, start, call_event_id, - ) -> None: + return_event=None, + ) -> HttpServerResponseEvent | None: if request_path == self.record_url: - return + return None env = Env.current if env.enables("requests"): rec = request_recorder.get() assert rec is not None + if return_event is None: + return self.after_request_main(rec, status, headers, start, call_event_id) + try: - self.after_request_main(rec, status, headers, start, call_event_id) + return_event.update(status, headers) output_dir = Env.current.output_dir / "requests" create_appmap_file( @@ -221,7 +227,11 @@ def after_request_hook( rec = Recorder.get_global() assert rec is not None if rec.get_enabled(): - self.after_request_main(rec, status, headers, start, call_event_id) + if return_event is None: + return self.after_request_main(rec, status, headers, start, call_event_id) + else: + return_event.update(status, headers) + return None def on_exception(self, rec, start, call_event_id, exc_info): duration = time.monotonic() - start diff --git a/appmap/flask.py b/appmap/flask.py index ebc56df6..530b74ea 100644 --- a/appmap/flask.py +++ b/appmap/flask.py @@ -2,9 +2,12 @@ import time from importlib.metadata import version +from blinker import signal + import jinja2 from flask import g, got_request_exception, request, request_finished, request_started from flask.cli import ScriptInfo +from flask.app import Flask # noqa: F401 from werkzeug.exceptions import BadRequest, UnsupportedMediaType from werkzeug.middleware.dispatcher import DispatcherMiddleware @@ -60,6 +63,11 @@ def request_params(req): NP_PARAMS = re.compile(r"", "{}") +_before_finalize = signal("_appmap_before_finalize") +_before_exception = signal("_appmap_before_exception") +_response_ready = signal("_appmap_response_ready") + + class AppmapFlask(AppmapMiddleware): """ A Flask extension to add remote recording to an application. @@ -89,8 +97,9 @@ def init_app(self): ) setattr(self.app, REMOTE_ENABLED_ATTR, remote_enabled) - request_started.connect(self.request_started, self.app, weak=False) - request_finished.connect(self.request_finished, self.app, weak=False) + request_started.connect(self.request_started, sender=self.app, weak=False) + _before_finalize.connect(self.before_finalize, sender=self.app, weak=False) + _response_ready.connect(self.response_ready, sender=self.app, weak=False) got_request_exception.connect(self.got_request_exception, self.app, weak=False) setattr(self.app, REQUEST_ENABLED_ATTR, True) @@ -135,10 +144,24 @@ def before_request_main(self, rec, req): g._appmap_request_start = time.monotonic() # pylint: disable=protected-access return None, None - def request_finished(self, _, response, **__): + def before_finalize(self, _, **__): if not self.should_record: - return response + return + + return_event = self.after_request_hook( + request.path, + request.method, + request.base_url, + None, + None, + g._appmap_request_start, # pylint: disable=protected-access + g._appmap_request_event.id, # pylint: disable=protected-access + ) + return return_event + def response_ready(self, _, **kw): + response = kw["response"] + return_event = kw["return_event"] self.after_request_hook( request.path, request.method, @@ -147,8 +170,8 @@ def request_finished(self, _, response, **__): response.headers, g._appmap_request_start, # pylint: disable=protected-access g._appmap_request_event.id, # pylint: disable=protected-access + return_event, ) - return response def got_request_exception(self, _, exception): self.on_exception( @@ -187,6 +210,18 @@ def install_extension(wrapped, _, args, kwargs): return app +def _finalize_request(wrapped, inst, args, kwargs): + if not Env.current.enabled: + return wrapped(*args, **kwargs) + + results = _before_finalize.send(inst) + resp = wrapped(*args, **kwargs) + if len(results) > 0: + assert len(results) == 1 + _, return_event = results[0] + _response_ready.send(inst, response=resp, return_event=return_event) + return resp + if Env.current.enabled: # ScriptInfo.load_app is the function that's used by the Flask cli to load an app, no matter how @@ -194,3 +229,4 @@ def install_extension(wrapped, _, args, kwargs): # it so it installs our extension on the app. load_app = wrapt.wrap_function_wrapper("flask.cli", "ScriptInfo.load_app", install_extension) ScriptInfo.load_app = load_app # type: ignore[method-assign] + wrapt.wrap_function_wrapper("flask.app", "Flask.finalize_request", _finalize_request)