commit 797756460a058c597aaabde6a8a69470c4fb2ad2
author Kelvin Zhang <zhangkelvin@google.com> Fri Feb 19 16:05:08 2021 -0500
committer Treehugger Robot <treehugger-gerrit@google.com> Thu Mar 11 20:58:50 2021 +0000
tree 7c39bf1be45641c06cfb4c66d1700212b5611805
parent 1c21eb8364f4d4468ac1f9aa562e402901f39bad

Improve load speed for python payload script

When using update_payload module to read a full OTA, it will attempt to
load the entire payload into memory first, which is slow. Use mmap
instead, as parsing manifest doesn't require the entire payload to be in
memory, only the first couple of MB.
In case of opening a zipfile, use zfp.open(), which returns a streaming
interface.

Test: import update_payload, load an OTA
Change-Id: I3f7c7515dbf5576ccfda8d844fe22e9be34fdaed

scripts/update_payload/payload.py

diff --git a/scripts/update_payload/payload.py b/scripts/update_payload/payload.py
index 6b28046..998703a 100644
--- a/scripts/update_payload/payload.py
+++ b/scripts/update_payload/payload.py
@@ -21,6 +21,7 @@
 
 import hashlib
 import io
+import mmap
 import struct
 import zipfile
 
@@ -123,10 +124,11 @@
     """
     if zipfile.is_zipfile(payload_file):
       with zipfile.ZipFile(payload_file) as zfp:
-        with zfp.open("payload.bin") as payload_fp:
-          self.payload_file = io.BytesIO(payload_fp.read())
+        self.payload_file = zfp.open("payload.bin", "r")
     elif isinstance(payload_file, str):
-      self.payload_file = open(payload_file, "rb")
+      payload_fp = open(payload_file, "rb")
+      payload_bytes = mmap.mmap(payload_fp.fileno(), 0, access=mmap.ACCESS_READ)
+      self.payload_file = io.BytesIO(payload_bytes)
     else:
       self.payload_file = payload_file
     self.payload_file_offset = payload_file_offset