From cbcf2bd84a174df9fb7fe17a9fcbf34efee37e9f Mon Sep 17 00:00:00 2001
From: occheung <dc@m-labs.hk>
Date: Mon, 26 Aug 2024 17:41:27 +0800
Subject: [PATCH] frontend: pass drtio destination during communication

---
 artiq/coredevice/comm_mgmt.py    | 4 +++-
 artiq/frontend/aqctl_corelog.py  | 8 ++++++--
 artiq/frontend/artiq_coremgmt.py | 7 ++++++-
 3 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/artiq/coredevice/comm_mgmt.py b/artiq/coredevice/comm_mgmt.py
index 870e2759d..1b218afec 100644
--- a/artiq/coredevice/comm_mgmt.py
+++ b/artiq/coredevice/comm_mgmt.py
@@ -46,15 +46,17 @@ class LogLevel(Enum):
 
 
 class CommMgmt:
-    def __init__(self, host, port=1380):
+    def __init__(self, host, port=1380, drtio_dest=0):
         self.host = host
         self.port = port
+        self.drtio_dest = drtio_dest
 
     def open(self):
         if hasattr(self, "socket"):
             return
         self.socket = create_connection(self.host, self.port)
         self.socket.sendall(b"ARTIQ management\n")
+        self._write_int8(self.drtio_dest)
         endian = self._read(1)
         if endian == b"e":
             self.endian = "<"
diff --git a/artiq/frontend/aqctl_corelog.py b/artiq/frontend/aqctl_corelog.py
index c5f5a8eee..d2dd33a4a 100755
--- a/artiq/frontend/aqctl_corelog.py
+++ b/artiq/frontend/aqctl_corelog.py
@@ -25,6 +25,9 @@ def get_argparser():
                         help="Simulation - does not connect to device")
     parser.add_argument("core_addr", metavar="CORE_ADDR",
                         help="hostname or IP address of the core device")
+    parser.add_argument("-s", "--satellite", default=0,
+                        metavar="DRTIO_ID", type=int,
+                        help="the logged DRTIO destination")
     return parser
 
 
@@ -39,7 +42,7 @@ async def get_logs_sim(host):
         log_with_name("firmware.simulation", logging.INFO, "hello " + host)
 
 
-async def get_logs(host):
+async def get_logs(host, drtio_dest):
     try:
         reader, writer = await async_open_connection(
             host,
@@ -49,6 +52,7 @@ async def get_logs(host):
             max_fails=3,
         )
         writer.write(b"ARTIQ management\n")
+        writer.write(drtio_dest.to_bytes(1))
         endian = await reader.readexactly(1)
         if endian == b"e":
             endian = "<"
@@ -96,7 +100,7 @@ def main():
         signal_handler.setup()
         try:
             get_logs_task = asyncio.ensure_future(
-                get_logs_sim(args.core_addr) if args.simulation else get_logs(args.core_addr),
+                get_logs_sim(args.core_addr) if args.simulation else get_logs(args.core_addr, args.satellite),
                 loop=loop)
             try:
                 server = Server({"corelog": PingTarget()}, None, True)
diff --git a/artiq/frontend/artiq_coremgmt.py b/artiq/frontend/artiq_coremgmt.py
index 789c26bf7..2fbf450c0 100755
--- a/artiq/frontend/artiq_coremgmt.py
+++ b/artiq/frontend/artiq_coremgmt.py
@@ -95,6 +95,11 @@ def get_argparser():
     p_allocator = subparsers.add_parser("allocator",
                                         help="show heap layout")
 
+    # manage target
+    p_drtio_dest = parser.add_argument("-s", "--satellite", default=0,
+                        metavar="DRTIO ID", type=int,
+                        help="specify DRTIO destination that receives this command")
+
     return parser
 
 
@@ -107,7 +112,7 @@ def main():
         core_addr = ddb.get("core", resolve_alias=True)["arguments"]["host"]
     else:
         core_addr = args.device
-    mgmt = CommMgmt(core_addr)
+    mgmt = CommMgmt(core_addr, drtio_dest=args.satellite)
 
     if args.tool == "log":
         if args.action == "set_level":