From ce5aefd3d8953fc528575d60e736b0e48c5ee7ee Mon Sep 17 00:00:00 2001
From: Jeff Bradberry <jeff.bradberry@gmail.com>
Date: Mon, 20 Dec 2021 15:23:32 -0500
Subject: [PATCH] Capture hop nodes and links in the automatic discovery
 machinery

Also, make sure that the control service is turned on in the dev
environment's hop node, so that it shows up in the Advertisements
list.
---
 awx/main/tasks.py                             | 37 +++++++++++++------
 .../sources/templates/receptor-hop.conf.j2    |  3 ++
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/awx/main/tasks.py b/awx/main/tasks.py
index 477cb1a47a..d936d3e6ae 100644
--- a/awx/main/tasks.py
+++ b/awx/main/tasks.py
@@ -67,6 +67,7 @@ from awx.main.models import (
     TowerScheduleState,
     Instance,
     InstanceGroup,
+    InstanceLink,
     UnifiedJob,
     Notification,
     Inventory,
@@ -495,28 +496,29 @@ def execution_node_health_check(node):
 
 def inspect_execution_nodes(instance_list):
     with advisory_lock('inspect_execution_nodes_lock', wait=False):
-        node_lookup = {}
-        for inst in instance_list:
-            if inst.node_type == 'execution':
-                node_lookup[inst.hostname] = inst
+        node_lookup = {inst.hostname: inst for inst in instance_list}
 
         ctl = get_receptor_ctl()
-        connections = ctl.simple_command('status')['Advertisements']
+        mesh_status = ctl.simple_command('status')
+
         nowtime = now()
-        for ad in connections:
+        workers = mesh_status['Advertisements']
+        for ad in workers:
             hostname = ad['NodeID']
-            commands = ad.get('WorkCommands') or []
-            worktypes = []
-            for c in commands:
-                worktypes.append(c["WorkType"])
-            if 'ansible-runner' not in worktypes:
+            if ad.get('WorkCommands') is None:
+                node_type = 'hop'
+            elif any(cmd['WorkType'] == 'ansible-runner' for cmd in ad['WorkCommands'] or []):
+                node_type = 'execution'
+            else:
                 continue
+
             changed = False
             if hostname in node_lookup:
                 instance = node_lookup[hostname]
             elif settings.MESH_AUTODISCOVERY_ENABLED:
                 defaults = dict(enabled=False)
-                (changed, instance) = Instance.objects.register(hostname=hostname, node_type='execution', defaults=defaults)
+                (changed, instance) = Instance.objects.register(hostname=hostname, node_type=node_type, defaults=defaults)
+                node_lookup[hostname] = instance
                 logger.warn(f"Registered execution node '{hostname}' (marked disabled by default)")
             else:
                 logger.warn(f"Unrecognized node on mesh advertising ansible-runner work type: {hostname}")
@@ -546,6 +548,17 @@ def inspect_execution_nodes(instance_list):
                     logger.debug(f'Restarting health check for execution node {hostname} with known errors.')
                     execution_node_health_check.apply_async([hostname])
 
+        links = {tuple(sorted((node, peer))) for node, peers in mesh_status['KnownConnectionCosts'].items() for peer in peers}
+        for a, b in links:
+            if a not in node_lookup:
+                logger.warn(f"Cannot link {a} to {b}: {a} not registered.")
+                continue
+            if b not in node_lookup:
+                logger.warn(f"Cannot link {a} to {b}: {b} not registered.")
+                continue
+            a_obj, b_obj = node_lookup[a], node_lookup[b]
+            InstanceLink.objects.get_or_create(source=a_obj, target=b_obj)
+
 
 @task(queue=get_local_queuename)
 def cluster_node_heartbeat():
diff --git a/tools/docker-compose/ansible/roles/sources/templates/receptor-hop.conf.j2 b/tools/docker-compose/ansible/roles/sources/templates/receptor-hop.conf.j2
index 4f055fa1f3..f0bd884f04 100644
--- a/tools/docker-compose/ansible/roles/sources/templates/receptor-hop.conf.j2
+++ b/tools/docker-compose/ansible/roles/sources/templates/receptor-hop.conf.j2
@@ -10,3 +10,6 @@
 
 - tcp-listener:
     port: 5555
+
+- control-service:
+    service: control