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

@@ -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

_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,
+    ) -> Optional[HttpServerResponseEvent]:
         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,10 @@ 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)
+                return_event.update(status, headers)
+        return None
 
     def on_exception(self, rec, start, call_event_id, exc_info):
         duration = time.monotonic() - start

appmap/flask.py

@@ -3,7 +3,8 @@
 from importlib.metadata import version
 
 import jinja2
-from flask import g, got_request_exception, request, request_finished, request_started
+from blinker import signal
+from flask import g, got_request_exception, request, request_started
 from flask.cli import ScriptInfo
 from werkzeug.exceptions import BadRequest, UnsupportedMediaType
 from werkzeug.middleware.dispatcher import DispatcherMiddleware
@@ -60,6 +61,11 @@ def request_params(req):
 NP_PARAMS = re.compile(r"<Rule '(.*?)'")
 NP_PARAM_DELIMS = str.maketrans("<>", "{}")
 
+_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 +95,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 +142,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 None
 
+        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 +168,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,10 +208,23 @@ 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
     # the app's module is specified (e.g. with the FLASK_APP env var, the `--app` flag, etc). Hook
     # 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)