-
Notifications
You must be signed in to change notification settings - Fork 557
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Propagate trace context to webhook and upload requests
Based on the implementation in #1698 for sync cog. If the request to /predict contains headers `traceparent` and `tracestate` defined by w3c Trace Context[^1] then these headers are forwarded on to the webhook and upload calls. This allows observability systems to link requests passing through cog. [^1]: https://www.w3.org/TR/trace-context/ Signed-off-by: technillogue <technillogue@gmail.com>
- Loading branch information
1 parent
8d834f0
commit 5d38ae7
Showing
4 changed files
with
163 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,54 @@ | ||
from contextlib import contextmanager | ||
from contextvars import ContextVar | ||
from typing import Generator, Optional | ||
|
||
# TypedDict was added in 3.8 | ||
from typing_extensions import TypedDict | ||
|
||
|
||
# See: https://www.w3.org/TR/trace-context/ | ||
class TraceContext(TypedDict, total=False): | ||
traceparent: str | ||
tracestate: str | ||
|
||
|
||
TRACE_CONTEXT: ContextVar[Optional[TraceContext]] = ContextVar( | ||
"trace_context", default=None | ||
) | ||
|
||
|
||
def make_trace_context( | ||
traceparent: Optional[str] = None, tracestate: Optional[str] = None | ||
) -> TraceContext: | ||
""" | ||
Creates a trace context dictionary from the given traceparent and tracestate | ||
headers. This is used to pass the trace context between services. | ||
""" | ||
ctx: TraceContext = {} | ||
if traceparent: | ||
ctx["traceparent"] = traceparent | ||
if tracestate: | ||
ctx["tracestate"] = tracestate | ||
return ctx | ||
|
||
|
||
def current_trace_context() -> Optional[TraceContext]: | ||
""" | ||
Returns the current trace context, this needs to be added via HTTP headers | ||
to all outgoing HTTP requests. | ||
""" | ||
return TRACE_CONTEXT.get() | ||
|
||
|
||
@contextmanager | ||
def trace_context(ctx: TraceContext) -> Generator[None, None, None]: | ||
""" | ||
A helper for managing the current trace context provided by the inbound | ||
HTTP request. This context is used to link requests across the system and | ||
needs to be added to all internal outgoing HTTP requests. | ||
""" | ||
t = TRACE_CONTEXT.set(ctx) | ||
try: | ||
yield | ||
finally: | ||
TRACE_CONTEXT.reset(t) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters