[Azure Functions v2] Propagate invocation IDs in client gRPC metadata

azure-functions-durable/CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 FIXED
 
+- Async and synchronous Functions-bound clients now propagate the calling
+function's invocation ID in gRPC metadata, enabling host client-operation log
+correlation even when synchronous clients are cached across invocations.
 - Functions clients now remove HTTP URL paths from the host RPC address before
 they connect through gRPC. Both async and sync clients keep the host and port.
 - Functions orchestration queries now request at most 1000 instances per page

azure-functions-durable/azure/durable_functions/decorators/durable_app.py

@@ -479,6 +479,7 @@ def sync_client_bound(*args: Any, **kwargs: Any) -> Any:
                 client_bound = sync_client_bound
 
             set_client_metadata(client_bound)
+            client_bound = wrap_invocation(client_bound)
             if isinstance(user_fn, FunctionBuilder):
                 user_fn._function._func = client_bound  # pyright: ignore[reportPrivateUsage]
                 return wrap(user_fn)

azure-functions-durable/azure/durable_functions/internal/azurefunctions_grpc_interceptor.py

@@ -3,11 +3,30 @@
 
 from importlib.metadata import version
 
+import grpc
+import grpc.aio
+
 from durabletask.internal.grpc_interceptor import (
     DefaultAsyncClientInterceptorImpl,
     DefaultClientInterceptorImpl,
+    _AsyncClientCallDetails,  # pyright: ignore[reportPrivateUsage]
+    _ClientCallDetails,  # pyright: ignore[reportPrivateUsage]
+    _MetadataLike,  # pyright: ignore[reportPrivateUsage]
 )
 
+from .invocation import get_invocation_id
+
+
+def _with_invocation_id(metadata: _MetadataLike | None) -> _MetadataLike | None:
+    invocation_id = get_invocation_id()
+    if not invocation_id:
+        return metadata
+    key = "x-azure-functions-invocationid"
+    # Resolve at call time: synchronous clients and their interceptors are cached.
+    headers = [(name, value) for name, value in (metadata or ()) if name.lower() != key]
+    headers.append((key, invocation_id))
+    return headers
+
 
 def _build_metadata(taskhub_name: str) -> list[tuple[str, str]]:
     """Build the gRPC metadata headers sent on every Durable Functions call."""
@@ -34,6 +53,12 @@ def __init__(self, taskhub_name: str, required_query_string_parameters: str):
         self._metadata = _build_metadata(taskhub_name)
         super().__init__(self._metadata)
 
+    def _intercept_call(self, client_call_details: grpc.ClientCallDetails) -> grpc.ClientCallDetails:
+        details = super()._intercept_call(client_call_details)
+        return _ClientCallDetails(
+            details.method, details.timeout, _with_invocation_id(details.metadata),
+            details.credentials, details.wait_for_ready, details.compression)
+
 
 class AzureFunctionsAsyncDefaultClientInterceptorImpl(DefaultAsyncClientInterceptorImpl):
     """Async version of AzureFunctionsDefaultClientInterceptorImpl for use with grpc.aio channels.
@@ -46,3 +71,9 @@ def __init__(self, taskhub_name: str, required_query_string_parameters: str):
         self.required_query_string_parameters = required_query_string_parameters
         self._metadata = _build_metadata(taskhub_name)
         super().__init__(self._metadata)
+
+    async def _intercept_call(self, client_call_details: grpc.aio.ClientCallDetails) -> grpc.aio.ClientCallDetails:
+        details = await super()._intercept_call(client_call_details)
+        return _AsyncClientCallDetails(
+            details.method, details.timeout, _with_invocation_id(details.metadata),
+            details.credentials, details.wait_for_ready)

azure-functions-durable/azure/durable_functions/internal/invocation.py

@@ -78,7 +78,13 @@ def invoke() -> _Result:
     return await asyncio.get_running_loop().run_in_executor(_executor(), copied_context.run, invoke)
 
 
-def wrap_invocation(function: Callable[..., Any], trigger_name: str,
+def get_invocation_id() -> str | None:
+    """Return the originating Functions invocation ID, if one is available."""
+    context = _invocation_context.get()
+    return context.invocation_id if context is not None else None
+
+
+def wrap_invocation(function: Callable[..., Any], trigger_name: str | None = None,
                     *, registered_trigger_name: str | None = None) -> Callable[..., Any]:
     signature = inspect.signature(function)
     registered_name = registered_trigger_name or trigger_name
@@ -87,7 +93,8 @@ def wrap_invocation(function: Callable[..., Any], trigger_name: str,
         while registered_name in signature.parameters:
             registered_name = "_" + registered_name
     parameters = [
-        parameter.replace(name=registered_name) if parameter.name == trigger_name else parameter
+        parameter.replace(name=registered_name)
+        if parameter.name == trigger_name and registered_name is not None else parameter
         for parameter in signature.parameters.values()
     ]
     inject_context = "context" not in signature.parameters or trigger_name == "context"
@@ -99,27 +106,48 @@ def wrap_invocation(function: Callable[..., Any], trigger_name: str,
         parameters.insert(insertion, context_parameter)
     host_signature = signature.replace(parameters=parameters)
 
-    @wraps(function)
-    async def wrapper(*args: Any, **kwargs: Any) -> Any:
+    def bind_invocation(args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple[inspect.BoundArguments, Any]:
         bound = host_signature.bind(*args, **kwargs)
         context = bound.arguments.get("context")
         if inject_context:
             bound.arguments.pop("context", None)
-        if registered_name != trigger_name:
+        if registered_name is not None and trigger_name is not None and registered_name != trigger_name:
             bound.arguments[trigger_name] = bound.arguments.pop(registered_name)
-        token = _invocation_context.set(context)
-        try:
-            original = inspect.BoundArguments(signature, bound.arguments)
-            return await function(*original.args, **original.kwargs)
-        finally:
-            _invocation_context.reset(token)
+        return inspect.BoundArguments(signature, bound.arguments), context
+
+    wrapper: Callable[..., Any]
+    if inspect.iscoroutinefunction(function):
+        @wraps(function)
+        async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+            original, context = bind_invocation(args, kwargs)
+            token = _invocation_context.set(context)
+            try:
+                return await function(*original.args, **original.kwargs)
+            finally:
+                _invocation_context.reset(token)
+
+        wrapper = async_wrapper
+    else:
+        @wraps(function)
+        def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+            original, context = bind_invocation(args, kwargs)
+            token = _invocation_context.set(context)
+            try:
+                return function(*original.args, **original.kwargs)
+            finally:
+                _invocation_context.reset(token)
+
+        wrapper = sync_wrapper
 
     annotations = dict(function.__annotations__)
-    if registered_name != trigger_name and trigger_name in annotations:
+    if registered_name is not None and trigger_name is not None and registered_name != trigger_name \
+            and trigger_name in annotations:
         annotations[registered_name] = annotations.pop(trigger_name)
     if inject_context:
         annotations["context"] = func.Context
     wrapper.__annotations__ = annotations
     setattr(wrapper, "__signature__", host_signature)
-    setattr(wrapper, "_df_trigger_name", registered_name)
+    if registered_name is not None:
+        setattr(wrapper, "_df_trigger_name", registered_name)
+    preserve_function_source(wrapper, function)
     return wrapper

tests/azure-functions-durable/e2e/apps/dtask_style/client_routes.py

@@ -106,6 +106,26 @@ async def start_logging_filtered(
         json.dumps({"id": instance_id}), status_code=202, mimetype="application/json")
 
 
+@bp.route(route="correlation/async", methods=["POST"])
+@bp.durable_client_input(client_name="client")
+async def correlation_async(
+        req: func.HttpRequest, client: df.DurableFunctionsClient, context: func.Context) -> func.HttpResponse:
+    instance_id = await client.schedule_new_orchestration("activity_chain")
+    return func.HttpResponse(
+        json.dumps({"id": instance_id, "invocationId": context.invocation_id}),
+        status_code=202, mimetype="application/json")
+
+
+@bp.route(route="correlation/sync", methods=["POST"])
+@bp.durable_client_input(client_name="client")
+def correlation_sync(
+        req: func.HttpRequest, client: df.SyncDurableFunctionsClient, context: func.Context) -> func.HttpResponse:
+    instance_id = client.schedule_new_orchestration("activity_chain")
+    return func.HttpResponse(
+        json.dumps({"id": instance_id, "invocationId": context.invocation_id}),
+        status_code=202, mimetype="application/json")
+
+
 @bp.route(route="status/{id}", methods=["GET"])
 @bp.durable_client_input(client_name="client")
 async def get_orchestration_status(

tests/azure-functions-durable/e2e/test_dtask_client_e2e.py

@@ -9,13 +9,34 @@
 ``wait_for_orchestration_completion``.
 """
 
+from concurrent.futures import ThreadPoolExecutor
+from uuid import UUID
+
 import pytest
 
 from ._harness import http_request
 
 pytestmark = pytest.mark.functions_e2e
 
 
+@pytest.mark.parametrize("client_type", ["async", "sync"])
+def test_invocation_correlation(dtask_app, client_type):
+    def invoke(_):
+        response = http_request("POST", f"{dtask_app.base_url}/api/correlation/{client_type}")
+        assert response.status == 202
+        return response.json()
+
+    with ThreadPoolExecutor(max_workers=2) as executor:
+        results = list(executor.map(invoke, range(2)))
+    assert results[0]["invocationId"] != results[1]["invocationId"]
+    for result in results:
+        assert str(UUID(result["invocationId"])) == result["invocationId"]
+        log = dtask_app.wait_for_host_log(
+            f"Client operation 'StartOrchestration' received for instance '{result['id']}'.")
+        assert f"FunctionInvocationId: {result['invocationId']}." in log
+        dtask_app.wait_for_completion(result["id"])
+
+
 def test_suspend_and_resume(dtask_app):
     instance_id = dtask_app.start_orchestration("wait_for_approval")