Merge pull request #459 from LedgerHQ/clean/better_shutdown

[clean] Better (and faster) shutdown process

CHANGELOG.md

@@ -5,6 +5,12 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.7.0] - 2024-02-26
+
+### Changed
+- Significative performance improvement on display/snapshot management
+- Simplified HTTP API thread management
+
 ## [0.6.0] - 2024-02-21
 
 ### Added

speculos/api/api.py

@@ -23,13 +23,13 @@
 class ApiRunner(IODevice):
     """Run the Speculos API server in a dedicated thread, with a notification when it stops"""
     def __init__(self, api_port: int) -> None:
-        self._app: Flask
+        self._api_wrapper: ApiWrapper
         # self.sock is used by Screen.add_notifier. Closing self._notify_exit
         # signals it that the API is no longer running.
         self.sock: socket.socket
         self._notify_exit: socket.socket
         self.sock, self._notify_exit = socket.socketpair()
-        self.api_port: int = api_port
+        self._port: int = api_port
         self._api_thread: threading.Thread
 
     @property
@@ -40,32 +40,25 @@ def can_read(self, screen: DisplayNotifier) -> None:
         # Being able to read from the socket only happens when the API server exited.
         raise ReadError("API server exited")
 
-    def _run(self) -> None:
-        try:
-            # threaded must be set to allow serving requests along events streaming
-            self._app.run(host="0.0.0.0", port=self.api_port, threaded=True, use_reloader=False)
-        except Exception as exc:
-            self._app.logger.error("An exception occurred in the Flask API server: %s", exc)
-            raise exc
-        finally:
-            self._notify_exit.close()
-
     def start_server_thread(self,
                             screen_: DisplayNotifier,
                             seph_: SeProxyHal,
                             automation_server: BroadcastInterface) -> None:
-        wrapper = ApiWrapper(screen_, seph_, automation_server)
-        self._app = wrapper.app
-        self._api_thread = threading.Thread(target=self._run, name="API-server", daemon=True)
+        self._api_wrapper = ApiWrapper(self._port, screen_, seph_, automation_server)
+        self._api_thread = threading.Thread(target=self._api_wrapper.run, name="API-server", daemon=True)
         self._api_thread.start()
 
     def stop(self):
-        self._api_thread.join(10)
+        self._notify_exit.close()
 
 
 class ApiWrapper:
-    def __init__(self, screen: DisplayNotifier, seph: SeProxyHal, automation_server: BroadcastInterface):
-
+    def __init__(self,
+                 api_port: int,
+                 screen: DisplayNotifier,
+                 seph: SeProxyHal,
+                 automation_server: BroadcastInterface):
+        self._port = api_port
         static_folder = pkg_resources.resource_filename(__name__, "/static")
         self._app = Flask(__name__, static_url_path="", static_folder=static_folder)
         self._app.env = "development"
@@ -75,7 +68,7 @@ def __init__(self, screen: DisplayNotifier, seph: SeProxyHal, automation_server:
         app_kwargs = {"app": self._app}
         event_kwargs: Dict[str, Any] = {**app_kwargs, "automation_server": automation_server}
 
-        self._api = Api(self.app)
+        self._api = Api(self._app)
 
         self._api.add_resource(APDU, "/apdu", resource_class_kwargs=seph_kwargs)
         self._api.add_resource(Automation, "/automation", resource_class_kwargs=seph_kwargs)
@@ -91,6 +84,6 @@ def __init__(self, screen: DisplayNotifier, seph: SeProxyHal, automation_server:
         self._api.add_resource(WebInterface, "/", resource_class_kwargs=app_kwargs)
         self._api.add_resource(Ticker, "/ticker/", resource_class_kwargs=seph_kwargs)
 
-    @property
-    def app(self) -> Flask:
-        return self._app
+    def run(self):
+        # threaded must be set to allow serving requests along events streaming
+        self._app.run(host="0.0.0.0", port=self._port, threaded=True, use_reloader=False)

speculos/main.py

@@ -527,12 +527,18 @@ def main(prog=None) -> int:
         assert automation_server is not None
         apirun.start_server_thread(screen_notifier, seph, automation_server)
 
-    screen_notifier.run()
-
-    if apirun is not None:
-        apirun.stop()
-
-    s2.close()
-    _, status = os.waitpid(qemu_pid, 0)
-    qemu_exit_status = os.WEXITSTATUS(status)
-    sys.exit(qemu_exit_status)
+    try:
+        screen_notifier.run()
+    except BaseException:
+        # Will deal with exception triggered in the ScreenNotifier, including
+        # KeyboardInterrupt (if not Qt display, else it will segfault)
+        logger.exception("An error occurred")
+        logger.critical("Stopping Speculos")
+    finally:
+        if apirun is not None:
+            apirun.stop()
+
+        s2.close()
+        _, status = os.waitpid(qemu_pid, 0)
+        qemu_exit_status = os.WEXITSTATUS(status)
+        sys.exit(qemu_exit_status)

speculos/mcu/display.py

@@ -8,7 +8,9 @@
 from typing import Any, Dict, IO, List, Optional, Tuple, Union
 
 from speculos.observer import TextEvent
-from .struct import DisplayArgs, MODELS, ServerArgs
+from .struct import DisplayArgs, MODELS, Pixel, ServerArgs
+
+PixelColorMapping = Dict[Pixel, int]
 
 
 class IODevice(ABC):
@@ -70,8 +72,8 @@ class FrameBuffer:
     }
 
     def __init__(self, model: str):
-        self.pixels: Dict[Tuple[int, int], int] = {}
-        self.screenshot_pixels: Dict[Tuple[int, int], int] = {}
+        self.pixels: PixelColorMapping = {}
+        self.screenshot_pixels: PixelColorMapping = {}
         self.default_color = 0
         self.draw_default_color = False
         self._public_screenshot_value = b''
@@ -131,7 +133,7 @@ def take_screenshot(self) -> Tuple[Tuple[int, int], bytes]:
         return self.current_screen_size, self._get_image()
 
     def update_screenshot(self) -> None:
-        self.screenshot_pixels = {**self.screenshot_pixels, **self.pixels}
+        self.screenshot_pixels.update(self.pixels)
 
     def update_public_screenshot(self) -> None:
         # Stax only

speculos/mcu/struct.py

@@ -1,19 +1,14 @@
 from dataclasses import dataclass
 from typing import Any, Dict, NamedTuple, Optional, Tuple
 
-# from speculos.mcu.apdu import ApduServer
-# from speculos.mcu.seproxyhal import SeProxyHal
-# from speculos.mcu.button_tcp import FakeButton
-# from speculos.mcu.finger_tcp import FakeFinger
-# from speculos.mcu.vnc import VNC
-# from speculos.api.api import ApiRunner
+Pixel = Tuple[int, int]
 
 
 @dataclass
 class Model:
     name: str
     screen_size: Tuple[int, int]
-    box_position: Tuple[int, int]
+    box_position: Pixel
     box_size: Tuple[int, int]
 
 

speculos/mcu/vnc.py

@@ -11,7 +11,7 @@
 import sys
 from typing import IO, Optional, Tuple
 
-from .display import DisplayNotifier, IODevice
+from .display import DisplayNotifier, IODevice, PixelColorMapping
 
 
 class VNC(IODevice):
@@ -22,13 +22,13 @@ def __init__(self,
                  verbose: bool = False):
         self.logger = logging.getLogger("vnc")
 
-        self.width, self.height = screen_size
+        self._width, self._height = screen_size
         path = os.path.dirname(os.path.realpath(__file__))
         server = os.path.join(path, '../resources/vnc_server')
         cmd = [server]
 
         # custom options
-        cmd += ['-s', f'{self.width}x{self.height}']
+        cmd += ['-s', f'{self._width}x{self._height}']
         if verbose:
             cmd += ['-v']
 
@@ -48,7 +48,7 @@ def file(self) -> IO[bytes]:
         assert self.subprocess.stdout is not None
         return self.subprocess.stdout
 
-    def redraw(self, pixels, default_color):
+    def redraw(self, pixels: PixelColorMapping, default_color: int) -> None:
         '''The framebuffer was updated, forward everything to the VNC server.'''
 
         # int.to_bytes() is super slow, hence the manual encoding
@@ -68,10 +68,11 @@ def redraw(self, pixels, default_color):
                 buf[i + 8] = 0x0a
                 i += 9
 
+        assert self.subprocess.stdin is not None
         self.subprocess.stdin.write(buf)
         self.subprocess.stdin.flush()
 
-    def can_read(self, screen: DisplayNotifier):
+    def can_read(self, screen: DisplayNotifier) -> None:
         '''Process a new keyboard or mouse event message from the VNC server.'''
 
         data = b''