hpr: aei: Add support for updating AEI HPR PSUs

Summary:
(1) Add config for HPR devices. This allows us to do different things
    for HPR (Example, HPR has different block sizes and binary embeds the
    block#)
(2) Add timeouts everywhere.
(3) Retry isp_enter.
(4) Move isp_enter and isp_exit into an isp context manager. So, we
    always exit ISP mode when we fail.
(5) If HPR do not prefix block# to the data. Assume its part of the 68b
    block. Also this implies we no longer validate the response.
(6) wait_update_complete now checks if the PSU has rebooted before
    waiting for bit-changes.
(7) Pass address to suppress_monitoring so it can pause PMM along with
    rackmon monitoring.

Test Plan:
Update AEI PSU: P1558749788
```
psu-update-aei.py --addr 144 --device hpr ./AEI/PSU/ORV3_PSU_HPR_v02K_FW_TEST.bin
```

Reviewed By: doranand

Differential Revision: D62052667

fbshipit-source-id: c702a3aa16db4c661b76417032f73f926bc366ec

common/recipes-core/psu-update/files/psu-update-aei.py

@@ -4,10 +4,11 @@
 import sys
 import time
 import traceback
+from contextlib import contextmanager
 
-from modbus_update_helper import get_parser, print_perc, suppress_monitoring
+from modbus_update_helper import get_parser, print_perc, retry, suppress_monitoring
 
-from pyrmd import ModbusException, RackmonInterface as rmd
+from pyrmd import ModbusException, ModbusTimeout, RackmonInterface as rmd
 
 
 ISP_CTRL_CMD_ENTER = 0x1
@@ -17,8 +18,24 @@
 BLOCK_WRITE_BLOCK_WRITTEN = 0x1
 BLOCK_WRITE_INVALID_RANGE = 0x2
 
-
+device_params = {
+    "orv3": {
+        "block_size": 64,
+        "embedded_block_no": False,
+    },
+    "hpr": {
+        "block_size": 68,
+        "embedded_block_no": True,
+    },
+}
 parser = get_parser()
+parser.add_argument(
+    "--device",
+    type=str,
+    default="orv3",
+    choices=list(device_params.keys()),
+    help="Pick device type",
+)
 
 
 class StatusRegister:
@@ -66,7 +83,7 @@ def getSet(self):
 
 
 class BadAEIResponse(ModbusException):
-    ...
+    pass
 
 
 def load_file(path):
@@ -77,16 +94,17 @@ def load_file(path):
 
 def isp_get_status(addr):
     req = struct.pack(">BBB", addr, 0x43, 2)
-    resp = rmd.raw(req, expected=7)
+    resp = rmd.raw(req, expected=7, timeout=2000)
     raddr, rfunc, bcount, status = struct.unpack(">BBBH", resp)
     if raddr != addr or rfunc != 0x43:
         raise BadAEIResponse()
     return StatusRegister(status)
 
 
+@retry(5, delay=1.0)
 def isp_enter(addr):
     req = struct.pack(">BBB", addr, 0x42, ISP_CTRL_CMD_ENTER)
-    resp = rmd.raw(req, expected=7)
+    resp = rmd.raw(req, expected=7, timeout=5000)
     raddr, rfun, rcmd, _ = struct.unpack(">BBBB", resp)
     if raddr != addr or rfun != 0x42 or rcmd != ISP_CTRL_CMD_ENTER:
         raise BadAEIResponse()
@@ -100,24 +118,51 @@ def isp_enter(addr):
 
 def isp_exit(addr):
     req = struct.pack(">BBB", addr, 0x42, ISP_CTRL_CMD_EXIT)
-    resp = rmd.raw(req, expected=7)
-    raddr, rfun, rcmd, status, _ = struct.unpack(">BBBBB", resp)
-    if raddr != addr or rfun != 0x42 or rcmd != ISP_CTRL_CMD_EXIT:
-        raise BadAEIResponse()
-    # 0 - success, 1 - failure
-    if status != 0:
-        print(f"Status from exit returned {status}")
-        raise BadAEIResponse()
-
-
-def isp_flash_block(addr, block_no, block):
-    bsize = 64
+    try:
+        resp = rmd.raw(req, expected=7, timeout=5000)
+        raddr, rfun, rcmd, status, _ = struct.unpack(">BBBBB", resp)
+        if raddr != addr or rfun != 0x42 or rcmd != ISP_CTRL_CMD_EXIT:
+            raise BadAEIResponse()
+        # 0 - success, 1 - failure
+        if status != 0:
+            print(f"Status from exit returned {status}")
+            raise BadAEIResponse()
+    except ModbusTimeout:
+        print("WARNING: Ignoring timeout of ISP_EXIT")
+        time.sleep(5)
+
+
+@contextmanager
+def isp(addr):
+    """
+    Ensure we always exit ISP mode
+    """
+    try:
+        print("Enter ISP Mode")
+        isp_enter(addr)
+        yield
+    finally:
+        print("Exit ISP Mode")
+        isp_exit(addr)
+
+
+def isp_flash_block(addr, block_no, block, params):
+    bsize = params["block_size"]
     if len(block) != bsize:
-        raise ValueError(f"Unexpected block of size {len(block)}")
-    req = struct.pack(">BBBH", addr, 0x45, bsize + 2, block_no) + block
-    resp = rmd.raw(req, expected=8)
+        print(f"Ignoring unexpected block size {len(block)}")
+        return
+    if params["embedded_block_no"]:
+        # Block# is embedded in the binary block read from file.
+        req = struct.pack(">BBB", addr, 0x45, bsize) + block
+    else:
+        # We need to prefix the 2 byte block#.
+        req = struct.pack(">BBBH", addr, 0x45, bsize + 2, block_no) + block
+    resp = rmd.raw(req, expected=8, timeout=2000)
     raddr, rfunc, rlen, rblock, rcode = struct.unpack(">BBBHB", resp)
-    if raddr != addr or rfunc != 0x45 or rlen != 0x3 or rblock != block_no:
+    if raddr != addr or rfunc != 0x45 or rlen != 0x3:
+        print("Bad Block write response:", resp)
+        raise BadAEIResponse()
+    if not params["embedded_block_no"] and rblock != block_no:
         print("Bad Block write response:", resp)
         raise BadAEIResponse()
     if rcode == BLOCK_WRITE_BLOCK_WRITTEN:
@@ -130,15 +175,14 @@ def isp_flash_block(addr, block_no, block):
         raise BadAEIResponse()
 
 
-def transfer_image(addr, image):
-    block_size = 64
+def transfer_image(addr, image, params):
+    block_size = params["block_size"]
     sent_blocks = 0
     total_blocks = len(image) // block_size
     if len(image) % block_size != 0:
-        print("Unsupported file of size {len(image)} has partial block size at the end")
-        raise BadAEIResponse()
+        print(f"Ignoring partial block at end size={len(image)}")
     for i in range(0, len(image), block_size):
-        isp_flash_block(addr, sent_blocks, image[i : i + block_size])
+        isp_flash_block(addr, sent_blocks, image[i : i + block_size], params)
         sent_blocks += 1
         print_perc(
             sent_blocks * 100.0 / total_blocks,
@@ -162,14 +206,14 @@ def wait_update_complete(addr):
     for _ in range(max_time):
         try:
             status = isp_get_status(addr)
+            if status.val == 0:
+                print("PSU Rebooted!")
+                break
             set_fields = status.getSet()
             errors = set_fields.intersection(error_fields)
             if errors:
                 raise ValueError("ERROR: " + " ".join(errors))
             if last_set != set_fields:
-                if status.val == 0:
-                    print("PSU Rebooted!")
-                    break
                 added = set_fields - last_set
                 last_set = set_fields
                 print("ALERT:", " ".join(added))
@@ -181,31 +225,30 @@ def wait_update_complete(addr):
         raise ValueError("Timed out waiting for PSU to reset")
 
 
-def update_device(addr, filename):
+def update_device(addr, filename, params):
     print("Parsing Firmware")
     binimg = load_file(filename)
-    print("Enter ISP Mode")
-    isp_enter(addr)
-    time.sleep(10.0)  # Wait for PSU to erase flash.
-    print("Transfer Image")
-    transfer_image(addr, binimg)
-    print("Check Status")
-    st = isp_get_status(addr)
-    if not st["FULL_IMAGE_RECEIVED"]:
-        print("PSU Did not receive the full image. Status:", str(st))
-        raise BadAEIResponse()
-    print("Exit ISP Mode")
-    isp_exit(addr)
+    with isp(addr):
+        time.sleep(10.0)  # Wait for PSU to erase flash.
+        print("Transfer Image")
+        transfer_image(addr, binimg, params)
+        print("Check Status")
+        st = isp_get_status(addr)
+        if not st["FULL_IMAGE_RECEIVED"]:
+            print("PSU Did not receive the full image. Status:", str(st))
+            raise BadAEIResponse()
     print("Wait update to complete. Sleeping for ~6min")
     wait_update_complete(addr)
     print("done")
 
 
 def main():
+    global device_params
     args = parser.parse_args()
-    with suppress_monitoring():
+    params = device_params[args.device]
+    with suppress_monitoring(args.addr):
         try:
-            update_device(args.addr, args.file)
+            update_device(args.addr, args.file, params)
         except Exception as e:
             print("Firmware update failed %s" % str(e))
             traceback.print_exc()