fix(celery): close celery.apply spans even without after_task_publish, when using apply_async

[wantsui]

The instrumentation for the Celery integration relies on various Celery signals in order to start and end the span when calling on apply_async.

The integration can fail if the expected signals don't trigger, which can lead to broken context propagation (and unexpected traces).

Example:

• dd-trace-py expects the signal before_task_publish to start the span then after_task_publish to close the span. If the after_task_publish signal never gets called (which can happen if a Celery exception occurs while processing the app), then the span won't finish.
• The same thing above can also happen to task_prerun and task_postrun.

Solution

This PR patches apply_async so that there is a check to see if there is a span lingering around and closes it when apply_task is called.

If an internal exception happens, the error will be marked on the celery.apply span.

To track this, I added new logs in debug mode:

The after_task_publish signal was not called, so manually closing span

and

The task_postrun signal was not called, so manually closing span

There's a related PR #10848 that works to improve how we extract information based on the protocols, that also affects when spans get closed or not.

Special Thanks:

• Thanks to @tabgok for going through this with me in great detail!
• @timmc-edx for helping us track it down!

APMS-13158

Checklist

• PR author has checked that all the criteria below are met
• The PR description includes an overview of the change
• The PR description articulates the motivation for the change
• The change includes tests OR the PR description describes a testing strategy
• The PR description notes risks associated with the change, if any
• Newly-added code is easy to change
• The change follows the library release note guidelines
• The change includes or references documentation updates if necessary
• Backport labels are set (if applicable)

Reviewer Checklist

• Reviewer has checked that all the criteria below are met
• Title is accurate
• All changes are related to the pull request's stated goal
• Avoids breaking API changes
• Testing strategy adequately addresses listed risks
• Newly-added code is easy to change
• Release note makes sense to a user of the library
• If necessary, author has acknowledged and discussed the performance implications of this PR as reported in the benchmarks PR comment
• Backport labels are set in a manner that is consistent with the release branch maintenance policy

APMS-13158

[datadog-dd-trace-py-rkomorn]

Datadog Report

Branch report: close-apply-async-celery-spans
Commit report: 81b2eb9
Test service: dd-trace-py

✅ 0 Failed, 820 Passed, 376 Skipped, 21m 13.4s Total duration (15m 15.32s time saved)

[github-actions]

CODEOWNERS have been resolved as:

releasenotes/notes/fix-celery-apply-async-span-close-b7a8db188459f5b5.yaml  @DataDog/apm-python
ddtrace/contrib/internal/celery/app.py                                  @DataDog/apm-core-python @DataDog/apm-idm-python
ddtrace/contrib/internal/celery/signals.py                              @DataDog/apm-core-python @DataDog/apm-idm-python
tests/contrib/celery/test_integration.py                                @DataDog/apm-core-python @DataDog/apm-idm-python

[pr-commenter]

Benchmarks

Benchmark execution time: 2024-10-01 20:02:33

Comparing candidate commit 81b2eb9 in PR branch close-apply-async-celery-spans with baseline commit 7f6f3db in branch main.

Found 0 performance improvements and 0 performance regressions! Performance is the same for 371 metrics, 53 unstable metrics.