Integration PROCS_PRIO cmd in lmkd

Creating new cmd in LMKD to batch PROCPRIO requests.

Test: atest lmkd_tests
Bug: 325525024
Change-Id: I5460446d4e968e80263aa25298e2a893863eece4
Signed-off-by: Carlos Galo <carlosgalo@google.com>

include/liblmkd_utils.h

@@ -39,6 +39,14 @@ int lmkd_connect();
  */
 int lmkd_register_proc(int sock, struct lmk_procprio *params);
 
+/*
+ * Registers a batch of processes with lmkd and sets its oomadj score.
+ * On success returns 0.
+ * On error, -1 is returned.
+ * In the case of error errno is set appropriately.
+ */
+int lmkd_register_procs(int sock, struct lmk_procs_prio* params, const int proc_count);
+
 /*
  * Unregisters a process previously registered with lmkd.
  * On success returns 0.

include/lmkd.h

@@ -38,6 +38,7 @@ enum lmk_cmd {
     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 */
+    LMK_PROCS_PRIO,         /* Register processes and set the same oom_adj_score */
 };
 
 /*
@@ -108,6 +109,8 @@ struct lmk_procprio {
     int oomadj;
     enum proc_type ptype;
 };
+#define LMK_PROCPRIO_FIELD_COUNT 4
+#define LMK_PROCPRIO_SIZE (LMK_PROCPRIO_FIELD_COUNT * sizeof(int))
 
 /*
  * For LMK_PROCPRIO packet get its payload.
@@ -324,6 +327,54 @@ static inline void lmkd_pack_get_boot_completed_notif_repl(
     params->result = ntohl(packet[1]);
 }
 
+#define PROCS_PRIO_MAX_RECORD_COUNT (CTRL_PACKET_MAX_SIZE / LMK_PROCPRIO_SIZE)
+
+struct lmk_procs_prio {
+    struct lmk_procprio procs[PROCS_PRIO_MAX_RECORD_COUNT];
+};
+
+/*
+ * For LMK_PROCS_PRIO packet get its payload.
+ * Warning: no checks performed, caller should ensure valid parameters.
+ */
+static inline int lmkd_pack_get_procs_prio(LMKD_CTRL_PACKET packet, struct lmk_procs_prio* params,
+                                           const int field_count) {
+    if (field_count < LMK_PROCPRIO_FIELD_COUNT || (field_count % LMK_PROCPRIO_FIELD_COUNT) != 0)
+        return -1;
+    const int procs_count = (field_count / LMK_PROCPRIO_FIELD_COUNT);
+
+    /* Start packet at 1 since 0 is cmd type */
+    int packetIdx = 1;
+    for (int procs_idx = 0; procs_idx < procs_count; procs_idx++) {
+        params->procs[procs_idx].pid = (pid_t)ntohl(packet[packetIdx++]);
+        params->procs[procs_idx].uid = (uid_t)ntohl(packet[packetIdx++]);
+        params->procs[procs_idx].oomadj = ntohl(packet[packetIdx++]);
+        params->procs[procs_idx].ptype = (enum proc_type)ntohl(packet[packetIdx++]);
+    }
+
+    return procs_count;
+}
+
+/*
+ * Prepare LMK_PROCS_PRIO packet and return packet size in bytes.
+ * Warning: no checks performed, caller should ensure valid parameters.
+ */
+static inline size_t lmkd_pack_set_procs_prio(LMKD_CTRL_PACKET packet,
+                                              struct lmk_procs_prio* params,
+                                              const int procs_count) {
+    packet[0] = htonl(LMK_PROCS_PRIO);
+    int packetIdx = 1;
+
+    for (int i = 0; i < procs_count; i++) {
+        packet[packetIdx++] = htonl(params->procs[i].pid);
+        packet[packetIdx++] = htonl(params->procs[i].uid);
+        packet[packetIdx++] = htonl(params->procs[i].oomadj);
+        packet[packetIdx++] = htonl((int)params->procs[i].ptype);
+    }
+
+    return packetIdx * sizeof(int);
+}
+
 __END_DECLS
 
 #endif /* _LMKD_H_ */

liblmkd_utils.cpp

@@ -43,6 +43,17 @@ int lmkd_register_proc(int sock, struct lmk_procprio *params) {
     return (ret < 0) ? -1 : 0;
 }
 
+int lmkd_register_procs(int sock, struct lmk_procs_prio* params, const int proc_count) {
+    LMKD_CTRL_PACKET packet;
+    size_t size;
+    int ret;
+
+    size = lmkd_pack_set_procs_prio(packet, params, proc_count);
+    ret = TEMP_FAILURE_RETRY(write(sock, packet, size));
+
+    return (ret < 0) ? -1 : 0;
+}
+
 int lmkd_unregister_proc(int sock, struct lmk_procremove *params) {
     LMKD_CTRL_PACKET packet;
     size_t size;

lmkd.cpp

@@ -1479,6 +1479,20 @@ static void cmd_target(int ntargets, LMKD_CTRL_PACKET packet) {
     }
 }
 
+static void cmd_procs_prio(LMKD_CTRL_PACKET packet, const int field_count, struct ucred* cred) {
+    struct lmk_procs_prio params;
+
+    const int procs_count = lmkd_pack_get_procs_prio(packet, &params, field_count);
+    if (procs_count < 0) {
+        ALOGE("LMK_PROCS_PRIO received invalid packet format");
+        return;
+    }
+
+    for (int i = 0; i < procs_count; i++) {
+        apply_proc_prio(params.procs[i], cred);
+    }
+}
+
 static void ctrl_command_handler(int dsock_idx) {
     LMKD_CTRL_PACKET packet;
     struct ucred cred;
@@ -1627,6 +1641,9 @@ static void ctrl_command_handler(int dsock_idx) {
             ALOGE("Failed to report boot-completed operation results");
         }
         break;
+    case LMK_PROCS_PRIO:
+        cmd_procs_prio(packet, nargs, &cred);
+        break;
     default:
         ALOGE("Received unknown command code %d", cmd);
         return;

tests/lmkd_tests.cpp

@@ -26,6 +26,7 @@
 #include <liblmkd_utils.h>
 #include <log/log_properties.h>
 #include <private/android_filesystem_config.h>
+#include <stdlib.h>
 
 using namespace android::base;
 
@@ -113,6 +114,16 @@ class LmkdTest : public ::testing::Test {
         }
     }
 
+    void SendProcsPrioRequest(struct lmk_procs_prio procs_prio_request, int procs_count) {
+        ASSERT_FALSE(lmkd_register_procs(sock, &procs_prio_request, procs_count) < 0)
+                << "Failed to communicate with lmkd, err=" << strerror(errno);
+    }
+
+    void SendGetKillCountRequest(struct lmk_getkillcnt* get_kill_cnt_request) {
+        ASSERT_GE(lmkd_get_kill_count(sock, get_kill_cnt_request), 0)
+                << "Failed fetching lmkd kill count";
+    }
+
     static std::string ExecCommand(const std::string& command) {
         FILE* fp = popen(command.c_str(), "r");
         std::string content;
@@ -170,6 +181,8 @@ class LmkdTest : public ::testing::Test {
                reap_pid == pid;
     }
 
+    uid_t getLmkdTestUid() const { return uid; }
+
   private:
     int sock;
     uid_t uid;
@@ -239,6 +252,88 @@ TEST_F(LmkdTest, TargetReaping) {
    }
 }
 
+/*
+ * Verify that the `PROCS_PRIO` cmd is able to receive a batch of processes and adjust their
+ * those processes' OOM score.
+ */
+TEST_F(LmkdTest, batch_procs_oom_score_adj) {
+    struct ChildProcessInfo {
+        pid_t pid;
+        int original_oom_score;
+        int req_new_oom_score;
+    };
+
+    struct ChildProcessInfo children_info[PROCS_PRIO_MAX_RECORD_COUNT];
+
+    for (unsigned int i = 0; i < PROCS_PRIO_MAX_RECORD_COUNT; i++) {
+        children_info[i].pid = fork();
+        if (children_info[i].pid < 0) {
+            for (const auto child : children_info)
+                if (child.pid >= 0) kill(child.pid, SIGKILL);
+            FAIL() << "Failed forking process in iteration=" << i;
+        } else if (children_info[i].pid == 0) {
+            /*
+             * Keep the children alive, the parent process will kill it
+             * once we are done with it.
+             */
+            while (true) {
+                sleep(20);
+            }
+        }
+    }
+
+    struct lmk_procs_prio procs_prio_request;
+    const uid_t parent_uid = getLmkdTestUid();
+
+    for (unsigned int i = 0; i < PROCS_PRIO_MAX_RECORD_COUNT; i++) {
+        if (children_info[i].pid < 0) continue;
+
+        const std::string process_oom_path =
+                "proc/" + std::to_string(children_info[i].pid) + "/oom_score_adj";
+        std::string curr_oom_score;
+        if (!ReadFileToString(process_oom_path, &curr_oom_score) || curr_oom_score.empty()) {
+            for (const auto child : children_info)
+                if (child.pid >= 0) kill(child.pid, SIGKILL);
+            FAIL() << "Failed reading original oom score for child process: "
+                   << children_info[i].pid;
+        }
+
+        children_info[i].original_oom_score = atoi(curr_oom_score.c_str());
+        children_info[i].req_new_oom_score =
+                ((unsigned int)children_info[i].original_oom_score != i) ? i : (i + 10);
+        procs_prio_request.procs[i] = {.pid = children_info[i].pid,
+                                       .uid = parent_uid,
+                                       .oomadj = children_info[i].req_new_oom_score,
+                                       .ptype = proc_type::PROC_TYPE_APP};
+    }
+
+    /*
+     * Submit batching, then send a new/different request and wait for LMKD
+     * to respond to it. This ensures that LMKD has finished the batching
+     * request and we can now read/validate the new OOM scores.
+     */
+    SendProcsPrioRequest(procs_prio_request, PROCS_PRIO_MAX_RECORD_COUNT);
+    struct lmk_getkillcnt kill_cnt_req = {.min_oomadj = -1000, .max_oomadj = 1000};
+    SendGetKillCountRequest(&kill_cnt_req);
+
+    for (auto child_info : children_info) {
+        if (child_info.pid < 0) continue;
+        const std::string process_oom_path =
+                "proc/" + std::to_string(child_info.pid) + "/oom_score_adj";
+        std::string curr_oom_score;
+        if (!ReadFileToString(process_oom_path, &curr_oom_score) || curr_oom_score.empty()) {
+            for (const auto child : children_info)
+                if (child.pid >= 0) kill(child.pid, SIGKILL);
+            FAIL() << "Failed reading new oom score for child process: " << child_info.pid;
+        }
+        kill(child_info.pid, SIGKILL);
+
+        const int actual_new_oom_score = atoi(curr_oom_score.c_str());
+        ASSERT_EQ(child_info.req_new_oom_score, actual_new_oom_score)
+                << "Child with pid=" << child_info.pid << " didn't update its OOM score";
+    }
+}
+
 int main(int argc, char** argv) {
     ::testing::InitGoogleTest(&argc, argv);
     InitLogging(argv, StderrLogger);