fix: Adapted instrumentations to the new get_tracer_tuple function

Signed-off-by: Cagri Yonca <cagri@ibm.com>

Author: CagriYonca

src/instana/instrumentation/aio_pika.py

@@ -15,7 +15,7 @@
 
     from instana.log import logger
     from instana.propagators.format import Format
-    from instana.util.traceutils import get_tracer_tuple, tracing_is_off
+    from instana.util.traceutils import get_tracer_tuple
     from instana.singletons import get_tracer
 
     if TYPE_CHECKING:
@@ -41,10 +41,10 @@ async def publish_with_instana(
         args: Tuple[object],
         kwargs: Dict[str, Any],
     ) -> Optional["ConfirmationFrameType"]:
-        if tracing_is_off():
+        tracer, parent_span, _ = get_tracer_tuple()
+        if not tracer:
             return await wrapped(*args, **kwargs)
 
-        tracer, parent_span, _ = get_tracer_tuple()
         parent_context = parent_span.get_span_context() if parent_span else None
 
         def _bind_args(

src/instana/instrumentation/aioamqp.py

@@ -8,7 +8,7 @@
     from opentelemetry.trace.status import StatusCode
 
     from instana.log import logger
-    from instana.util.traceutils import get_tracer_tuple, tracing_is_off
+    from instana.util.traceutils import get_tracer_tuple
 
     @wrapt.patch_function_wrapper("aioamqp.channel", "Channel.basic_publish")
     async def basic_publish_with_instana(
@@ -17,10 +17,10 @@ async def basic_publish_with_instana(
         argv: Tuple[object, Tuple[object, ...]],
         kwargs: Dict[str, Any],
     ) -> object:
-        if tracing_is_off():
+        tracer, parent_span, _ = get_tracer_tuple()
+        if not tracer:
             return await wrapped(*argv, **kwargs)
 
-        tracer, parent_span, _ = get_tracer_tuple()
         parent_context = parent_span.get_span_context() if parent_span else None
         with tracer.start_as_current_span(
             "aioamqp-publisher", span_context=parent_context
@@ -57,11 +57,11 @@ async def basic_consume_with_instana(
         argv: Tuple[object, Tuple[object, ...]],
         kwargs: Dict[str, Any],
     ) -> object:
-        if tracing_is_off():
+        tracer, parent_span, _ = get_tracer_tuple()
+        if not tracer:
             return await wrapped(*argv, **kwargs)
 
         callback = argv[0]
-        tracer, parent_span, _ = get_tracer_tuple()
         parent_context = parent_span.get_span_context() if parent_span else None
 
         @wrapt.decorator

src/instana/instrumentation/aiohttp/client.py

@@ -12,7 +12,7 @@
 from instana.propagators.format import Format
 from instana.singletons import agent
 from instana.util.secrets import strip_secrets_from_query
-from instana.util.traceutils import get_tracer_tuple, tracing_is_off, extract_custom_headers
+from instana.util.traceutils import get_tracer_tuple, extract_custom_headers
 
 try:
     import aiohttp
@@ -21,17 +21,16 @@
         from aiohttp.client import ClientSession
         from instana.span.span import InstanaSpan
 
-
     async def stan_request_start(
         session: "ClientSession", trace_config_ctx: SimpleNamespace, params
     ) -> Awaitable[None]:
         try:
+            tracer, parent_span, _ = get_tracer_tuple()
             # If we're not tracing, just return
-            if tracing_is_off():
+            if not tracer:
                 trace_config_ctx.span_context = None
                 return
 
-            tracer, parent_span, _ = get_tracer_tuple()
             parent_context = parent_span.get_span_context() if parent_span else None
 
             span = tracer.start_span("aiohttp-client", span_context=parent_context)

src/instana/instrumentation/asyncio.py

@@ -4,7 +4,7 @@
 
 import time
 from contextlib import contextmanager
-from typing import Any, Callable, Dict, Iterator, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, Tuple
 
 import wrapt
 from opentelemetry.trace import use_span
@@ -13,25 +13,26 @@
 from instana.configurator import config
 from instana.log import logger
 from instana.span.span import InstanaSpan
-from instana.util.traceutils import get_tracer_tuple, tracing_is_off
+from instana.util.traceutils import get_tracer_tuple
 
 try:
     import asyncio
 
+    if TYPE_CHECKING:
+        from instana.tracer import InstanaTracer
+
     @wrapt.patch_function_wrapper("asyncio", "ensure_future")
     def ensure_future_with_instana(
         wrapped: Callable[..., asyncio.ensure_future],
         instance: object,
         argv: Tuple[object, Tuple[object, ...]],
         kwargs: Dict[str, Any],
     ) -> object:
-        if (
-            not config["asyncio_task_context_propagation"]["enabled"]
-            or tracing_is_off()
-        ):
+        tracer, parent_span, _ = get_tracer_tuple()
+        if not config["asyncio_task_context_propagation"]["enabled"] or not tracer:
             return wrapped(*argv, **kwargs)
 
-        with _start_as_current_async_span() as span:
+        with _start_as_current_async_span(tracer, parent_span) as span:
             try:
                 span.set_status(StatusCode.OK)
                 return wrapped(*argv, **kwargs)
@@ -47,26 +48,26 @@ def create_task_with_instana(
             argv: Tuple[object, Tuple[object, ...]],
             kwargs: Dict[str, Any],
         ) -> object:
-            if (
-                not config["asyncio_task_context_propagation"]["enabled"]
-                or tracing_is_off()
-            ):
+            tracer, parent_span, _ = get_tracer_tuple()
+            if not config["asyncio_task_context_propagation"]["enabled"] or not tracer:
                 return wrapped(*argv, **kwargs)
 
-            with _start_as_current_async_span() as span:
+            with _start_as_current_async_span(tracer, parent_span) as span:
                 try:
                     span.set_status(StatusCode.OK)
                     return wrapped(*argv, **kwargs)
                 except Exception as exc:
                     logger.debug(f"asyncio create_task_with_instana error: {exc}")
 
     @contextmanager
-    def _start_as_current_async_span() -> Iterator[InstanaSpan]:
+    def _start_as_current_async_span(
+        tracer: "InstanaTracer",
+        parent_span: "InstanaSpan",
+    ) -> Iterator[InstanaSpan]:
         """
         Creates and yield a special InstanaSpan to only propagate the Asyncio
         context.
         """
-        tracer, parent_span, _ = get_tracer_tuple()
         parent_context = parent_span.get_span_context() if parent_span else None
 
         _time = time.time_ns()

src/instana/instrumentation/aws/boto3.py

@@ -14,6 +14,7 @@
         from botocore.client import BaseClient
 
         from instana.span.span import InstanaSpan
+        from instana.tracer import InstanaTracer
 
     import json
 
@@ -22,14 +23,16 @@
     from instana.log import logger
     from instana.propagators.format import Format
     from instana.singletons import get_tracer
-    from instana.span.span import get_current_span
     from instana.util.traceutils import (
         extract_custom_headers,
         get_tracer_tuple,
-        tracing_is_off,
     )
 
-    def lambda_inject_context(payload: Dict[str, Any], span: "InstanaSpan") -> None:
+    def lambda_inject_context(
+        tracer: "InstanaTracer",
+        payload: Dict[str, Any],
+        span: "InstanaSpan",
+    ) -> None:
         """
         When boto3 lambda client 'Invoke' is called, we want to inject the tracing context.
         boto3/botocore has specific requirements:
@@ -54,9 +57,9 @@ def emit_add_auth_with_instana(
         args: Tuple[object],
         kwargs: Dict[str, Any],
     ) -> Callable[..., None]:
-        current_span = get_current_span()
-        if not tracing_is_off() and current_span and current_span.is_recording():
-            extract_custom_headers(current_span, args[0].headers)
+        _, parent_span, _ = get_tracer_tuple()
+        if parent_span:
+            extract_custom_headers(parent_span, args[0].headers)
         return wrapped(*args, **kwargs)
 
     @wrapt.patch_function_wrapper("botocore.client", "BaseClient._make_api_call")
@@ -66,12 +69,11 @@ def make_api_call_with_instana(
         args: Sequence[Dict[str, Any]],
         kwargs: Dict[str, Any],
     ) -> Dict[str, Any]:
+        tracer, parent_span, _ = get_tracer_tuple()
         # If we're not tracing, just return
-        if tracing_is_off():
+        if not tracer:
             return wrapped(*args, **kwargs)
 
-        tracer, parent_span, _ = get_tracer_tuple()
-
         parent_context = parent_span.get_span_context() if parent_span else None
 
         if instance.meta.service_model.service_name == "dynamodb":
@@ -101,7 +103,7 @@ def make_api_call_with_instana(
 
                 # Inject context when invoking lambdas
                 if "lambda" in instance._endpoint.host and operation == "Invoke":
-                    lambda_inject_context(payload, span)
+                    lambda_inject_context(tracer, payload, span)
 
                 try:
                     result = wrapped(*args, **kwargs)

src/instana/instrumentation/aws/s3.py

@@ -15,7 +15,6 @@
     from instana.singletons import get_tracer
     from instana.util.traceutils import (
         get_tracer_tuple,
-        tracing_is_off,
     )
 
     operations = {
@@ -48,12 +47,11 @@ def collect_s3_injected_attributes(
         args: Sequence[object],
         kwargs: Dict[str, Any],
     ) -> Callable[..., object]:
+        tracer, parent_span, _ = get_tracer_tuple()
         # If we're not tracing, just return
-        if tracing_is_off():
+        if not tracer:
             return wrapped(*args, **kwargs)
 
-        tracer, parent_span, _ = get_tracer_tuple()
-
         parent_context = parent_span.get_span_context() if parent_span else None
 
         with tracer.start_as_current_span("s3", span_context=parent_context) as span:

src/instana/instrumentation/cassandra.py

@@ -14,7 +14,7 @@
     import wrapt
 
     from instana.log import logger
-    from instana.util.traceutils import get_tracer_tuple, tracing_is_off
+    from instana.util.traceutils import get_tracer_tuple
 
     if TYPE_CHECKING:
         from cassandra.cluster import ResponseFuture, Session
@@ -73,11 +73,11 @@ def request_init_with_instana(
         fn: "ResponseFuture",
     ) -> None:
         tracer, parent_span, _ = get_tracer_tuple()
-        parent_context = parent_span.get_span_context() if parent_span else None
-
-        if tracing_is_off():
+        if not tracer:
             return
 
+        parent_context = parent_span.get_span_context() if parent_span else None
+
         attributes = {}
         if isinstance(fn.query, cassandra.query.SimpleStatement):
             attributes["cassandra.query"] = fn.query.query_string

src/instana/instrumentation/celery.py

@@ -145,40 +145,42 @@ def before_task_publish(
     ) -> None:
         try:
             tracer, parent_span, _ = get_tracer_tuple()
+            if not tracer:
+                return
+
             parent_context = parent_span.get_span_context() if parent_span else None
 
-            if tracer:
-                body = kwargs["body"]
-                headers = kwargs["headers"]
-                task_name = kwargs["sender"]
-                task = registry.tasks.get(task_name)
-                task_id = _get_task_id(headers, body)
-
-                span = tracer.start_span("celery-client", span_context=parent_context)
-                span.set_attribute("task", task_name)
-                span.set_attribute("task_id", task_id)
-                add_broker_attributes(span, task.app.conf["broker_url"])
-
-                # Context propagation
-                context_headers = {}
-                tracer.inject(
-                    span.context,
-                    Format.HTTP_HEADERS,
-                    context_headers,
-                    disable_w3c_trace_context=True,
-                )
+            body = kwargs["body"]
+            headers = kwargs["headers"]
+            task_name = kwargs["sender"]
+            task = registry.tasks.get(task_name)
+            task_id = _get_task_id(headers, body)
+
+            span = tracer.start_span("celery-client", span_context=parent_context)
+            span.set_attribute("task", task_name)
+            span.set_attribute("task_id", task_id)
+            add_broker_attributes(span, task.app.conf["broker_url"])
+
+            # Context propagation
+            context_headers = {}
+            tracer.inject(
+                span.context,
+                Format.HTTP_HEADERS,
+                context_headers,
+                disable_w3c_trace_context=True,
+            )
+
+            # Fix for broken header propagation
+            # https://github.com/celery/celery/issues/4875
+            task_headers = kwargs.get("headers") or {}
+            task_headers.setdefault("headers", {})
+            task_headers["headers"].update(context_headers)
+            kwargs["headers"] = task_headers
 
-                # Fix for broken header propagation
-                # https://github.com/celery/celery/issues/4875
-                task_headers = kwargs.get("headers") or {}
-                task_headers.setdefault("headers", {})
-                task_headers["headers"].update(context_headers)
-                kwargs["headers"] = task_headers
-
-                ctx = trace.set_span_in_context(span)
-                token = context.attach(ctx)
-                client_token["token"] = token
-                client_span.set(span)
+            ctx = trace.set_span_in_context(span)
+            token = context.attach(ctx)
+            client_token["token"] = token
+            client_span.set(span)
         except Exception:
             logger.debug("celery-client before_task_publish: ", exc_info=True)
 

src/instana/instrumentation/couchbase.py

@@ -25,7 +25,7 @@
     import wrapt
 
     from instana.span.span import InstanaSpan
-    from instana.util.traceutils import get_tracer_tuple, tracing_is_off
+    from instana.util.traceutils import get_tracer_tuple
 
     # List of operations to instrument
     # incr, incr_multi, decr, decr_multi, retrieve_in are wrappers around operations above
@@ -94,12 +94,12 @@ def wrapper(
             kwargs: Dict[str, Any],
         ) -> object:
             tracer, parent_span, _ = get_tracer_tuple()
-            parent_context = parent_span.get_span_context() if parent_span else None
-
             # If we're not tracing, just return
-            if tracing_is_off():
+            if not tracer:
                 return wrapped(*args, **kwargs)
 
+            parent_context = parent_span.get_span_context() if parent_span else None
+
             with tracer.start_as_current_span(
                 "couchbase", span_context=parent_context
             ) as span:
@@ -120,12 +120,12 @@ def query_with_instana(
         kwargs: Dict[str, Any],
     ) -> object:
         tracer, parent_span, _ = get_tracer_tuple()
-        parent_context = parent_span.get_span_context() if parent_span else None
-
         # If we're not tracing, just return
-        if tracing_is_off():
+        if not tracer:
             return wrapped(*args, **kwargs)
 
+        parent_context = parent_span.get_span_context() if parent_span else None
+
         with tracer.start_as_current_span(
             "couchbase", span_context=parent_context
         ) as span: