pc_rpc: server identification support

sbourdeauducq · sbourdeauducq · commit 4cae5531ec05 · 2014-10-28T15:45:56.000+08:00
diff --git a/artiq/devices/pdq2/pdq2-client b/artiq/devices/pdq2/pdq2-client
@@ -54,7 +54,7 @@ def _get_args():
 
 def _main():
     args = _get_args()
-    dev = Client(args.server, args.port)
+    dev = Client(args.server, args.port, "pdq2")
     dev.init()
 
     if args.reset:
diff --git a/artiq/devices/pdq2/pdq2-controller b/artiq/devices/pdq2/pdq2-controller
@@ -355,7 +355,8 @@ def main():
 
     dev = Pdq2(serial=args.serial)
     try:
-        simple_server_loop(dev, args.bind, args.port)
+        simple_server_loop(dev, "pdq2", args.bind, args.port,
+                           id_parameters="serial="+str(args.serial))
     finally:
         dev.close()
 
diff --git a/artiq/management/pc_rpc.py b/artiq/management/pc_rpc.py
@@ -14,8 +14,16 @@
 
 
 class RemoteError(Exception):
-    """Exception raised when a RPC failed or raised an exception on the
-    remote (server) side.
+    """Raised when a RPC failed or raised an exception on the remote (server)
+    side.
+
+    """
+    pass
+
+
+class IncompatibleServer(Exception):
+    """Raised by the client when attempting to connect to a server that does
+    not have the expected type.
 
     """
     pass
@@ -44,10 +52,21 @@ class Client:
         hostname or a IPv4 or IPv6 address (see
         ``socket.create_connection`` in the Python standard library).
     :param port: TCP port to use.
+    :param expected_id_type: Server type to expect. ``IncompatibleServer`` is
+        raised when the types do not match. Use ``None`` to accept any server
+        type.
 
     """
-    def __init__(self, host, port):
+    def __init__(self, host, port, expected_id_type):
         self.socket = socket.create_connection((host, port))
+        self._identify(expected_id_type)
+
+    def get_rpc_id(self):
+        """Returns a dictionary containing the identification information of
+        the server.
+
+        """
+        return self._server_identification
 
     def close_rpc(self):
         """Closes the connection to the RPC server.
@@ -57,8 +76,7 @@ def close_rpc(self):
         """
         self.socket.close()
 
-    def _do_rpc(self, name, args, kwargs):
-        obj = {"action": "call", "name": name, "args": args, "kwargs": kwargs}
+    def _send_recv(self, obj):
         line = pyon.encode(obj) + "\n"
         self.socket.sendall(line.encode())
 
@@ -69,6 +87,19 @@ def _do_rpc(self, name, args, kwargs):
                 break
             buf += more.decode()
         obj = pyon.decode(buf)
+
+        return obj
+
+    def _identify(self, expected_id_type):
+        obj = {"action": "identify"}
+        self._server_identification = self._send_recv(obj)
+        if (expected_id_type is not None
+                and self._server_identification["type"] != expected_id_type):
+            raise IncompatibleServer
+
+    def _do_rpc(self, name, args, kwargs):
+        obj = {"action": "call", "name": name, "args": args, "kwargs": kwargs}
+        obj = self._send_recv(obj)
         if obj["result"] == "ok":
             return obj["ret"]
         elif obj["result"] == "error":
@@ -94,10 +125,16 @@ class Server:
 
     :param target: Object providing the RPC methods to be exposed to the
         client.
+    :param id_type: A string identifying the server type. Clients use it to
+        verify that they are connected to the proper server.
+    :param id_parameters: An optional human-readable string giving more
+        information about the parameters of the server.
 
     """
-    def __init__(self, target):
+    def __init__(self, target, id_type, id_parameters=None):
         self.target = target
+        self.id_type = id_type
+        self.id_parameters = id_parameters
         self._client_tasks = set()
 
     @asyncio.coroutine
@@ -156,22 +193,25 @@ def _handle_connection_task(self, reader, writer):
                                "traceback": traceback.format_exc()}
                     line = pyon.encode(obj) + "\n"
                     writer.write(line.encode())
+                elif action == "identify":
+                    obj = {"type": self.id_type}
+                    if self.id_parameters is not None:
+                        obj["parameters"] = self.id_parameters
+                    line = pyon.encode(obj) + "\n"
+                    writer.write(line.encode())
         finally:
             writer.close()
 
 
-def simple_server_loop(target, host, port):
+def simple_server_loop(target, id_type, host, port, id_parameters=None):
     """Runs a server until an exception is raised (e.g. the user hits Ctrl-C).
 
-    :param target: Object providing the RPC methods to be exposed to the
-        client.
-    :param host: Bind address of the server.
-    :param port: TCP port to bind to.
+    See ``Server`` for a description of the parameters.
 
     """
     loop = asyncio.get_event_loop()
     try:
-        server = Server(target)
+        server = Server(target, id_type, id_parameters)
         loop.run_until_complete(server.start(host, port))
         try:
             loop.run_forever()
diff --git a/doc/manual/writing_a_driver.rst b/doc/manual/writing_a_driver.rst
@@ -23,7 +23,7 @@ To turn it into a server, we use :class:`artiq.management.pc_rpc`. Import the fu
 and add a ``main`` function that is run when the program is executed: ::
 
     def main():
-        simple_server_loop(Hello(), "::1", 7777)
+        simple_server_loop(Hello(), "hello", "::1", 7777)
 
     if __name__ == "__main__":
         main()
@@ -49,6 +49,11 @@ and verify that you can connect to the TCP port: ::
 
 :tip: Use the key combination Ctrl-AltGr-9 to get the ``telnet>`` prompt, and enter ``close`` to quit Telnet. Quit the controller with Ctrl-C.
 
+Also verify that you can get the type of the server (the "hello" string passed to ``simple_server_loop``) using the ``identify-controller`` program from the ARTIQ front-end tools: ::
+
+    ./identify-controller ::1 7777
+    Type:       hello
+
 The client
 ----------
 
@@ -62,7 +67,7 @@ Create a ``hello-client`` file with the following contents: ::
 
 
     def main():
-        remote = Client("::1", 7777)
+        remote = Client("::1", 7777, "hello")
         try:
             remote.message("Hello World!")
         finally:
diff --git a/frontend/artiq b/frontend/artiq
@@ -22,7 +22,7 @@ def _get_args():
 
 def main():
     args = _get_args()
-    remote = Client(args.server, args.port)
+    remote = Client(args.server, args.port, "master")
     try:
         for path, name, timeout in args.run_once:
             remote.run_once(
diff --git a/frontend/artiqd b/frontend/artiqd
@@ -25,7 +25,7 @@ def main():
         scheduler = Scheduler()
         loop.run_until_complete(scheduler.start())
         try:
-            server = Server(scheduler)
+            server = Server(scheduler, "master")
             loop.run_until_complete(server.start(args.bind, args.port))
             try:
                 loop.run_forever()
diff --git a/frontend/identify-controller b/frontend/identify-controller
@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+
+import argparse
+
+from artiq.management.pc_rpc import Client
+
+
+def _get_args():
+    parser = argparse.ArgumentParser(
+        description="ARTIQ controller identification tool")
+    parser.add_argument("server",
+                        help="hostname or IP of the controller to connect to")
+    parser.add_argument("port", type=int,
+                        help="TCP port to use to connect to the controller")
+    return parser.parse_args()
+
+
+def main():
+    args = _get_args()
+    remote = Client(args.server, args.port, None)
+    try:
+        ident = remote.get_rpc_id()
+    finally:
+        remote.close_rpc()
+    print("Type:       " + ident["type"])
+    if "parameters" in ident:
+        print("Parameters: " + ident["parameters"])
+
+if __name__ == "__main__":
+    main()
diff --git a/test/pc_rpc.py b/test/pc_rpc.py
@@ -25,7 +25,8 @@ def test_echo(self):
                 for attempt in range(100):
                     time.sleep(.2)
                     try:
-                        remote = pc_rpc.Client(test_address, test_port)
+                        remote = pc_rpc.Client(test_address, test_port,
+                                               "test")
                     except ConnectionRefusedError:
                         pass
                     else:
@@ -65,7 +66,7 @@ def run_server():
     loop = asyncio.get_event_loop()
     try:
         echo = Echo()
-        server = pc_rpc.Server(echo)
+        server = pc_rpc.Server(echo, "test")
         loop.run_until_complete(server.start(test_address, test_port))
         try:
             loop.run_until_complete(echo.wait_quit())
