Turn the display off reliably when the program exits or the Windows session ends

library/display.py

@@ -149,6 +149,15 @@ def turn_off(self):
         # Turn off backplate RGB LED
         self.lcd.SetBackplateLedColor(led_color=(0, 0, 0))
 
+    def close(self):
+        # Close the communication with the display. Called explicitly on exit: the program ends
+        # with os._exit(), so the LcdComm* driver destructors that would close it never run.
+        try:
+            self.lcd.closeSerial()
+        except Exception as e:
+            # Closing must not prevent the program from exiting
+            logger.warning("Failed to close communication with display: %s" % str(e), exc_info=True)
+
     def display_static_images(self):
         if config.THEME_DATA.get('static_images', False):
             for image in config.THEME_DATA['static_images']:

library/lcd/lcd_simulated.py

@@ -70,6 +70,7 @@ def __init__(self, com_port: str = "AUTO", display_width: int = 320, display_hei
             logger.debug("To see your simulated screen, open http://%s:%d in a browser" % ("localhost", WEBSERVER_PORT))
             threading.Thread(target=self.webServer.serve_forever).start()
         except OSError:
+            self.webServer = None
             logger.error("Error starting webserver! An instance might already be running on port %d." % WEBSERVER_PORT)
 
     def __del__(self):
@@ -80,8 +81,9 @@ def auto_detect_com_port() -> Optional[str]:
         return None
 
     def closeSerial(self):
-        logger.debug("Shutting down web server")
-        self.webServer.shutdown()
+        if self.webServer is not None:
+            logger.debug("Shutting down web server")
+            self.webServer.shutdown()
 
     def InitializeComm(self):
         pass

library/scheduler.py

@@ -20,6 +20,7 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
+import queue
 import sched
 import threading
 import time
@@ -28,9 +29,13 @@
 
 import library.config as config
 import library.stats as stats
+from library.log import logger
 
 STOPPING = False
 
+# Time the queue handler waits for a new request before checking whether to stop, in seconds
+QUEUE_GET_TIMEOUT = 0.1
+
 
 def async_job(threadname=None):
     """ wrapper to handle asynchronous threads """
@@ -58,10 +63,14 @@ def decorator(func):
 
         def periodic(scheduler, periodic_interval, action, actionargs=()):
             """ Wrap the scheduler with our periodic interval """
-            if not STOPPING:
-                # If the program is not stopping: re-schedule the task for future execution
-                scheduler.enter(periodic_interval, 1, periodic,
-                                (scheduler, periodic_interval, action, actionargs))
+            if STOPPING:
+                # The program is stopping: do not run the action, it would queue a new frame
+                # behind the "turn off screen" requests that clean_stop() just queued
+                return
+
+            # Re-schedule the task for future execution
+            scheduler.enter(periodic_interval, 1, periodic,
+                            (scheduler, periodic_interval, action, actionargs))
             action(*actionargs)
 
         @wraps(func)
@@ -186,20 +195,34 @@ def PingStats():
 
 
 @async_job("Queue_Handler")
-@schedule(timedelta(milliseconds=1).total_seconds())
 def QueueHandler():
-    # Do next action waiting in the queue
-    if STOPPING:
-        # Empty the action queue to allow program to exit cleanly
-        while not config.update_queue.empty():
-            f, args = config.update_queue.get()
-            f(*args)
-    else:
-        # Execute first action in the queue
-        f, args = config.update_queue.get()
-        if f:
+    # Not scheduled with @schedule: the scheduled job blocked in an untimed get(), so it could
+    # never notice STOPPING, and it was not re-scheduled once STOPPING was set. This loop keeps
+    # draining until the queue is empty, then ends.
+    while True:
+        try:
+            f, args = config.update_queue.get(timeout=QUEUE_GET_TIMEOUT)
+        except queue.Empty:
+            if STOPPING:
+                return
+            continue
+
+        try:
             f(*args)
+        except Exception:
+            # Keep draining: this thread is the only consumer, and if it dies the queue is
+            # never emptied again and every later clean_stop() waits for its full timeout
+            logger.error("Failed to send a request to the display", exc_info=True)
+        finally:
+            config.update_queue.task_done()
 
 
 def is_queue_empty() -> bool:
-    return config.update_queue.empty()
+    # The queue being empty is not enough: the request taken from it may still be being sent
+    # to the display, and exiting at that point would cut the data in the middle. The count of
+    # unfinished tasks only drops back to zero once task_done() is called, i.e. once the request
+    # has actually been processed. That counter is what queue.Queue.join() waits on, but only
+    # its behaviour is documented, not the attribute name: reading it directly relies on a
+    # CPython implementation detail, in exchange for the timeout that join() does not offer.
+    # https://docs.python.org/3/library/queue.html#queue.Queue.join
+    return config.update_queue.unfinished_tasks == 0

main.py

@@ -39,14 +39,18 @@
     import platform
     import signal
     import subprocess
+    import threading
     import time
     from pathlib import Path
     from PIL import Image
 
     if platform.system() == 'Windows':
+        import ctypes
         import win32api
         import win32con
         import win32gui
+        import win32process
+        import win32ts
 
     from library.log import logger
     import library.scheduler as scheduler
@@ -90,16 +94,43 @@ def wait_for_empty_queue(timeout: int = 5):
 
         logger.debug("(Waited %.1fs)" % wait_time)
 
-    def clean_stop(tray_icon=None):
-        # Turn screen and LEDs off before stopping
-        display.turn_off()
+    # Set when the queue handler thread is started, so clean_stop() can wait for it
+    queue_handler_thread = None
 
-        # Do not stop the program now in case data transmission was in progress
-        # Instead, ask the scheduler to empty the action queue before stopping
-        scheduler.STOPPING = True
+    # Guards the shutdown of the display: the end of the session, the log off and the tray
+    # Exit can fire at once, and running the sequence twice would cut the data being sent
+    stop_display_lock = threading.Lock()
+    display_stopped = False
 
-        # Waiting for all pending request to be sent to display
-        wait_for_empty_queue(5)
+    def stop_display(timeout: int = 5):
+        """Turn the display off and close the communication with it. Only runs once."""
+        global display_stopped
+
+        with stop_display_lock:
+            if display_stopped:
+                return
+            display_stopped = True
+
+            # Turn screen and LEDs off before stopping
+            display.turn_off()
+
+            # Do not stop the program now in case data transmission was in progress
+            # Instead, ask the scheduler to empty the action queue before stopping
+            scheduler.STOPPING = True
+
+            # Waiting for all pending request to be sent to display
+            wait_for_empty_queue(timeout)
+
+            # Wait for the queue handler to stop before closing: a request still being sent
+            # would fail on the closed port, and WriteLine() reopens the port when that happens
+            if queue_handler_thread is not None:
+                queue_handler_thread.join(timeout=1)
+
+            # Close the communication with the display
+            display.close()
+
+    def clean_stop(tray_icon=None, timeout: int = 5):
+        stop_display(timeout)
 
         # Remove tray icon just before exit
         if tray_icon:
@@ -145,30 +176,82 @@ def on_clean_exit(*args):
     if platform.system() == "Windows":
         def on_win32_ctrl_event(event):
             """Handle Windows console control events (like Ctrl-C)."""
-            if event in (win32con.CTRL_C_EVENT, win32con.CTRL_BREAK_EVENT, win32con.CTRL_CLOSE_EVENT):
+            if event in (win32con.CTRL_C_EVENT, win32con.CTRL_BREAK_EVENT, win32con.CTRL_CLOSE_EVENT,
+                         win32con.CTRL_LOGOFF_EVENT, win32con.CTRL_SHUTDOWN_EVENT):
                 logger.debug("Caught Windows control event %s, exiting" % event)
                 clean_stop()
             return 0
 
 
+        # Message sent to a window registered with WTSRegisterSessionNotification, and the
+        # session change it reports on log off. Neither of them is exposed by pywin32.
+        WM_WTSSESSION_CHANGE = 0x02B1
+        WTS_SESSION_LOGOFF = 0x5
+
+
+        def shutdown_block_reason(hWnd, reason=None):
+            """Tell the system why the shutdown is being held, or clear that reason."""
+            try:
+                if reason:
+                    ctypes.windll.user32.ShutdownBlockReasonCreate(ctypes.c_void_p(hWnd),
+                                                                   ctypes.c_wchar_p(reason))
+                else:
+                    ctypes.windll.user32.ShutdownBlockReasonDestroy(ctypes.c_void_p(hWnd))
+            except Exception as e:
+                # Asking the system to wait is a courtesy: the shutdown goes on without it
+                logger.warning("Failed to set the shutdown block reason: %s" % str(e))
+
+
         def on_win32_wm_event(hWnd, msg, wParam, lParam):
             """Handle Windows window message events (like ENDSESSION, CLOSE, DESTROY)."""
-            logger.debug("Caught Windows window message event %s" % msg)
+            logger.debug("Caught Windows window message event 0x%04X (wParam 0x%X)" % (msg, wParam))
+
             if msg == win32con.WM_POWERBROADCAST:
                 # WM_POWERBROADCAST is used to detect computer going to/resuming from sleep
                 if wParam == win32con.PBT_APMSUSPEND:
                     logger.info("Computer is going to sleep, display will turn off")
                     display.turn_off()
+                    # Suspending interrupts the transfer, so wait for the command to be sent.
+                    # The scheduler keeps running: the display is used again on resume
+                    wait_for_empty_queue(2)
                 elif wParam == win32con.PBT_APMRESUMEAUTOMATIC:
                     logger.info("Computer is resuming from sleep, display will turn on")
                     display.turn_on()
                     # Some models have troubles displaying back the previous bitmap after being turned off/on
                     display.display_static_images()
                     display.display_static_text()
-            else:
-                # For any other events, the program will stop
-                logger.info("Program will now exit")
-                clean_stop()
+                return True
+
+            if msg == win32con.WM_QUERYENDSESSION:
+                # The session is ending: turn the display off while the system is still
+                # waiting for an answer, but do not exit here. Returning TRUE is what lets the
+                # shutdown carry on, and the program stops on the WM_ENDSESSION that follows.
+                # The work is done now rather than on WM_ENDSESSION because that message is not
+                # guaranteed to be delivered: once every application has answered, the system is
+                # free to terminate them. The cost is a display left off if the shutdown ends up
+                # being cancelled, which a restart of the program undoes.
+                logger.info("Session is ending, display will turn off")
+                shutdown_block_reason(hWnd, "Turning the display off")
+                stop_display(2)
+                shutdown_block_reason(hWnd)
+                return True
+
+            if msg == win32con.WM_ENDSESSION and not wParam:
+                # The session is not ending after all: another application refused the shutdown.
+                # The display was already turned off by WM_QUERYENDSESSION and cannot be brought
+                # back without restarting the program, but there is no reason to exit either.
+                logger.warning("Session is not ending after all, display stays off")
+                return 0
+
+            if msg == WM_WTSSESSION_CHANGE:
+                # This notification also reports lock, unlock and fast user switching
+                if wParam != WTS_SESSION_LOGOFF:
+                    return 0
+                logger.info("User is logging off")
+
+            # WM_ENDSESSION and any other event: the program will stop
+            logger.info("Program will now exit")
+            clean_stop(timeout=2)
 
     # Create a tray icon for the program, with an Exit entry in menu
     try:
@@ -214,14 +297,20 @@ def on_win32_wm_event(hWnd, msg, wParam, lParam):
         signal.signal(signal.SIGQUIT, on_signal_caught)
     if platform.system() == "Windows":
         win32api.SetConsoleCtrlHandler(on_win32_ctrl_event, True)
+        # Be notified before regular applications when the session ends: the turn-off command
+        # travels over a serial port and has to reach the display before the system goes down
+        try:
+            win32process.SetProcessShutdownParameters(0x3FF, 0)
+        except Exception as e:
+            logger.warning("Failed to raise the shutdown notification priority: %s" % str(e))
 
     # Initialize the display
     logger.info("Initialize display")
     display.initialize_display()
 
     # Start serial queue handler
     if not args.theme_screenshots:
-        scheduler.QueueHandler()
+        queue_handler_thread = scheduler.QueueHandler()
 
     # Create all static images
     display.display_static_images()
@@ -298,7 +387,8 @@ def on_win32_wm_event(hWnd, msg, wParam, lParam):
                       win32con.WM_QUIT: on_win32_wm_event,
                       win32con.WM_DESTROY: on_win32_wm_event,
                       win32con.WM_CLOSE: on_win32_wm_event,
-                      win32con.WM_POWERBROADCAST: on_win32_wm_event}
+                      win32con.WM_POWERBROADCAST: on_win32_wm_event,
+                      WM_WTSSESSION_CHANGE: on_win32_wm_event}
 
         wndclass.lpfnWndProc = messageMap
 
@@ -316,10 +406,21 @@ def on_win32_wm_event(hWnd, msg, wParam, lParam):
                                            0,
                                            hinst,
                                            None)
-            while True:
-                # Receive and dispatch window messages
-                win32gui.PumpWaitingMessages()
-                time.sleep(0.5)
+            logger.debug("Windows event window created (handle %s)" % hwnd)
+
+            # An application with no visible window is not a reliable target for the end of
+            # session broadcast: this notification goes to the window that asked for it
+            try:
+                win32ts.WTSRegisterSessionNotification(hwnd, win32ts.NOTIFY_FOR_THIS_SESSION)
+            except Exception as e:
+                logger.warning("Failed to register for session notifications: %s" % str(e))
+
+            # Receive and dispatch window messages. PumpMessages blocks until the loop is
+            # over, so events are handled as they arrive instead of on the next poll
+            win32gui.PumpMessages()
+
+            logger.info("Program will now exit")
+            clean_stop()
 
         except Exception as e:
             logger.error("Exception while creating event window: %s" % str(e))