fix(pm): Restore sleep suspension of devices.

* After the move to `sys_poweroff`, restore the behavior of
  suspending devices before entering sleep state.

3 files changed, 79 insertions, 0 deletions

diff --git a/app/CMakeLists.txt b/app/CMakeLists.txt
index 314714844f..6ef0031102 100644
--- a/app/CMakeLists.txt
+++ b/app/CMakeLists.txt
@@ -8,6 +8,10 @@ set(ZEPHYR_EXTRA_MODULES "${ZMK_EXTRA_MODULES};${CMAKE_CURRENT_SOURCE_DIR}/modul
 find_package(Zephyr REQUIRED HINTS ../zephyr)
 project(zmk)
 
+if(CONFIG_ZMK_SLEEP)
+  zephyr_linker_sources(SECTIONS include/linker/zmk-pm-devices.ld)
+endif()
+
 zephyr_linker_sources(SECTIONS include/linker/zmk-behaviors.ld)
 zephyr_linker_sources(RODATA include/linker/zmk-events.ld)
 
diff --git a/app/include/linker/zmk-pm-devices.ld b/app/include/linker/zmk-pm-devices.ld
new file mode 100644
index 0000000000..93ec5025de
--- /dev/null
+++ b/app/include/linker/zmk-pm-devices.ld
@@ -0,0 +1,9 @@
+/*
+ * Copyright (c) 2023 The ZMK Contributors
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+#include <zephyr/linker/linker-defs.h>
+
+ITERABLE_SECTION_RAM(zmk_pm_device_slots, 4)
diff --git a/app/src/activity.c b/app/src/activity.c
index 1140915168..8f421f85d0 100644
--- a/app/src/activity.c
+++ b/app/src/activity.c
@@ -7,6 +7,8 @@
 #include <zephyr/device.h>
 #include <zephyr/init.h>
 #include <zephyr/kernel.h>
+#include <zephyr/pm/device.h>
+#include <zephyr/pm/device_runtime.h>
 #include <zephyr/sys/poweroff.h>
 
 #include <zephyr/logging/log.h>
@@ -24,6 +26,63 @@ LOG_MODULE_DECLARE(zmk, CONFIG_ZMK_LOG_LEVEL);
 #include <zmk/usb.h>
 #endif
 
+// Reimplement some of the device work from Zephyr PM to work with the new `sys_poweroff` API.
+// TODO: Tweak this to smarter runtime PM of subsystems on sleep.
+
+#ifdef CONFIG_PM_DEVICE
+TYPE_SECTION_START_EXTERN(const struct device *, zmk_pm_device_slots);
+
+#if !defined(CONFIG_PM_DEVICE_RUNTIME_EXCLUSIVE)
+/* Number of devices successfully suspended. */
+static size_t zmk_num_susp;
+
+static int zmk_pm_suspend_devices(void) {
+    const struct device *devs;
+    size_t devc;
+
+    devc = z_device_get_all_static(&devs);
+
+    zmk_num_susp = 0;
+
+    for (const struct device *dev = devs + devc - 1; dev >= devs; dev--) {
+        int ret;
+
+        /*
+         * Ignore uninitialized devices, busy devices, wake up sources, and
+         * devices with runtime PM enabled.
+         */
+        if (!device_is_ready(dev) || pm_device_is_busy(dev) || pm_device_state_is_locked(dev) ||
+            pm_device_wakeup_is_enabled(dev) || pm_device_runtime_is_enabled(dev)) {
+            continue;
+        }
+
+        ret = pm_device_action_run(dev, PM_DEVICE_ACTION_SUSPEND);
+        /* ignore devices not supporting or already at the given state */
+        if ((ret == -ENOSYS) || (ret == -ENOTSUP) || (ret == -EALREADY)) {
+            continue;
+        } else if (ret < 0) {
+            LOG_ERR("Device %s did not enter %s state (%d)", dev->name,
+                    pm_device_state_str(PM_DEVICE_STATE_SUSPENDED), ret);
+            return ret;
+        }
+
+        TYPE_SECTION_START(zmk_pm_device_slots)[zmk_num_susp] = dev;
+        zmk_num_susp++;
+    }
+
+    return 0;
+}
+
+static void zmk_pm_resume_devices(void) {
+    for (int i = (zmk_num_susp - 1); i >= 0; i--) {
+        pm_device_action_run(TYPE_SECTION_START(zmk_pm_device_slots)[i], PM_DEVICE_ACTION_RESUME);
+    }
+
+    zmk_num_susp = 0;
+}
+#endif /* !CONFIG_PM_DEVICE_RUNTIME_EXCLUSIVE */
+#endif /* CONFIG_PM_DEVICE */
+
 bool is_usb_power_present(void) {
 #if IS_ENABLED(CONFIG_USB_DEVICE_STACK)
     return zmk_usb_is_powered();
@@ -70,6 +129,13 @@ void activity_work_handler(struct k_work *work) {
     if (inactive_time > MAX_SLEEP_MS && !is_usb_power_present()) {
         // Put devices in suspend power mode before sleeping
         set_state(ZMK_ACTIVITY_SLEEP);
+
+        if (zmk_pm_suspend_devices() < 0) {
+            LOG_ERR("Failed to suspend all the devices");
+            zmk_pm_resume_devices();
+            return;
+        }
+
         sys_poweroff();
     } else
 #endif /* IS_ENABLED(CONFIG_ZMK_SLEEP) */