fix: improve async execution ID context propagation using contextvars

- Use contextvars.copy_context() to properly propagate execution context in async functions
- Implement AsyncExecutionIdHandler to handle JSON logging with execution_id
- Redirect logging output from stderr to stdout for consistency
- Add build dependency to dev dependencies
- Update tests to reflect new logging output location

Author: taeold

pyproject.toml

@@ -65,6 +65,7 @@ functions_framework = ["py.typed"]
 [dependency-groups]
 dev = [
     "black>=23.3.0",
+    "build>=1.1.1",
     "isort>=5.11.5",
     "pretend>=1.0.9",
     "pytest>=7.4.4",

src/functions_framework/aio/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import contextvars
 import functools
 import inspect
 import json
@@ -69,11 +70,7 @@ async def _crash_handler(request, exc):
         f"Exception on {request.url.path} [{request.method}]\n{tb_text}".rstrip()
     )
 
-    if _enable_execution_id_logging():
-        log_entry = {"message": error_msg, "levelname": "ERROR"}
-        logger.error(json.dumps(log_entry))
-    else:
-        logger.error(error_msg)
+    logger.error(error_msg)
 
     headers = {_FUNCTION_STATUS_HEADER_FIELD: _CRASH}
     return Response("Internal Server Error", status_code=500, headers=headers)
@@ -132,15 +129,14 @@ async def handler(request):
             result = await function(request)
         else:
             # TODO: Use asyncio.to_thread when we drop Python 3.8 support
-            # Python 3.8 compatible version of asyncio.to_thread
             loop = asyncio.get_event_loop()
-            result = await loop.run_in_executor(None, function, request)
+            ctx = contextvars.copy_context()
+            result = await loop.run_in_executor(None, ctx.run, function, request)
         if isinstance(result, str):
             return Response(result)
         elif isinstance(result, dict):
             return JSONResponse(result)
         elif isinstance(result, tuple) and len(result) == 2:
-            # Support Flask-style tuple response
             content, status_code = result
             if isinstance(content, dict):
                 return JSONResponse(content, status_code=status_code)
@@ -187,16 +183,37 @@ def _configure_app_execution_id_logging():
     import logging
     import logging.config
 
+    class AsyncExecutionIdHandler(logging.StreamHandler):
+        def emit(self, record):
+            context = execution_id.execution_context_var.get(None)
+
+            log_entry = {
+                "message": self.format(record),
+                "severity": record.levelname,
+            }
+
+            if context and context.execution_id:
+                log_entry["logging.googleapis.com/labels"] = {
+                    "execution_id": context.execution_id
+                }
+
+            if context and context.span_id:
+                log_entry["logging.googleapis.com/spanId"] = context.span_id
+
+            try:
+                self.stream.write(json.dumps(log_entry) + "\n")
+                self.stream.flush()
+            except Exception:
+                super().emit(record)
+
     root_logger = logging.getLogger()
     root_logger.setLevel(logging.INFO)
 
-    # Remove existing handlers
     for handler in root_logger.handlers[:]:
         root_logger.removeHandler(handler)
 
-    handler = logging.StreamHandler(
-        execution_id.LoggingHandlerAddExecutionId(sys.stderr)
-    )
+    handler = AsyncExecutionIdHandler(sys.stdout)
+    handler.setFormatter(logging.Formatter("%(message)s"))
     handler.setLevel(logging.NOTSET)
     root_logger.addHandler(handler)
 

src/functions_framework/execution_id.py

@@ -88,7 +88,6 @@ def _extract_context_from_headers(headers):
     return ExecutionContext(execution_id, span_id)
 
 
-# Middleware to add execution id to request header if one does not already exist
 class WsgiMiddleware:
     def __init__(self, wsgi_app):
         self.wsgi_app = wsgi_app
@@ -101,7 +100,6 @@ def __call__(self, environ, start_response):
         return self.wsgi_app(environ, start_response)
 
 
-# ASGI Middleware to add execution id to request header if one does not already exist
 class AsgiMiddleware:
     def __init__(self, app):
         self.app = app
@@ -187,8 +185,6 @@ async def async_wrapper(request, *args, **kwargs):
             with stderr_redirect, stdout_redirect:
                 result = await func(request, *args, **kwargs)
 
-                # Only reset context on successful completion
-                # On exception, leave context available for exception handlers
                 execution_context_var.reset(token)
                 return result
 
@@ -200,8 +196,6 @@ def sync_wrapper(request, *args, **kwargs):
             with stderr_redirect, stdout_redirect:
                 result = func(request, *args, **kwargs)
 
-                # Only reset context on successful completion
-                # On exception, leave context available for exception handlers
                 execution_context_var.reset(token)
                 return result
 

tests/test_execution_id_async.py

@@ -62,7 +62,9 @@ def test_async_uncaught_exception_in_user_function_sets_execution_id(
     )
     assert resp.status_code == 500
     record = capsys.readouterr()
-    assert f'"execution_id": "{TEST_EXECUTION_ID}"' in record.err
+    assert f'"execution_id": "{TEST_EXECUTION_ID}"' in record.out
+    assert '"logging.googleapis.com/labels"' in record.out
+    assert "ZeroDivisionError" in record.out
 
 
 def test_async_print_from_user_function_sets_execution_id(capsys, monkeypatch):
@@ -99,8 +101,9 @@ def test_async_log_from_user_function_sets_execution_id(capsys, monkeypatch):
         json={"message": json.dumps({"custom-field": "some-message"})},
     )
     record = capsys.readouterr()
-    assert f'"execution_id": "{TEST_EXECUTION_ID}"' in record.err
-    assert '"custom-field": "some-message"' in record.err
+    assert f'"execution_id": "{TEST_EXECUTION_ID}"' in record.out
+    assert '\\"custom-field\\": \\"some-message\\"' in record.out
+    assert '"logging.googleapis.com/labels"' in record.out
 
 
 def test_async_user_function_can_retrieve_generated_execution_id(monkeypatch):
@@ -175,11 +178,9 @@ def make_request(client, message, exec_id):
     thread2.join()
 
     record = capsys.readouterr()
-    logs = record.err.strip().split("\n")
-    logs_as_json = [json.loads(log) for log in logs if log]
+    logs = record.out.strip().split("\n")
+    logs_as_json = [json.loads(log) for log in logs if log and log.startswith("{")]
 
-    # Check that each message appears twice (once at start, once at end of async_sleep)
-    # and that each has the correct execution ID
     message1_logs = [log for log in logs_as_json if log.get("message") == "message1"]
     message2_logs = [log for log in logs_as_json if log.get("message") == "message2"]
 
@@ -190,7 +191,6 @@ def make_request(client, message, exec_id):
         len(message2_logs) == 2
     ), f"Expected 2 logs for message2, got {len(message2_logs)}"
 
-    # Check that all message1 logs have exec-id-1
     for log in message1_logs:
         assert log["logging.googleapis.com/labels"]["execution_id"] == "exec-id-1"