def _trace_failure(*args, **kwargs):
        task = utils.retrieve_task_from_sender(kwargs)
        task_id = utils.retrieve_task_id(kwargs)

        if task is None or task_id is None:
            return

        # retrieve and pass exception info to activation
        span, _ = utils.retrieve_span(task, task_id)
        if span is None or not span.is_recording():
            return

        status_kwargs = {"status_code": StatusCode.ERROR}

        ex = kwargs.get("einfo")

        if (
            hasattr(task, "throws")
            and ex is not None
            and isinstance(ex.exception, task.throws)
        ):
            return

        if ex is not None:
            status_kwargs["description"] = str(ex)
        span.set_status(Status(**status_kwargs))
    def test_span_propagation(self):
        # ensure spans getter and setter works properly
        @self.app.task
        def fn_task():
            return 42

        # propagate and retrieve a Span
        task_id = "7c6731af-9533-40c3-83a9-25b58f0d837f"
        span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext))
        utils.attach_span(fn_task, task_id, span)
        span_after = utils.retrieve_span(fn_task, task_id)
        self.assertIs(span, span_after)
    def test_span_delete(self):
        # ensure the helper removes properly a propagated Span
        @self.app.task
        def fn_task():
            return 42

        # propagate a Span
        task_id = "7c6731af-9533-40c3-83a9-25b58f0d837f"
        span = trace._Span("name", mock.Mock(spec=trace_api.SpanContext))
        utils.attach_span(fn_task, task_id, span)
        # delete the Span
        utils.detach_span(fn_task, task_id)
        self.assertEqual(utils.retrieve_span(fn_task, task_id), (None, None))
    def test_span_delete_empty(self):
        # ensure detach_span doesn't raise an exception if span is not present
        @self.app.task
        def fn_task():
            return 42

        # delete the Span
        task_id = "7c6731af-9533-40c3-83a9-25b58f0d837f"
        try:
            utils.detach_span(fn_task, task_id)
            self.assertEqual(utils.retrieve_span(fn_task, task_id),
                             (None, None))
        except Exception as ex:  # pylint: disable=broad-except
            self.fail("Exception was raised: %s" % ex)
    def _trace_after_publish(*args, **kwargs):
        task = utils.retrieve_task_from_sender(kwargs)
        task_id = utils.retrieve_task_id_from_message(kwargs)

        if task is None or task_id is None:
            return

        # retrieve and finish the Span
        _, activation = utils.retrieve_span(task, task_id, is_publish=True)
        if activation is None:
            logger.warning("no existing span found for task_id=%s", task_id)
            return

        activation.__exit__(None, None, None)  # pylint: disable=E1101
        utils.detach_span(task, task_id, is_publish=True)
    def _trace_retry(*args, **kwargs):
        task = utils.retrieve_task_from_sender(kwargs)
        task_id = utils.retrieve_task_id_from_request(kwargs)
        reason = utils.retrieve_reason(kwargs)

        if task is None or task_id is None or reason is None:
            return

        span, _ = utils.retrieve_span(task, task_id)
        if span is None or not span.is_recording():
            return

        # Add retry reason metadata to span
        # Use `str(reason)` instead of `reason.message` in case we get
        # something that isn't an `Exception`
        span.set_attribute(_TASK_RETRY_REASON_KEY, str(reason))
    def _trace_postrun(*args, **kwargs):
        task = utils.retrieve_task(kwargs)
        task_id = utils.retrieve_task_id(kwargs)

        if task is None or task_id is None:
            return

        logger.debug("postrun signal task_id=%s", task_id)

        # retrieve and finish the Span
        span, activation = utils.retrieve_span(task, task_id)
        if span is None:
            logger.warning("no existing span found for task_id=%s", task_id)
            return

        # request context tags
        span.set_attribute(_TASK_TAG_KEY, _TASK_RUN)
        utils.set_attributes_from_context(span, kwargs)
        utils.set_attributes_from_context(span, task.request)
        span.set_attribute(_TASK_NAME_KEY, task.name)

        activation.__exit__(None, None, None)
        utils.detach_span(task, task_id)