From 296dc3b0c42d0bbd49188e06fe080b748388ae3b Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Tue, 11 Apr 2017 03:24:24 +0000
Subject: [PATCH] artiq_coreboot: allow hot-rebooting the device.

---
 artiq/coredevice/comm_mgmt.py         |  9 +++++++--
 artiq/firmware/libproto/mgmt_proto.rs |  6 ++++--
 artiq/firmware/runtime/mgmt.rs        |  9 ++++++++-
 artiq/frontend/artiq_coreboot.py      | 23 +++++++++++++++++++----
 artiq/frontend/artiq_devtool.py       |  2 +-
 5 files changed, 39 insertions(+), 10 deletions(-)

diff --git a/artiq/coredevice/comm_mgmt.py b/artiq/coredevice/comm_mgmt.py
index 36793957b..9f638e780 100644
--- a/artiq/coredevice/comm_mgmt.py
+++ b/artiq/coredevice/comm_mgmt.py
@@ -13,6 +13,7 @@ class Request(Enum):
     SetLogFilter = 3
 
     Hotswap = 4
+    Reboot = 5
 
 
 class Reply(Enum):
@@ -20,7 +21,7 @@ class Reply(Enum):
 
     LogContent = 2
 
-    HotswapImminent = 3
+    RebootImminent = 3
 
 
 class LogLevel(Enum):
@@ -130,4 +131,8 @@ class CommMgmt:
     def hotswap(self, firmware):
         self._write_header(Request.Hotswap)
         self._write_bytes(firmware)
-        self._read_expect(Reply.HotswapImminent)
+        self._read_expect(Reply.RebootImminent)
+
+    def reboot(self):
+        self._write_header(Request.Reboot)
+        self._read_expect(Reply.RebootImminent)
diff --git a/artiq/firmware/libproto/mgmt_proto.rs b/artiq/firmware/libproto/mgmt_proto.rs
index 75e480840..7e299c02b 100644
--- a/artiq/firmware/libproto/mgmt_proto.rs
+++ b/artiq/firmware/libproto/mgmt_proto.rs
@@ -12,6 +12,7 @@ pub enum Request {
     SetLogFilter(LogLevelFilter),
 
     Hotswap(Vec<u8>),
+    Reboot,
 }
 
 pub enum Reply<'a> {
@@ -19,7 +20,7 @@ pub enum Reply<'a> {
 
     LogContent(&'a str),
 
-    HotswapImminent,
+    RebootImminent,
 }
 
 impl Request {
@@ -42,6 +43,7 @@ impl Request {
                 Request::SetLogFilter(level)
             }
             4 => Request::Hotswap(reader.read_bytes()?),
+            5 => Request::Reboot,
             _  => return Err(io::Error::new(io::ErrorKind::InvalidData, "unknown request type"))
         })
     }
@@ -59,7 +61,7 @@ impl<'a> Reply<'a> {
                 writer.write_string(log)?;
             },
 
-            Reply::HotswapImminent => {
+            Reply::RebootImminent => {
                 writer.write_u8(3)?;
             },
         }
diff --git a/artiq/firmware/runtime/mgmt.rs b/artiq/firmware/runtime/mgmt.rs
index f195a6846..8beae791d 100644
--- a/artiq/firmware/runtime/mgmt.rs
+++ b/artiq/firmware/runtime/mgmt.rs
@@ -45,11 +45,18 @@ fn worker(mut stream: &mut TcpStream) -> io::Result<()> {
             },
 
             Request::Hotswap(firmware) => {
-                Reply::HotswapImminent.write_to(stream)?;
+                Reply::RebootImminent.write_to(stream)?;
                 stream.close()?;
                 warn!("hotswapping firmware");
                 unsafe { board::boot::hotswap(&firmware) }
             },
+
+            Request::Reboot => {
+                Reply::RebootImminent.write_to(stream)?;
+                stream.close()?;
+                warn!("rebooting");
+                unsafe { board::boot::reboot() }
+            }
         };
     }
 }
diff --git a/artiq/frontend/artiq_coreboot.py b/artiq/frontend/artiq_coreboot.py
index ca5b7888a..45c317247 100755
--- a/artiq/frontend/artiq_coreboot.py
+++ b/artiq/frontend/artiq_coreboot.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 
 import argparse
+import sys
 import struct
 
 from artiq.tools import verbosity_args, init_logger
@@ -10,14 +11,22 @@ from artiq.coredevice.comm_mgmt import CommMgmt
 
 
 def get_argparser():
-    parser = argparse.ArgumentParser(description="ARTIQ core device hotswap tool")
+    parser = argparse.ArgumentParser(description="ARTIQ core device boot tool")
 
     verbosity_args(parser)
     parser.add_argument("--device-db", default="device_db.pyon",
                        help="device database file (default: '%(default)s')")
 
-    parser.add_argument("image", metavar="IMAGE", type=argparse.FileType('rb'),
-                        help="runtime image to be executed")
+    subparsers = parser.add_subparsers(dest="action")
+
+    p_reboot = subparsers.add_parser("reboot",
+                                     help="reboot the currently running firmware")
+
+    p_hotswap = subparsers.add_parser("hotswap",
+                                      help="load the specified firmware in RAM")
+
+    p_hotswap.add_argument("image", metavar="IMAGE", type=argparse.FileType('rb'),
+                           help="runtime image to be executed")
 
     return parser
 
@@ -29,7 +38,13 @@ def main():
     try:
         core_addr = device_mgr.get_desc("comm")["arguments"]["host"]
         mgmt = CommMgmt(device_mgr, core_addr)
-        mgmt.hotswap(args.image.read())
+        if args.action == "reboot":
+            mgmt.reboot()
+        elif args.action == "hotswap":
+            mgmt.hotswap(args.image.read())
+        else:
+            print("An action needs to be specified.", file=sys.stderr)
+            sys.exit(1)
     finally:
         device_mgr.close_devices()
 
diff --git a/artiq/frontend/artiq_devtool.py b/artiq/frontend/artiq_devtool.py
index 02f189a1c..d93130e51 100755
--- a/artiq/frontend/artiq_devtool.py
+++ b/artiq/frontend/artiq_devtool.py
@@ -193,7 +193,7 @@ def main():
             logger.info("Hotswapping firmware")
             try:
                 subprocess.check_call(["python3",
-                    "-m", "artiq.frontend.artiq_coreboot",
+                    "-m", "artiq.frontend.artiq_coreboot", "hotswap",
                     "/tmp/{target}/software/{firmware}/{firmware}.bin"
                         .format(target=args.target, firmware=firmware)])
             except subprocess.CalledProcessError: