Add OpenTelemetry integration

[ducdetronquito]

Hi !

I'm currently evaluating procrastinate for my company as a replacement for Celery and I needed to have an OpenTelemetry integration where a span would be created when a task is deferred and when a task is ran.

Like many other OpenTelemetry integrations it relies on monkey patching, which is not very satisfying to me, but hey "it works™" !
I just wanted to share with you the task middleware that I made in order to discuss about it and have your feedbacks:

• Is it useful to others ?
• How could it be improved ?
• Would it make sens to ship an OpenTelemetry integration with procrastinate, like the one for django?
• Would it make sens to ship it as an external package (like procrastinate-opentelemetry or something) ?
• How to make it works for async stack (I only work on a sync Django codebase right now, so I have not tested it at all) ?

⚠️ Bare in mind that it's more a proof of concept than a final implementation

import functools
from opentelemetry import trace, propagate, context
from procrastinate import Blueprint, JobContext, tasks, types


def traced_task(blueprint: Blueprint):
    """
    Declare a procrastinate task that produce OpenTelemetry spans on defer/run.

    It's a drop-in replacement for procrastinate's `app.task` or `blueprint.task` decorators.

    This works by monkey patching procrastinate Task instances:

    - We monkey patch `Task.defer` method to retrieve the current trace context,
      create a child span to inform the task has been defered, and pass the trace context
      silently as the second argument of the task (`trace_context`).
    - We decorate the user defined task function to retrieve the trace context for the
      `trace_context` task parameters, and create a child span to inform the task is being
      processed.
    """

    def traced_task_wrapper(user_func=None, **kwargs):
        def task_wrapper(func):
            def traced_task_run(
                job_context: JobContext,
                trace_context: dict[str, str],
                *job_args,
                **job_kwargs,
            ):
                task = job_context.task
                if task is None:
                    raise Exception(
                        "Procrastinate: job context has no related task; wtf!"
                    )

                token = context.attach(propagate.extract(trace_context))
                tracer = _get_or_set_tracer_on_blueprint(task.blueprint)
                try:
                    with tracer.start_as_current_span(
                        f"run/{task.name}",
                        kind=trace.SpanKind.CONSUMER,
                        record_exception=True,
                    ) as span:
                        if span.is_recording():
                            _set_procrastinate_task_attributes(span, task)
                        result = func(*job_args, **job_kwargs)
                    return result
                finally:
                    context.detach(token)

            wrapped_func = functools.update_wrapper(traced_task_run, func, updated=())
            kwargs["pass_context"] = True
            task: tasks.Task = blueprint.task(**kwargs)(wrapped_func)

            task.__original_defer = task.defer

            def traced_defer(**task_kwargs: types.JSONValue) -> int:
                tracer = _get_or_set_tracer_on_blueprint(task.blueprint)
                trace_context = {}
                with tracer.start_as_current_span(
                    f"defer/{task.name}",
                    kind=trace.SpanKind.PRODUCER,
                    record_exception=True,
                ) as span:
                    if span.is_recording():
                        propagate.inject(trace_context)
                        _set_procrastinate_task_attributes(span, task)

                    task_kwargs["trace_context"] = trace_context
                    return task.__original_defer(**task_kwargs)

            task.defer = traced_defer

            return task

        if not user_func:
            return task_wrapper

        return task_wrapper(user_func)

    return traced_task_wrapper


def _get_or_set_tracer_on_blueprint(blueprint: Blueprint) -> trace.Tracer:
    """
    We cache the OpenTelemetry tracer on the blueprint, in order to be reused for the tasks of this blueprint.
    In all honesty, I'm not even sure it's worth it or not.
    """
    tracer = getattr(blueprint, "_tracer", None)
    if tracer is None:
        tracer = trace.get_tracer(__name__)
        setattr(blueprint, "_tracer", tracer)

    return tracer


def _set_procrastinate_task_attributes(span: trace.Span, task: tasks.Task) -> None:
    attributes = {
        "procrastinate.name": task.name,
        "procrastinate.queue": task.queue,
    }
    if task.lock is not None:
        attributes["procrastinate.lock"] = task.lock

    if task.queueing_lock is not None:
        attributes["procrastinate.queueing_lock"] = task.lock

    span.set_attributes(attributes)

Have a nice day 🌴

[ewjoachim]

Hello !

It's an excellent initiative, opening this discussion. I'm currently dissatisfied by the current state of observability in Procrastinate. I think investing in making it easier to integrate OpenTelemetry or other tools should definitely be a step in the right direction, for traces, but also metrics and logs (as far as I can tell, there's no code to modify to integrate logs properly, but it's worth testing that autoinstrumentation works well and maybe point to the relevant part of OTel's doc).

Of course, this can work be split in multiple parts. I've never played with Open Telemetry before, so it's all very abstract for me, but I'll do my best to answer your questions:

Is it useful to others ?

OpenTelemetry in general, yes, though I believe it's probably sane to keep it optional. I believe some people won't have the rest of the infra ready for this.
This middleware in particular: I'd tend to favor either integrating a solution that wouldn't monkey patch. Maybe that could lead to making the worker (and/or the deferrer) include hooks that would allow integrating other observability tools too.

How could it be improved ?

Let's discuss what we can do in procrastinate so that this is less invasive, whether we decide to ship opentelemetry optional code or just make it easier to people to integrate these kind of libs without commitments to otel in particular.
We have a contrib namespace that would make sense for this to live outside of core but be shipped in the package.

Would it make sens to ship an OpenTelemetry integration with procrastinate, like the one for django?

(damn, I'm answering questions in advance)

Would it make sens to ship it as an external package (like procrastinate-opentelemetry or something) ?

I'm perfectly ok with that too, and I'll add multiple steps in between:

• It could live in this codebase but be provided as a different PyPI project
• It could live in another repo under the procrastinate-org umbrella
• It could live elsewhere entirely, as a 3rd-party package

That said, for now, I'm ok having live in the contrib section, with potential optional dependencies (extra etc)

How to make it works for async stack (I only work on a sync Django codebase right now, so I have not tested it at all) ?

Aha, that's a complex question. I don't know otel enough to know if it's simple or not, so it's hard to say. In procrastinate, we usually try to make things async at core, and use asgiref.async_to_sync to make it available to sync programs. Note that some codepaths we try to leave purely sync (such as deferring tasks from a sync context). I believe I can help with that.

Let's keep this discussion moving. It's a bit blurry in my mind so far, and I'll let this sink in a bit. I guess I'd love to gather opinions from other frequent users/occasional contributors and get a sense of how we can work together on this to manage expectations and avoid spending time on some code until we know it's going to get merged.

Maybe something that could be cool to better understand this could be to dust off the docker-compose file and setup a small observability stack that would gather logs, metrics, traces in a toy app (such as the Django app in the demos folder). That could also give us a playground for testing an equivalent async stack later.

Let us play :) I don't expect an immediate decision but it's an opportunity to evolve the app in a good future-friendly way.

[ashleyheath]

I'd be keen to see this as it was on my todo list, so +1 from me.

One suggestion - it would be useful to go a level higher and wrap the execution of the job so that the job.id can be included in the span attributes (and perhaps also whether the job will retry).

[frisia-mtz]

Hello everybody!

My take on an OpenTelemetry integration is the following:
I am monkey patching internal methods from the procrastinate.worker.Worker class and wrap them with a trace context. This way the whole job execution is traced, and the job ID can be included in the span attributes. What's missing is the ability to connect task defers and runs.

This is the example setup:

def _set_task_attributes(span, task) -> None:
    """Set span attributes of the Procrastinate task.

    Args:
        span (opentelemetry.trace.Span): Current recording OTel span.
        task (procrastinate.tasks.Task): Running Procrastinate task in the span.
    """

    attributes = {
        "task.name": task.name,
        "task.priority": task.priority,
    }
    if task.lock is not None:
        attributes["task.lock"] = task.lock

    if task.queueing_lock is not None:
        attributes["task.queueing_lock"] = task.lock

    span.set_attributes(attributes)


def _set_job_context_attributes(span, context) -> None:
    """Set span attributes of the Procrastinate job context, including its task.

    Args:
        span (opentelemetry.trace.Span): Current recording OTel span.
        context (procrastinate.job_context.JobContext): Job context.
    """
    job = context.job
    attributes = {
        "job.id": job.id,
        "job.priority": job.priority,
        "job.queue": job.queue,
        "worker.name": context.worker_name,
    }
    if job.queueing_lock is not None:
        attributes["job.queueing_lock"] = job.queueing_lock
    if job.scheduled_at is not None:
        attributes["job.scheduled_at"] = job.scheduled_at
    span.set_attributes(attributes)
    _set_task_attributes(span, context.task)


def setup_worker_tracing():
    """Setup tracing of worker jobs with OpenTelemetry.

    Monkey patches internal Procrastinate Worker functions.
    """
    from opentelemetry import trace
    from procrastinate import exceptions, job_context, jobs, worker

    tracer = trace.get_tracer("worker")

    # monkey patch job processing for starting spans
    original_process_job = worker.Worker._process_job
    # monkey patch job outcome for recording exceptions and duration
    original_log_job_outcome = worker.Worker._log_job_outcome

    async def traced_process_job(self, context: job_context.JobContext):
        # start OTel span with context
        with tracer.start_as_current_span(
            context.task.name,
            kind=trace.SpanKind.CONSUMER,
            record_exception=True,
            end_on_exit=False,  # ends in traced_log_job_outcome
        ) as span:
            if span.is_recording():
                _set_job_context_attributes(span, context)

            await original_process_job(self, context)

    def traced_log_job_outcome(
        self,
        status: jobs.Status,
        context: job_context.JobContext,
        job_result: job_context.JobResult | None,
        job_retry: exceptions.JobRetry | None,
        exc_info: bool | BaseException = False,
    ):
        task_span = trace.get_current_span()
        if isinstance(exc_info, BaseException):
            task_span.record_exception(exc_info)
            task_span.set_status(trace.StatusCode.ERROR)
            if job_retry is not None:
                task_span.set_attribute("job.retry", str(job_retry.retry_decision.retry_at))

        # the current time is precise enough for task duration
        task_span.end()
        original_log_job_outcome(self, status, context, job_result, job_retry, exc_info)

    worker.Worker._process_job = traced_process_job
    worker.Worker._log_job_outcome = traced_log_job_outcome