Skip to content

Commit 2e49481

Browse files
andystaplesCopilot
andauthored
[Azure Functions v2] Propagate invocation IDs in client gRPC metadata (#328)
* Propagate Functions invocation IDs in client gRPC metadata Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> * Preserve bindings named context during invocation correlation Only capture runtime-shaped invocation contexts, leaving client and activity bindings unchanged when no runtime context is available. Cover sync and async RPCs and both activity decorator orders. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>
1 parent 2854f2a commit 2e49481

7 files changed

Lines changed: 401 additions & 14 deletions

File tree

‎azure-functions-durable/CHANGELOG.md‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
99

1010
FIXED
1111

12+
- Async and synchronous Functions-bound clients now propagate the calling
13+
function's invocation ID in gRPC metadata, enabling host client-operation log
14+
correlation even when synchronous clients are cached across invocations.
1215
- Functions clients now remove HTTP URL paths from the host RPC address before
1316
they connect through gRPC. Both async and sync clients keep the host and port.
1417
- Functions orchestration queries now request at most 1000 instances per page

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -479,6 +479,7 @@ def sync_client_bound(*args: Any, **kwargs: Any) -> Any:
479479
client_bound = sync_client_bound
480480

481481
set_client_metadata(client_bound)
482+
client_bound = wrap_invocation(client_bound)
482483
if isinstance(user_fn, FunctionBuilder):
483484
user_fn._function._func = client_bound # pyright: ignore[reportPrivateUsage]
484485
return wrap(user_fn)

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,30 @@
33

44
from importlib.metadata import version
55

6+
import grpc
7+
import grpc.aio
8+
69
from durabletask.internal.grpc_interceptor import (
710
DefaultAsyncClientInterceptorImpl,
811
DefaultClientInterceptorImpl,
12+
_AsyncClientCallDetails, # pyright: ignore[reportPrivateUsage]
13+
_ClientCallDetails, # pyright: ignore[reportPrivateUsage]
14+
_MetadataLike, # pyright: ignore[reportPrivateUsage]
915
)
1016

17+
from .invocation import get_invocation_id
18+
19+
20+
def _with_invocation_id(metadata: _MetadataLike | None) -> _MetadataLike | None:
21+
invocation_id = get_invocation_id()
22+
if not invocation_id:
23+
return metadata
24+
key = "x-azure-functions-invocationid"
25+
# Resolve at call time: synchronous clients and their interceptors are cached.
26+
headers = [(name, value) for name, value in (metadata or ()) if name.lower() != key]
27+
headers.append((key, invocation_id))
28+
return headers
29+
1130

1231
def _build_metadata(taskhub_name: str) -> list[tuple[str, str]]:
1332
"""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):
3453
self._metadata = _build_metadata(taskhub_name)
3554
super().__init__(self._metadata)
3655

56+
def _intercept_call(self, client_call_details: grpc.ClientCallDetails) -> grpc.ClientCallDetails:
57+
details = super()._intercept_call(client_call_details)
58+
return _ClientCallDetails(
59+
details.method, details.timeout, _with_invocation_id(details.metadata),
60+
details.credentials, details.wait_for_ready, details.compression)
61+
3762

3863
class AzureFunctionsAsyncDefaultClientInterceptorImpl(DefaultAsyncClientInterceptorImpl):
3964
"""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):
4671
self.required_query_string_parameters = required_query_string_parameters
4772
self._metadata = _build_metadata(taskhub_name)
4873
super().__init__(self._metadata)
74+
75+
async def _intercept_call(self, client_call_details: grpc.aio.ClientCallDetails) -> grpc.aio.ClientCallDetails:
76+
details = await super()._intercept_call(client_call_details)
77+
return _AsyncClientCallDetails(
78+
details.method, details.timeout, _with_invocation_id(details.metadata),
79+
details.credentials, details.wait_for_ready)

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

Lines changed: 48 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,13 @@ def invoke() -> _Result:
7878
return await asyncio.get_running_loop().run_in_executor(_executor(), copied_context.run, invoke)
7979

8080

81-
def wrap_invocation(function: Callable[..., Any], trigger_name: str,
81+
def get_invocation_id() -> str | None:
82+
"""Return the originating Functions invocation ID, if one is available."""
83+
context = _invocation_context.get()
84+
return context.invocation_id if context is not None else None
85+
86+
87+
def wrap_invocation(function: Callable[..., Any], trigger_name: str | None = None,
8288
*, registered_trigger_name: str | None = None) -> Callable[..., Any]:
8389
signature = inspect.signature(function)
8490
registered_name = registered_trigger_name or trigger_name
@@ -87,7 +93,8 @@ def wrap_invocation(function: Callable[..., Any], trigger_name: str,
8793
while registered_name in signature.parameters:
8894
registered_name = "_" + registered_name
8995
parameters = [
90-
parameter.replace(name=registered_name) if parameter.name == trigger_name else parameter
96+
parameter.replace(name=registered_name)
97+
if parameter.name == trigger_name and registered_name is not None else parameter
9198
for parameter in signature.parameters.values()
9299
]
93100
inject_context = "context" not in signature.parameters or trigger_name == "context"
@@ -99,27 +106,54 @@ def wrap_invocation(function: Callable[..., Any], trigger_name: str,
99106
parameters.insert(insertion, context_parameter)
100107
host_signature = signature.replace(parameters=parameters)
101108

102-
@wraps(function)
103-
async def wrapper(*args: Any, **kwargs: Any) -> Any:
109+
def bind_invocation(
110+
args: tuple[Any, ...], kwargs: dict[str, Any]) -> tuple[inspect.BoundArguments, func.Context | None]:
104111
bound = host_signature.bind(*args, **kwargs)
105-
context = bound.arguments.get("context")
112+
context: func.Context | None = bound.arguments.get("context")
113+
# A binding may also be named context. Worker contexts implement the
114+
# runtime contract without inheriting from azure.functions.Context.
115+
if not isinstance(getattr(context, "invocation_id", None), str) \
116+
or not hasattr(context, "thread_local_storage"):
117+
context = None
106118
if inject_context:
107119
bound.arguments.pop("context", None)
108-
if registered_name != trigger_name:
120+
if registered_name is not None and trigger_name is not None and registered_name != trigger_name:
109121
bound.arguments[trigger_name] = bound.arguments.pop(registered_name)
110-
token = _invocation_context.set(context)
111-
try:
112-
original = inspect.BoundArguments(signature, bound.arguments)
113-
return await function(*original.args, **original.kwargs)
114-
finally:
115-
_invocation_context.reset(token)
122+
return inspect.BoundArguments(signature, bound.arguments), context
123+
124+
wrapper: Callable[..., Any]
125+
if inspect.iscoroutinefunction(function):
126+
@wraps(function)
127+
async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
128+
original, context = bind_invocation(args, kwargs)
129+
token = _invocation_context.set(context)
130+
try:
131+
return await function(*original.args, **original.kwargs)
132+
finally:
133+
_invocation_context.reset(token)
134+
135+
wrapper = async_wrapper
136+
else:
137+
@wraps(function)
138+
def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
139+
original, context = bind_invocation(args, kwargs)
140+
token = _invocation_context.set(context)
141+
try:
142+
return function(*original.args, **original.kwargs)
143+
finally:
144+
_invocation_context.reset(token)
145+
146+
wrapper = sync_wrapper
116147

117148
annotations = dict(function.__annotations__)
118-
if registered_name != trigger_name and trigger_name in annotations:
149+
if registered_name is not None and trigger_name is not None and registered_name != trigger_name \
150+
and trigger_name in annotations:
119151
annotations[registered_name] = annotations.pop(trigger_name)
120152
if inject_context:
121153
annotations["context"] = func.Context
122154
wrapper.__annotations__ = annotations
123155
setattr(wrapper, "__signature__", host_signature)
124-
setattr(wrapper, "_df_trigger_name", registered_name)
156+
if registered_name is not None:
157+
setattr(wrapper, "_df_trigger_name", registered_name)
158+
preserve_function_source(wrapper, function)
125159
return wrapper

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

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -106,6 +106,26 @@ async def start_logging_filtered(
106106
json.dumps({"id": instance_id}), status_code=202, mimetype="application/json")
107107

108108

109+
@bp.route(route="correlation/async", methods=["POST"])
110+
@bp.durable_client_input(client_name="client")
111+
async def correlation_async(
112+
req: func.HttpRequest, client: df.DurableFunctionsClient, context: func.Context) -> func.HttpResponse:
113+
instance_id = await client.schedule_new_orchestration("activity_chain")
114+
return func.HttpResponse(
115+
json.dumps({"id": instance_id, "invocationId": context.invocation_id}),
116+
status_code=202, mimetype="application/json")
117+
118+
119+
@bp.route(route="correlation/sync", methods=["POST"])
120+
@bp.durable_client_input(client_name="client")
121+
def correlation_sync(
122+
req: func.HttpRequest, client: df.SyncDurableFunctionsClient, context: func.Context) -> func.HttpResponse:
123+
instance_id = client.schedule_new_orchestration("activity_chain")
124+
return func.HttpResponse(
125+
json.dumps({"id": instance_id, "invocationId": context.invocation_id}),
126+
status_code=202, mimetype="application/json")
127+
128+
109129
@bp.route(route="status/{id}", methods=["GET"])
110130
@bp.durable_client_input(client_name="client")
111131
async def get_orchestration_status(

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,13 +9,34 @@
99
``wait_for_orchestration_completion``.
1010
"""
1111

12+
from concurrent.futures import ThreadPoolExecutor
13+
from uuid import UUID
14+
1215
import pytest
1316

1417
from ._harness import http_request
1518

1619
pytestmark = pytest.mark.functions_e2e
1720

1821

22+
@pytest.mark.parametrize("client_type", ["async", "sync"])
23+
def test_invocation_correlation(dtask_app, client_type):
24+
def invoke(_):
25+
response = http_request("POST", f"{dtask_app.base_url}/api/correlation/{client_type}")
26+
assert response.status == 202
27+
return response.json()
28+
29+
with ThreadPoolExecutor(max_workers=2) as executor:
30+
results = list(executor.map(invoke, range(2)))
31+
assert results[0]["invocationId"] != results[1]["invocationId"]
32+
for result in results:
33+
assert str(UUID(result["invocationId"])) == result["invocationId"]
34+
log = dtask_app.wait_for_host_log(
35+
f"Client operation 'StartOrchestration' received for instance '{result['id']}'.")
36+
assert f"FunctionInvocationId: {result['invocationId']}." in log
37+
dtask_app.wait_for_completion(result["id"])
38+
39+
1940
def test_suspend_and_resume(dtask_app):
2041
instance_id = dtask_app.start_orchestration("wait_for_approval")
2142

0 commit comments

Comments
 (0)