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

[wantsui]

Credits go to @tabgok for going through this with me in great detail!

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 in two known places, which can lead to broken context propagation (and unexpected traces).

Issue 1: Closing signals don't get called.

• 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.

Issue 2: Even if after_task_publish is called, the integration previously only closed the span if the task or task id was present. However, it is possible that there is a lingering span kept open but doesn't return a task_id and instead just return nothing. This is bad because that means the remnants of that celery.apply span is never closed either.

To address both, 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.

Meanwhile, the fix to celery/utils.py is to address the fact that some times, the returning response returns an empty body, in which case, we should not try to run body.get('id') and just return None.

Open Item

1. At the moment, the span is closed but any underlying errors won't be surfaced on the span. It is possible to do this in an except clause, but I'm not sure if I should.

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

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

[datadog-dd-trace-py-rkomorn]

Datadog Report

Branch report: close-apply-async-celery-spans
Commit report: 5b96ae7
Test service: dd-trace-py

✅ 0 Failed, 673 Passed, 523 Skipped, 22m 41.31s Total duration (14m 7.95s 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
ddtrace/contrib/internal/celery/utils.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-09-27 14:33:35

Comparing candidate commit 5b96ae7 in PR branch close-apply-async-celery-spans with baseline commit 0bd5ea0 in branch main.

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

[tabgok]

nit: usually good to have one PR per change you're making. This is a one-liner so arguable, but it makes it easier to rollback/make sure you're not missing something in the future

[tabgok]

nit: would add a little bit about why they weren't called (exception generated in handling)