Provide an alternative OpenTelemetry implementation for traces that follows standard otel practices

airflow/cli/cli_config.py

@@ -551,6 +551,9 @@ def string_lower_type(val):
     action="store_true",
 )
 ARG_RAW = Arg(("-r", "--raw"), argparse.SUPPRESS, "store_true")
+ARG_CARRIER = Arg(
+    ("-c", "--carrier"), help="Context Carrier, containing the injected context for the task span", nargs="?"
+)
 ARG_IGNORE_ALL_DEPENDENCIES = Arg(
     ("-A", "--ignore-all-dependencies"),
     help="Ignores all non-critical dependencies, including ignore_ti_state and ignore_task_deps",
@@ -1311,6 +1314,7 @@ class GroupCommand(NamedTuple):
             ARG_CFG_PATH,
             ARG_LOCAL,
             ARG_RAW,
+            ARG_CARRIER,
             ARG_IGNORE_ALL_DEPENDENCIES,
             ARG_IGNORE_DEPENDENCIES,
             ARG_DEPENDS_ON_PAST,

airflow/cli/commands/remote_commands/task_command.py

@@ -66,6 +66,7 @@
 from airflow.utils.net import get_hostname
 from airflow.utils.providers_configuration_loader import providers_configuration_loaded
 from airflow.utils.session import NEW_SESSION, create_session, provide_session
+from airflow.utils.span_status import SpanStatus
 from airflow.utils.state import DagRunState
 from airflow.utils.task_instance_session import set_current_task_instance_session
 from airflow.utils.types import DagRunTriggeredByType
@@ -462,6 +463,13 @@ def task_run(args, dag: DAG | None = None) -> TaskReturnCode | None:
 
     log.info("Running %s on host %s", ti, hostname)
 
+    if args.carrier is not None:
+        log.info("Found args.carrier: %s. Setting the value on the ti instance.", args.carrier)
+        # The arg value is a dict string, and it needs to be converted back to a dict.
+        carrier_dict = json.loads(args.carrier)
+        ti.set_context_carrier(context_carrier=carrier_dict, with_commit=True)
+        ti.set_span_status(status=SpanStatus.ACTIVE, with_commit=True)
+
     task_return_code = None
     try:
         if args.interactive:

airflow/executors/base_executor.py

@@ -18,6 +18,7 @@
 
 from __future__ import annotations
 
+import json
 import logging
 import sys
 from collections import defaultdict, deque
@@ -39,6 +40,7 @@
 from airflow.traces.utils import gen_span_id_from_ti_key, gen_trace_id
 from airflow.utils.log.logging_mixin import LoggingMixin
 from airflow.utils.state import TaskInstanceState
+from airflow.utils.thread_safe_dict import ThreadSafeDict
 
 PARALLELISM: int = conf.getint("core", "PARALLELISM")
 
@@ -115,6 +117,8 @@ class BaseExecutor(LoggingMixin):
     :param parallelism: how many jobs should run at one time. Set to ``0`` for infinity.
     """
 
+    active_spans = ThreadSafeDict()
+
     supports_ad_hoc_ti_run: bool = False
     supports_sentry: bool = False
 
@@ -150,6 +154,10 @@ def __init__(self, parallelism: int = PARALLELISM):
     def __repr__(self):
         return f"{self.__class__.__name__}(parallelism={self.parallelism})"
 
+    @classmethod
+    def set_active_spans(cls, active_spans: ThreadSafeDict):
+        cls.active_spans = active_spans
+
     def start(self):  # pragma: no cover
         """Executors may need to get things started."""
 
@@ -337,6 +345,35 @@ def trigger_tasks(self, open_slots: int) -> None:
         for _ in range(min((open_slots, len(self.queued_tasks)))):
             key, (command, _, queue, ti) = sorted_queue.pop(0)
 
+            # If it's None, then the span for the current TaskInstanceKey hasn't been started.
+            if self.active_spans is not None and self.active_spans.get(key) is None:
+                from airflow.models.taskinstance import SimpleTaskInstance
+
+                if isinstance(ti, SimpleTaskInstance):
+                    parent_context = Trace.extract(ti.parent_context_carrier)
+                else:
+                    parent_context = Trace.extract(ti.dag_run.context_carrier)
+                # Start a new span using the context from the parent.
+                # Attributes will be set once the task has finished so that all
+                # values will be available (end_time, duration, etc.).
+                span = Trace.start_child_span(
+                    span_name=f"{ti.task_id}",
+                    parent_context=parent_context,
+                    component="task",
+                    start_time=ti.queued_dttm,
+                    start_as_current=False,
+                )
+                self.active_spans.set(key, span)
+                # Inject the current context into the carrier.
+                carrier = Trace.inject()
+                # The carrier needs to be set on the ti, but it can't happen here because db calls are expensive.
+                # By the time the db update has finished, another heartbeat will have started
+                # and the tasks will have been triggered again.
+                # So set the carrier as an argument to the command.
+                # The command execution will set it on the ti, and it will be propagated to the task itself.
+                command.append("--carrier")
+                command.append(json.dumps(carrier))
+
             # If a task makes it here but is still understood by the executor
             # to be running, it generally means that the task has been killed
             # externally and not yet been marked as failed.