mirror of
https://github.com/langgenius/dify.git
synced 2026-11-13 23:06:40 +08:00
308 lines
13 KiB
Python
308 lines
13 KiB
Python
"""SQL comment scope through real Celery requests and OTel task instrumentation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Iterator
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import dataclass
|
|
from threading import Event
|
|
from typing import Protocol, cast
|
|
from uuid import uuid4
|
|
|
|
import pytest
|
|
from celery import Celery, Task
|
|
from celery.app.task import Context
|
|
from celery.result import EagerResult
|
|
from celery.signals import task_postrun, task_prerun
|
|
from opentelemetry import context, trace
|
|
from opentelemetry.instrumentation.celery import CeleryInstrumentor
|
|
from opentelemetry.sdk.metrics import MeterProvider
|
|
from opentelemetry.sdk.trace import TracerProvider
|
|
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
|
|
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
|
|
from opentelemetry.sdk.trace.sampling import ALWAYS_OFF, ALWAYS_ON
|
|
|
|
from core.app.apps.exc import GenerateTaskStoppedError
|
|
from extensions.otel.celery_sqlcommenter import (
|
|
_SQLCOMMENTER_CONTEXT_KEY,
|
|
_TOKEN_ATTR,
|
|
CelerySqlcommenterTagsDict,
|
|
_build_celery_sqlcommenter_tags,
|
|
_on_task_postrun,
|
|
_on_task_prerun,
|
|
setup_celery_sqlcommenter,
|
|
)
|
|
|
|
|
|
class _DisconnectableSignal(Protocol):
|
|
"""Celery's signal stubs omit the public disconnect method."""
|
|
|
|
def disconnect(self, receiver: Callable[..., None]) -> bool: ...
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class _Invocation:
|
|
task_object: int
|
|
request: Context
|
|
span_context: trace.SpanContext
|
|
tags: CelerySqlcommenterTagsDict
|
|
|
|
|
|
@pytest.fixture
|
|
def celery_task() -> Iterator[Task[..., object]]:
|
|
app = Celery("sqlcommenter-unit-test", broker="memory://", set_as_current=False)
|
|
|
|
@app.task(name=f"sqlcommenter.{uuid4()}")
|
|
def unused() -> None:
|
|
pass
|
|
|
|
task = app.tasks[unused.name]
|
|
task.push_request(id=str(uuid4()), retries=0, delivery_info={})
|
|
try:
|
|
yield task
|
|
finally:
|
|
task.pop_request()
|
|
app.close()
|
|
|
|
|
|
def test_build_tags_includes_celery_framework_and_task_name(celery_task: Task[..., object]) -> None:
|
|
tags = _build_celery_sqlcommenter_tags(celery_task)
|
|
|
|
assert tags["framework"].startswith("celery:")
|
|
assert tags["task_name"] == celery_task.name
|
|
|
|
|
|
@pytest.mark.parametrize("retries", [0, 3])
|
|
def test_build_tags_includes_only_nonzero_retries(celery_task: Task[..., object], retries: int) -> None:
|
|
celery_task.request.retries = retries
|
|
|
|
tags = _build_celery_sqlcommenter_tags(celery_task)
|
|
|
|
if retries:
|
|
assert tags["celery_retries"] == retries
|
|
else:
|
|
assert "celery_retries" not in tags
|
|
|
|
|
|
def test_build_tags_includes_request_routing_key(celery_task: Task[..., object]) -> None:
|
|
celery_task.request.delivery_info = {"routing_key": "workflow_based_app_execution"}
|
|
|
|
assert _build_celery_sqlcommenter_tags(celery_task)["routing_key"] == "workflow_based_app_execution"
|
|
|
|
|
|
def test_build_tags_includes_current_traceparent(celery_task: Task[..., object]) -> None:
|
|
span_context = trace.SpanContext(trace_id=0x123, span_id=0x456, is_remote=True, trace_flags=trace.TraceFlags(1))
|
|
span = trace.NonRecordingSpan(span_context)
|
|
|
|
with trace.use_span(span):
|
|
tags = _build_celery_sqlcommenter_tags(celery_task)
|
|
|
|
assert tags["traceparent"] == f"00-{span_context.trace_id:032x}-{span_context.span_id:016x}-01"
|
|
|
|
|
|
def test_handlers_restore_context_and_remove_request_token(celery_task: Task[..., object]) -> None:
|
|
previous_context = context.get_current()
|
|
|
|
_on_task_prerun(task=celery_task)
|
|
assert hasattr(celery_task.request, _TOKEN_ATTR)
|
|
assert not hasattr(celery_task, _TOKEN_ATTR)
|
|
tags = context.get_value(_SQLCOMMENTER_CONTEXT_KEY)
|
|
assert isinstance(tags, dict)
|
|
assert tags["task_name"] == celery_task.name
|
|
|
|
_on_task_postrun(task=celery_task)
|
|
_on_task_postrun(task=celery_task)
|
|
|
|
assert context.get_current() is previous_context
|
|
assert not hasattr(celery_task.request, _TOKEN_ATTR)
|
|
|
|
|
|
def test_handlers_skip_missing_task_or_request_token(celery_task: Task[..., object]) -> None:
|
|
previous_context = context.get_current()
|
|
|
|
_on_task_prerun()
|
|
_on_task_postrun()
|
|
_on_task_postrun(task=celery_task)
|
|
|
|
assert context.get_current() is previous_context
|
|
|
|
|
|
@pytest.mark.parametrize("sampled", [False, True])
|
|
@pytest.mark.parametrize("outcome", ["success", "failure", "retry", "cancel"])
|
|
@pytest.mark.parametrize("finish_order", [(0, 1), (1, 0)])
|
|
def test_overlapping_task_invocations_restore_otel_parent(
|
|
sampled: bool,
|
|
outcome: str,
|
|
finish_order: tuple[int, int],
|
|
caplog: pytest.LogCaptureFixture,
|
|
) -> None:
|
|
"""Real eager task tracing includes signals, exception and retry handling.
|
|
|
|
Two threads execute the same registered Task singleton. Each Celery request
|
|
has its own token and incoming parent. Cancellation models the Workflow's
|
|
cooperative stop returning through Celery postrun, not killing a process.
|
|
"""
|
|
provider = TracerProvider(sampler=ALWAYS_ON if sampled else ALWAYS_OFF)
|
|
exporter = InMemorySpanExporter()
|
|
provider.add_span_processor(SimpleSpanProcessor(exporter))
|
|
meter_provider = MeterProvider()
|
|
setup_celery_sqlcommenter(tracer_provider=provider, meter_provider=meter_provider)
|
|
app = Celery("sqlcommenter-concurrency-test", broker="memory://", set_as_current=False)
|
|
started = [Event(), Event()]
|
|
release = [Event(), Event()]
|
|
attempts: dict[int, list[_Invocation]] = {0: [], 1: []}
|
|
restored: dict[int, bool] = {}
|
|
parents: dict[int, trace.SpanContext] = {}
|
|
name = f"sqlcommenter.concurrent.{uuid4()}"
|
|
|
|
@app.task(bind=True, name=name, max_retries=1)
|
|
def execute_task(task: Task[..., object], marker: int) -> dict[str, str]:
|
|
span = trace.get_current_span()
|
|
attempts[marker].append(
|
|
_Invocation(
|
|
task_object=id(task),
|
|
request=task.request,
|
|
span_context=span.get_span_context(),
|
|
tags=cast(CelerySqlcommenterTagsDict, context.get_value(_SQLCOMMENTER_CONTEXT_KEY)),
|
|
)
|
|
)
|
|
if task.request.retries == 0:
|
|
started[marker].set()
|
|
assert release[marker].wait(timeout=10), "concurrent invocation was never released"
|
|
if outcome == "failure":
|
|
raise ValueError(f"synthetic failure {marker}")
|
|
if outcome == "retry" and task.request.retries == 0:
|
|
raise task.retry(countdown=0)
|
|
if outcome == "cancel":
|
|
try:
|
|
raise GenerateTaskStoppedError("synthetic cooperative stop")
|
|
except GenerateTaskStoppedError:
|
|
return {"status": "stopped"}
|
|
return {"status": "succeeded"}
|
|
|
|
task = app.tasks[name]
|
|
|
|
def invoke(marker: int) -> EagerResult[object]:
|
|
parent = provider.get_tracer(__name__).start_span(f"caller-{marker}")
|
|
parents[marker] = parent.get_span_context()
|
|
baseline = context.set_value(_SQLCOMMENTER_CONTEXT_KEY, {"caller": marker}, trace.set_span_in_context(parent))
|
|
token = context.attach(baseline)
|
|
try:
|
|
result = task.apply(
|
|
kwargs={"marker": marker}, task_id=str(uuid4()), throw=False, routing_key=f"route-{marker}"
|
|
)
|
|
restored[marker] = context.get_current() is baseline and trace.get_current_span() is parent
|
|
assert context.get_value(_SQLCOMMENTER_CONTEXT_KEY) == {"caller": marker}
|
|
return result
|
|
finally:
|
|
context.detach(token)
|
|
parent.end()
|
|
|
|
try:
|
|
with ThreadPoolExecutor(max_workers=2) as executor:
|
|
futures = [executor.submit(invoke, 0)]
|
|
try:
|
|
assert started[0].wait(timeout=10)
|
|
futures.append(executor.submit(invoke, 1))
|
|
assert started[1].wait(timeout=10)
|
|
for marker in finish_order:
|
|
release[marker].set()
|
|
result = futures[marker].result(timeout=10)
|
|
if outcome == "failure":
|
|
assert result.state == "FAILURE"
|
|
assert isinstance(result.result, ValueError)
|
|
assert str(result.result) == f"synthetic failure {marker}"
|
|
else:
|
|
assert result.state == "SUCCESS"
|
|
assert result.result == {"status": "stopped" if outcome == "cancel" else "succeeded"}
|
|
finally:
|
|
for event in release:
|
|
event.set()
|
|
|
|
assert restored == {0: True, 1: True}
|
|
assert not hasattr(task, _TOKEN_ATTR)
|
|
assert parents[0].trace_id != parents[1].trace_id
|
|
expected_attempts = 2 if outcome == "retry" else 1
|
|
for marker, invocations in attempts.items():
|
|
assert len(invocations) == expected_attempts
|
|
for index, invocation in enumerate(invocations):
|
|
span_context = invocation.span_context
|
|
tags = invocation.tags
|
|
assert invocation.task_object == id(task)
|
|
assert not hasattr(invocation.request, _TOKEN_ATTR)
|
|
assert span_context.trace_id == parents[marker].trace_id
|
|
assert span_context.trace_flags.sampled is sampled
|
|
assert tags["routing_key"] == f"route-{marker}"
|
|
assert tags["traceparent"] == (
|
|
f"00-{span_context.trace_id:032x}-{span_context.span_id:016x}-{int(span_context.trace_flags):02x}"
|
|
)
|
|
if index:
|
|
assert tags["celery_retries"] == index
|
|
else:
|
|
assert "celery_retries" not in tags
|
|
task_spans = [
|
|
span
|
|
for span in exporter.get_finished_spans()
|
|
if span.attributes is not None and span.attributes.get("celery.task_name") == name
|
|
]
|
|
assert len(task_spans) == (2 * expected_attempts if sampled else 0)
|
|
for span in task_spans:
|
|
assert span.parent is not None
|
|
assert span.context is not None
|
|
parent = next(parent for parent in parents.values() if parent.trace_id == span.context.trace_id)
|
|
assert span.parent.span_id == parent.span_id
|
|
assert "Failed to detach context" not in caplog.text
|
|
finally:
|
|
cast(_DisconnectableSignal, task_prerun).disconnect(_on_task_prerun)
|
|
cast(_DisconnectableSignal, task_postrun).disconnect(_on_task_postrun)
|
|
CeleryInstrumentor().uninstrument()
|
|
app.close()
|
|
provider.shutdown()
|
|
meter_provider.shutdown()
|
|
|
|
|
|
def test_external_instrumentation_keeps_its_provider_and_context(caplog: pytest.LogCaptureFixture) -> None:
|
|
"""An existing instrumentor has no public signal-ordering API to nest under."""
|
|
external_provider = TracerProvider()
|
|
external_exporter = InMemorySpanExporter()
|
|
external_provider.add_span_processor(SimpleSpanProcessor(external_exporter))
|
|
unused_provider = TracerProvider()
|
|
unused_exporter = InMemorySpanExporter()
|
|
unused_provider.add_span_processor(SimpleSpanProcessor(unused_exporter))
|
|
meter_provider = MeterProvider()
|
|
instrumentor = CeleryInstrumentor()
|
|
instrumentor.instrument(tracer_provider=external_provider, meter_provider=meter_provider)
|
|
setup_celery_sqlcommenter(tracer_provider=unused_provider, meter_provider=meter_provider)
|
|
app = Celery("sqlcommenter-external-instrumentation", broker="memory://", set_as_current=False)
|
|
baseline_tags = {"outer": "external-owner"}
|
|
|
|
@app.task(name=f"sqlcommenter.external.{uuid4()}")
|
|
def execute_task() -> None:
|
|
assert context.get_value(_SQLCOMMENTER_CONTEXT_KEY) == baseline_tags
|
|
assert not hasattr(app.tasks[execute_task.name].request, _TOKEN_ATTR)
|
|
|
|
parent = external_provider.get_tracer(__name__).start_span("external-parent")
|
|
baseline = context.set_value(_SQLCOMMENTER_CONTEXT_KEY, baseline_tags, trace.set_span_in_context(parent))
|
|
token = context.attach(baseline)
|
|
try:
|
|
result = execute_task.apply(throw=True)
|
|
|
|
assert result.state == "SUCCESS"
|
|
assert context.get_current() is baseline
|
|
assert trace.get_current_span() is parent
|
|
spans = external_exporter.get_finished_spans()
|
|
assert len(spans) == 1
|
|
assert spans[0].parent is not None
|
|
assert spans[0].parent.span_id == parent.get_span_context().span_id
|
|
assert not unused_exporter.get_finished_spans()
|
|
assert "skipping SQL commenter signal registration" in caplog.text
|
|
assert "Failed to detach context" not in caplog.text
|
|
finally:
|
|
context.detach(token)
|
|
parent.end()
|
|
instrumentor.uninstrument()
|
|
app.close()
|
|
external_provider.shutdown()
|
|
unused_provider.shutdown()
|
|
meter_provider.shutdown()
|