Fix incorrect propagation of task.connection (#86121)

Co-authored-by: Matt Clay <matt@mystile.com>

changelogs/fragments/task_connection_propagation_redux.yml

@@ -0,0 +1,3 @@
+bugfixes:
+  - callbacks - The value of ``TaskResult.task.connection`` properly reflects the loaded connection name used.
+    Previously, incorrect values were reported in some cases.

lib/ansible/executor/task_executor.py

@@ -31,6 +31,7 @@ from ansible.module_utils.common.text.converters import to_text, to_native
 from ansible.module_utils.connection import write_to_stream
 from ansible.playbook.task import Task
 from ansible.plugins import get_plugin_class
+from ansible.plugins.connection import ConnectionBase
 from ansible.plugins.loader import become_loader, cliconf_loader, connection_loader, httpapi_loader, netconf_loader, terminal_loader
 from ansible._internal._templating._jinja_plugins import _invoke_lookup, _DirectCall
 from ansible._internal._templating._engine import TemplateEngine
@@ -318,6 +319,9 @@ class TaskExecutor:
             (self._task, tmp_task) = (tmp_task, self._task)
             (self._play_context, tmp_play_context) = (tmp_play_context, self._play_context)
 
+            # loop item task copies got connection field updated already; this updates the original outer task
+            self._update_task_connection()
+
             # now update the result with the item info, and append the result
             # to the list of results
             res[loop_var] = item
@@ -341,10 +345,6 @@ class TaskExecutor:
                     'msg': 'Failed to template loop_control.label: %s' % to_text(e)
                 })
 
-            # if plugin is loaded, get resolved name, otherwise leave original task connection
-            if self._connection and not isinstance(self._connection, str):
-                task_fields['connection'] = getattr(self._connection, 'ansible_name')
-
             tr = _RawTaskResult(
                 host=self._host,
                 task=self._task,
@@ -448,6 +448,18 @@ class TaskExecutor:
 
         return result
 
+    def _update_task_connection(self, task: Task | None = None) -> None:
+        """If a connection plugin is loaded, ensure the resolved name is propagated back to the controller as the task's connection."""
+
+        if not task:
+            task = self._task
+
+        # FUTURE: What value should be reported when there is no connection?
+        #         This is currently not possible, but it should be.
+
+        if isinstance(self._connection, ConnectionBase):
+            task.connection = self._connection.ansible_name
+
     def _execute_internal(self, templar: TemplateEngine, variables: dict[str, t.Any]) -> dict[str, t.Any]:
         """
         The primary workhorse of the executor system, this runs the task
@@ -603,6 +615,9 @@ class TaskExecutor:
         # get handler
         self._handler, _module_context = self._get_action_handler_with_module_context(templar=templar)
 
+        # self._connection should have its final value for this task/loop-item by this point; record on the task object
+        self._update_task_connection()
+
         retries = 1  # includes the default actual run + retries set by user/default
         if self._task.retries is not None:
             retries += max(0, self._task.retries)
@@ -866,6 +881,9 @@ class TaskExecutor:
             environment=self._task.environment,
         ))
 
+        # ensure that the synthetic async task has the resolved connection recorded on it
+        self._update_task_connection(async_task)
+
         # FIXME: this is no longer the case, normal takes care of all, see if this can just be generalized
         # Because this is an async task, the action handler is async. However,
         # we need the 'normal' action handler for the status check, so get it

test/integration/targets/callback_results/callback_plugins/track_connections.py

@@ -11,7 +11,10 @@ DOCUMENTATION = """
     type: aggregate
 """
 
+import functools
+import inspect
 import json
+
 from collections import defaultdict
 
 from ansible.plugins.callback import CallbackBase
@@ -22,20 +25,46 @@ class CallbackModule(CallbackBase):
     CALLBACK_VERSION = 2.0
     CALLBACK_TYPE = 'aggregate'
     CALLBACK_NAME = 'track_connections'
+    CALLBACK_NEEDS_ENABLED = True
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self._conntrack = defaultdict(lambda : defaultdict(int))
 
-    def _track(self, result: CallbackTaskResult, *args, **kwargs):
+        self._conntrack = defaultdict(lambda : defaultdict(list))
+
+        # dynamically implement all v2 callback methods that accept `result`
+        for name, sig in ((cb, inspect.signature(getattr(self, cb))) for cb in dir(self) if cb.startswith('v2_')):
+            if 'result' in sig.parameters:
+                setattr(self, name, functools.partial(self._track, event_name=name))
+
+    def _track(self, result: CallbackTaskResult, *_args, event_name: str, **_kwargs):
         host = result.host.get_name()
         task = result.task
 
-        self._conntrack[host][task.connection] += 1
-
-    v2_runner_on_ok = v2_runner_on_failed = _track
-    v2_runner_on_async_poll = v2_runner_on_async_ok = v2_runner_on_async_failed = _track
-    v2_runner_item_on_ok = v2_runner_item_on_failed = _track
+        self._conntrack[host][task.connection].append(f'{event_name}: {task.name}')
 
     def v2_playbook_on_stats(self, stats):
-        self._display.display(json.dumps(self._conntrack, indent=4))
+        expected = {
+            "testhost": {
+                "ansible.builtin.local": [
+                    "v2_runner_on_ok: execute a successful non-loop task with the local connection",
+                    "v2_runner_on_failed: execute a failing non-loop task with the local connection",
+                    "v2_runner_item_on_ok: execute a successful looped task with the local connection",
+                    "v2_runner_on_ok: execute a successful looped task with the local connection",
+                    "v2_runner_item_on_failed: execute a failing looped task with the local connection",
+                    "v2_runner_on_failed: execute a failing looped task with the local connection",
+                    "v2_runner_on_async_ok: execute a successful async task with the local connection",
+                    "v2_runner_on_ok: execute a successful async task with the local connection",
+                    "v2_runner_on_async_failed: execute a failing async task with the local connection",
+                    "v2_runner_on_failed: execute a failing async task with the local connection"
+                ],
+            }
+        }
+
+        if self._conntrack == expected:
+            self._display.display('FOUND EXPECTED EVENTS')
+            return
+
+        # pragma: nocover
+        self._display.display(f'ACTUAL\n{json.dumps(self._conntrack, indent=4)}')
+        self._display.display(f'EXPECTED\n{json.dumps(expected, indent=4)}')

test/integration/targets/callback_results/connection_name.yml

@@ -0,0 +1,37 @@
+- hosts: testhost
+  gather_facts: no
+  vars:
+    ansible_connection: '{{ "ansible.legacy.local" }}'  # use a templated value with the non-canonical name of the connection; this should resolve to `ansible.builtin.local`
+  name: validate that task.connection is always overwritten with the templated and resolved name of the connection
+  tasks:
+    - name: execute a successful non-loop task with the local connection
+      raw: echo hi
+
+    - name: execute a failing non-loop task with the local connection
+      raw: exit 1
+      ignore_errors: true
+
+    - name: execute a successful looped task with the local connection
+      raw: echo hi {{ item }}
+      loop: [1]
+
+    - name: execute a failing looped task with the local connection
+      raw: echo hi {{ item }}; exit 1
+      ignore_errors: true
+      loop: [1]
+
+    - name: execute a successful async task with the local connection
+      command: echo hi
+      async: 5
+      poll: 1
+
+    - name: execute a failing async task with the local connection
+      command: exit 1
+      async: 5
+      poll: 1
+
+    - name: execute a looped async task with the local connection
+      command: echo hi {{ item }}
+      async: 5
+      poll: 1
+      loop: [1]

test/integration/targets/callback_results/runme.sh

@@ -13,14 +13,13 @@ ansible-playbook "$@" -i ../../inventory task_name.yml | tee "${OUTFILE}"
 echo "Grepping for ${EXPECTED_REGEX} in stdout."
 grep -e "${EXPECTED_REGEX}" "${OUTFILE}"
 
-# test connection tracking
-EXPECTED_CONNECTION='{"testhost":{"ssh":4}}'
-OUTPUT_TAIL=$(tail -n5 ${OUTFILE} | tr -d '[:space:]')
-echo "Checking for connection string ${OUTPUT_TAIL} in stdout."
-[ "${EXPECTED_CONNECTION}" == "${OUTPUT_TAIL}" ]
-echo $?
-
 # check variables are interpolated in 'started'
 UNTEMPLATED_STARTED="^.*\[started .*{{.*}}.*$"
 echo "Checking we dont have untemplated started in stdout."
-grep -e "${UNTEMPLATED_STARTED}" "${OUTFILE}" || exit 0
+if grep -e "${UNTEMPLATED_STARTED}" "${OUTFILE}"; then
+  exit 1
+fi
+
+# test connection tracking
+ANSIBLE_CALLBACKS_ENABLED=track_connections ansible-playbook "$@" -i ../../inventory connection_name.yml | tee "${OUTFILE}"
+grep "FOUND EXPECTED EVENTS" "${OUTFILE}"