Remove pin from celery.app.task.Task and update patch and tests accor…

…dingly.

ddtrace/contrib/internal/celery/app.py

@@ -44,7 +44,6 @@ def patch_app(app, pin=None):
 
     # Patch apply_async
     trace_utils.wrap("celery.app.task", "Task.apply_async", _traced_apply_async_function(config.celery, "apply_async"))
-    pin.onto(celery.app.task.Task)
 
     # connect to the Signal framework
     signals.task_prerun.connect(trace_prerun, weak=False)
@@ -70,6 +69,7 @@ def unpatch_app(app):
 
     trace_utils.unwrap(celery.beat.Scheduler, "apply_entry")
     trace_utils.unwrap(celery.beat.Scheduler, "tick")
+    trace_utils.unwrap(celery.app.task.Task, "apply_async")
 
     signals.task_prerun.disconnect(trace_prerun)
     signals.task_postrun.disconnect(trace_postrun)
@@ -105,6 +105,7 @@ def _traced_beat_inner(func, instance, args, kwargs):
 
 def _traced_apply_async_function(integration_config, fn_name, resource_fn=None):
     def _traced_apply_async_inner(func, instance, args, kwargs):
+
         with core.context_with_data("task_context"):
             try:
                 return func(*args, **kwargs)
@@ -113,4 +114,8 @@ def _traced_apply_async_inner(func, instance, args, kwargs):
                 if task_span:
                     task_span.finish()
 
+                prerun_span = core.get_item("prerun_span")
+                if prerun_span:
+                    prerun_span.finish()
+
     return _traced_apply_async_inner

ddtrace/contrib/internal/celery/signals.py

@@ -49,6 +49,9 @@ def trace_prerun(*args, **kwargs):
     service = config.celery["worker_service_name"]
     span = pin.tracer.trace(c.WORKER_ROOT_SPAN, service=service, resource=task.name, span_type=SpanTypes.WORKER)
 
+    # Store an item called "prerun span" in case task_postrun doesn't get called
+    core.set_item("prerun_span", span)
+
     # set span.kind to the type of request being performed
     span.set_tag_str(SPAN_KIND, SpanKind.CONSUMER)
 
@@ -112,7 +115,7 @@ def trace_before_publish(*args, **kwargs):
     service = config.celery["producer_service_name"]
     span = pin.tracer.trace(c.PRODUCER_ROOT_SPAN, service=service, resource=task_name)
 
-    # Store an item called "task span" in case AFTER_TASK_PUBLISH doesn't get called
+    # Store an item called "task span" in case after_task_publish doesn't get called
     core.set_item("task_span", span)
 
     span.set_tag_str(COMPONENT, config.celery.integration_name)

tests/contrib/celery/base.py

@@ -85,7 +85,6 @@ def setUp(self):
         # override pins to use our Dummy Tracer
         Pin.override(self.app, tracer=self.tracer)
         Pin.override(celery.beat.Scheduler, tracer=self.tracer)
-        Pin.override(celery.app.task.Task, tracer=self.tracer)
 
     def tearDown(self):
         self.app = None