Distributed tracing with contextvars in Python

import contextvars
import uuid
import time

_trace_context = contextvars.ContextVar("trace_context", default=None)


class TraceContext:
    def __init__(self, trace_id, parent_span_id):
        self.trace_id = trace_id
        self.parent_span_id = parent_span_id
        self.span_id = uuid.uuid4().hex[:16]
        self.started_at = time.time()

    def __repr__(self):
        return f"trace={self.trace_id} parent_span={self.parent_span_id} span={self.span_id}"


def start_span(name):
    parent = _trace_context.get()
    if parent:
        context = TraceContext(parent.trace_id, parent.span_id)
    else:
        context = TraceContext(uuid.uuid4().hex[:16], "root")

    token = _trace_context.set(context)
    print(f"START {name}: {context}")
    return token


def end_span(name, token):
    context = _trace_context.get()
    duration_ms = (time.time() - context.started_at) * 1000
    print(f"END {name}: {context} took {duration_ms:.2f}ms")
    _trace_context.reset(token)


def process_payment():
    token = start_span("payment")
    time.sleep(0.01)
    end_span("payment", token)


def main():
    token = start_span("order")
    process_payment()
    end_span("order", token)


if __name__ == "__main__":
    main()