Add boot completed cmd to LMKD

Add new command to allow for post-boot actions to occur. This
will allow for the MemEventListener to start after the BPF files are
loaded, removing the need to stall during boot-up until they are loaded
in the device.

Test: Verified memevent listener initialized post-boot
Test: Verified LMKD no longer stalls until BPF progs are loaded
Bug: 331008250
Bug: 244232958
Change-Id: I55f97b41349ea7693cff81b1170d33712b820292
Signed-off-by: Carlos Galo <carlosgalo@google.com>

5 files changed, 180 insertions, 39 deletions

diff --git a/include/liblmkd_utils.h b/include/liblmkd_utils.h
index a84db84..1ca5d89 100644
--- a/include/liblmkd_utils.h
+++ b/include/liblmkd_utils.h
@@ -71,6 +71,20 @@ enum update_props_result lmkd_update_props(int sock);
  */
 int create_memcg(uid_t uid, pid_t pid);
 
+enum boot_completed_notification_result {
+    BOOT_COMPLETED_NOTIF_SUCCESS,
+    BOOT_COMPLETED_NOTIF_FAILS,
+    BOOT_COMPLETED_NOTIF_ALREADY_HANDLED,
+    BOOT_COMPLETED_NOTIF_SEND_ERR,
+    BOOT_COMPLETED_NOTIF_RECV_ERR,
+    BOOT_COMPLETED_NOTIF_FORMAT_ERR,
+};
+
+/*
+ * Notify LMKD the device has finished booting up.
+ */
+enum boot_completed_notification_result lmkd_notify_boot_completed(int sock);
+
 __END_DECLS
 
 #endif /* _LIBLMKD_UTILS_H_ */
diff --git a/include/lmkd.h b/include/lmkd.h
index aca5c34..c57b244 100644
--- a/include/lmkd.h
+++ b/include/lmkd.h
@@ -37,6 +37,7 @@ enum lmk_cmd {
     LMK_UPDATE_PROPS,       /* Reinit properties */
     LMK_STAT_KILL_OCCURRED, /* Unsolicited msg to subscribed clients on proc kills for statsd log */
     LMK_START_MONITORING,   /* Start psi monitoring if it was skipped earlier */
+    LMK_BOOT_COMPLETED,     /* Notify LMKD boot is completed */
 };
 
 /*
@@ -281,7 +282,7 @@ static inline size_t lmkd_pack_set_update_props_repl(LMKD_CTRL_PACKET packet, in
     return 2 * sizeof(int);
 }
 
-/* LMK_PROCPRIO reply payload */
+/* LMK_UPDATE_PROPS reply payload */
 struct lmk_update_props_reply {
     int result;
 };
@@ -295,6 +296,39 @@ static inline void lmkd_pack_get_update_props_repl(LMKD_CTRL_PACKET packet,
     params->result = ntohl(packet[1]);
 }
 
+/*
+ * Prepare LMK_BOOT_COMPLETED packet and return packet size in bytes.
+ * Warning: no checks performed, caller should ensure valid parameters.
+ */
+static inline size_t lmkd_pack_set_boot_completed_notif(LMKD_CTRL_PACKET packet) {
+    packet[0] = htonl(LMK_BOOT_COMPLETED);
+    return sizeof(int);
+}
+
+/*
+ * Prepare LMK_BOOT_COMPLETED reply packet and return packet size in bytes.
+ * Warning: no checks performed, caller should ensure valid parameters.
+ */
+static inline size_t lmkd_pack_set_boot_completed_notif_repl(LMKD_CTRL_PACKET packet, int result) {
+    packet[0] = htonl(LMK_BOOT_COMPLETED);
+    packet[1] = htonl(result);
+    return 2 * sizeof(int);
+}
+
+/* LMK_BOOT_COMPLETED reply payload */
+struct lmk_boot_completed_notif_reply {
+    int result;
+};
+
+/*
+ * For LMK_BOOT_COMPLETED reply payload.
+ * Warning: no checks performed, caller should ensure valid parameters.
+ */
+static inline void lmkd_pack_get_boot_completed_notif_repl(
+        LMKD_CTRL_PACKET packet, struct lmk_boot_completed_notif_reply* params) {
+    params->result = ntohl(packet[1]);
+}
+
 __END_DECLS
 
 #endif /* _LMKD_H_ */
diff --git a/liblmkd_utils.cpp b/liblmkd_utils.cpp
index e5e99de..f517004 100644
--- a/liblmkd_utils.cpp
+++ b/liblmkd_utils.cpp
@@ -78,6 +78,46 @@ enum update_props_result lmkd_update_props(int sock) {
     return params.result == 0 ? UPDATE_PROPS_SUCCESS : UPDATE_PROPS_FAIL;
 }
 
+enum boot_completed_notification_result lmkd_notify_boot_completed(int sock) {
+    LMKD_CTRL_PACKET packet;
+    int size;
+
+    size = lmkd_pack_set_boot_completed_notif(packet);
+    if (TEMP_FAILURE_RETRY(write(sock, packet, size)) < 0) {
+        return BOOT_COMPLETED_NOTIF_SEND_ERR;
+    }
+
+    size = TEMP_FAILURE_RETRY(read(sock, packet, CTRL_PACKET_MAX_SIZE));
+    if (size < 0) {
+        return BOOT_COMPLETED_NOTIF_RECV_ERR;
+    }
+
+    if (size != 2 * sizeof(int) || lmkd_pack_get_cmd(packet) != LMK_BOOT_COMPLETED) {
+        return BOOT_COMPLETED_NOTIF_FORMAT_ERR;
+    }
+
+    struct lmk_boot_completed_notif_reply params;
+    lmkd_pack_get_boot_completed_notif_repl(packet, &params);
+
+    enum boot_completed_notification_result res;
+    switch (params.result) {
+        case -1:
+            res = BOOT_COMPLETED_NOTIF_FAILS;
+            break;
+        case 0:
+            res = BOOT_COMPLETED_NOTIF_SUCCESS;
+            break;
+        case 1:
+            res = BOOT_COMPLETED_NOTIF_ALREADY_HANDLED;
+            break;
+        default:
+            /* This should never be reached */
+            res = BOOT_COMPLETED_NOTIF_FAILS;
+    }
+
+    return res;
+}
+
 int create_memcg(uid_t uid, pid_t pid) {
     return createProcessGroup(uid, pid, true) == 0 ? 0 : -1;
 }
diff --git a/lmkd.cpp b/lmkd.cpp
index 30b4338..40cc617 100644
--- a/lmkd.cpp
+++ b/lmkd.cpp
@@ -205,6 +205,7 @@ static bool pidfd_supported;
 static int last_kill_pid_or_fd = -1;
 static struct timespec last_kill_tm;
 static bool monitors_initialized;
+static bool boot_completed_handled = false;
 
 /* lmkd configurable parameters */
 static bool debug_process_killing;
@@ -569,6 +570,7 @@ static long page_k; /* page size in kB */
 static bool update_props();
 static bool init_monitors();
 static void destroy_monitors();
+static bool init_direct_reclaim_monitoring();
 
 static int clamp(int low, int high, int value) {
     return std::max(std::min(value, high), low);
@@ -1537,6 +1539,11 @@ static void ctrl_command_handler(int dsock_idx) {
             } else {
                 result = 0;
             }
+
+            if (direct_reclaim_threshold_ms > 0 && !memevent_listener) {
+                ALOGW("Kernel support for direct_reclaim_threshold_ms is not found");
+                direct_reclaim_threshold_ms = 0;
+            }
         }
 
         len = lmkd_pack_set_update_props_repl(packet, result);
@@ -1569,6 +1576,38 @@ static void ctrl_command_handler(int dsock_idx) {
         }
         ALOGI("Initialized monitors after boot completed.");
         break;
+    case LMK_BOOT_COMPLETED:
+        if (nargs != 0) goto wronglen;
+
+        if (boot_completed_handled) {
+            /* Notify we have already handled post boot-up operations */
+            result = 1;
+        } else if (!property_get_bool("sys.boot_completed", false)) {
+            ALOGE("LMK_BOOT_COMPLETED cannot be handled before boot completed");
+            result = -1;
+        } else {
+            /*
+             * Initialize the memevent listener after boot is completed to prevent
+             * waiting, during boot-up, for BPF programs to be loaded.
+             */
+            if (init_direct_reclaim_monitoring()) {
+                ALOGI("Using memevents for direct reclaim detection");
+            } else {
+                ALOGI("Using vmstats for direct reclaim detection");
+                if (direct_reclaim_threshold_ms > 0) {
+                    ALOGW("Kernel support for direct_reclaim_threshold_ms is not found");
+                    direct_reclaim_threshold_ms = 0;
+                }
+            }
+            result = 0;
+            boot_completed_handled = true;
+        }
+
+        len = lmkd_pack_set_boot_completed_notif_repl(packet, result);
+        if (ctrl_data_write(dsock_idx, (char*)packet, len) != len) {
+            ALOGE("Failed to report boot-completed operation results");
+        }
+        break;
     default:
         ALOGE("Received unknown command code %d", cmd);
         return;
@@ -3294,12 +3333,12 @@ static void direct_reclaim_state_change(int data __unused, uint32_t events __unu
 static bool init_direct_reclaim_monitoring() {
     static struct event_handler_info direct_reclaim_poll_hinfo = {0, direct_reclaim_state_change};
 
-    if (!memevent_listener) {
-        // Make sure bpf programs are loaded
-        android::bpf::waitForProgsLoaded();
-        memevent_listener = std::make_unique<android::bpf::memevents::MemEventListener>(
-                android::bpf::memevents::MemEventClient::LMKD);
-    }
+    if (memevent_listener) return true;
+
+    // Make sure bpf programs are loaded, else we'll wait until they are loaded
+    android::bpf::waitForProgsLoaded();
+    memevent_listener = std::make_unique<android::bpf::memevents::MemEventListener>(
+            android::bpf::memevents::MemEventClient::LMKD);
 
     if (!memevent_listener->ok()) {
         ALOGE("Failed to initialize memevents listener");
@@ -3326,10 +3365,6 @@ static bool init_direct_reclaim_monitoring() {
     epev.data.ptr = (void*)&direct_reclaim_poll_hinfo;
     if (epoll_ctl(epollfd, EPOLL_CTL_ADD, memevent_listener_fd, &epev) < 0) {
         ALOGE("Failed registering direct reclaim fd: %d; errno=%d", memevent_listener_fd, errno);
-        /*
-         * Reset the fd to let `destroy_direct_reclaim_monitoring` know we failed adding this fd,
-         * therefore it won't try to close the `memevent_listener_fd`.
-         */
         memevent_listener.reset();
         return false;
     }
@@ -3341,19 +3376,6 @@ static bool init_direct_reclaim_monitoring() {
     return true;
 }
 
-static void destroy_direct_reclaim_monitoring() {
-    if (!memevent_listener) return;
-
-    if (epoll_ctl(epollfd, EPOLL_CTL_DEL, memevent_listener->getRingBufferFd(), NULL) < 0) {
-        ALOGE("Failed to unregister direct reclaim monitoring; errno=%d", errno);
-    }
-
-    maxevents--;
-    memevent_listener.reset();
-    direct_reclaim_start_tm.tv_sec = 0;
-    direct_reclaim_start_tm.tv_nsec = 0;
-}
-
 static bool init_psi_monitors() {
     /*
      * When PSI is used on low-ram devices or on high-end devices without memfree levels
@@ -3505,16 +3527,6 @@ static bool init_monitors() {
         ALOGI("Using vmpressure for memory pressure detection");
     }
 
-    if (init_direct_reclaim_monitoring()) {
-        ALOGI("Using memevents for direct reclaim detection");
-    } else {
-        ALOGI("Using vmstats for direct reclaim detection");
-        if (direct_reclaim_threshold_ms > 0) {
-            ALOGW("Kernel support for direct_reclaim_threshold_ms is not found");
-            direct_reclaim_threshold_ms = 0;
-        }
-    }
-
     monitors_initialized = true;
     return true;
 }
@@ -3529,7 +3541,6 @@ static void destroy_monitors() {
         destroy_mp_common(VMPRESS_LEVEL_MEDIUM);
         destroy_mp_common(VMPRESS_LEVEL_LOW);
     }
-    destroy_direct_reclaim_monitoring();
 }
 
 static void drop_reaper_comm() {
@@ -3885,6 +3896,41 @@ int issue_reinit() {
     return res == UPDATE_PROPS_SUCCESS ? 0 : -1;
 }
 
+static int on_boot_completed() {
+    int sock;
+
+    sock = lmkd_connect();
+    if (sock < 0) {
+        ALOGE("failed to connect to lmkd: %s", strerror(errno));
+        return -1;
+    }
+
+    enum boot_completed_notification_result res = lmkd_notify_boot_completed(sock);
+
+    switch (res) {
+        case BOOT_COMPLETED_NOTIF_SUCCESS:
+            break;
+        case BOOT_COMPLETED_NOTIF_ALREADY_HANDLED:
+            ALOGW("lmkd already handled boot-completed operations");
+            break;
+        case BOOT_COMPLETED_NOTIF_SEND_ERR:
+            ALOGE("failed to send lmkd request: %m");
+            break;
+        case BOOT_COMPLETED_NOTIF_RECV_ERR:
+            ALOGE("failed to receive request: %m");
+            break;
+        case BOOT_COMPLETED_NOTIF_FORMAT_ERR:
+            ALOGE("lmkd reply is invalid");
+            break;
+        case BOOT_COMPLETED_NOTIF_FAILS:
+            ALOGE("lmkd failed to receive boot-completed notification");
+            break;
+    }
+
+    close(sock);
+    return res == BOOT_COMPLETED_NOTIF_SUCCESS ? 0 : -1;
+}
+
 static bool update_props() {
     /* By default disable low level vmpressure events */
     level_oomadj[VMPRESS_LEVEL_LOW] =
@@ -3945,11 +3991,15 @@ static bool update_props() {
 }
 
 int main(int argc, char **argv) {
-    if ((argc > 1) && argv[1] && !strcmp(argv[1], "--reinit")) {
-        if (property_set(LMKD_REINIT_PROP, "")) {
-            ALOGE("Failed to reset " LMKD_REINIT_PROP " property");
+    if ((argc > 1) && argv[1]) {
+        if (!strcmp(argv[1], "--reinit")) {
+            if (property_set(LMKD_REINIT_PROP, "")) {
+                ALOGE("Failed to reset " LMKD_REINIT_PROP " property");
+            }
+            return issue_reinit();
+        } else if (!strcmp(argv[1], "--boot_completed")) {
+            return on_boot_completed();
         }
-        return issue_reinit();
     }
 
     if (!update_props()) {
diff --git a/lmkd.rc b/lmkd.rc
index ffe0bc6..73ed79c 100644
--- a/lmkd.rc
+++ b/lmkd.rc
@@ -10,6 +10,9 @@ service lmkd /system/bin/lmkd
 on property:lmkd.reinit=1
     exec_background /system/bin/lmkd --reinit
 
+on property:sys.boot_completed=1
+    exec_background /system/bin/lmkd --boot_completed
+
 # reinitialize lmkd after device finished booting if experiments set any flags during boot
 on property:sys.boot_completed=1 && property:lmkd.reinit=0
     setprop lmkd.reinit 1