init: use std::function for epoll handling

Also allow unregistering of epoll handlers.

Bug: 64114943
Test: boot
Change-Id: I2abe6a56fd451839931d607dddb91669a7d02ff1

1 files changed, 27 insertions, 7 deletions

diff --git a/init/init.cpp b/init/init.cpp
index 0d5690b07..42032caeb 100644
--- a/init/init.cpp
+++ b/init/init.cpp
@@ -31,6 +31,10 @@
 #include <sys/types.h>
 #include <unistd.h>
 
+#include <map>
+#include <memory>
+#include <optional>
+
 #include <android-base/chrono_utils.h>
 #include <android-base/file.h>
 #include <android-base/logging.h>
@@ -43,9 +47,6 @@
 #include <private/android_filesystem_config.h>
 #include <selinux/android.h>
 
-#include <memory>
-#include <optional>
-
 #include "action_parser.h"
 #include "import_parser.h"
 #include "init_first_stage.h"
@@ -130,12 +131,31 @@ static void LoadBootScripts(ActionManager& action_manager, ServiceList& service_
     }
 }
 
-void register_epoll_handler(int fd, void (*fn)()) {
+static std::map<int, std::function<void()>> epoll_handlers;
+
+void register_epoll_handler(int fd, std::function<void()> handler) {
+    auto[it, inserted] = epoll_handlers.emplace(fd, std::move(handler));
+    if (!inserted) {
+        LOG(ERROR) << "Cannot specify two epoll handlers for a given FD";
+        return;
+    }
     epoll_event ev;
     ev.events = EPOLLIN;
-    ev.data.ptr = reinterpret_cast<void*>(fn);
+    // std::map's iterators do not get invalidated until erased, so we use the pointer to the
+    // std::function in the map directly for epoll_ctl.
+    ev.data.ptr = reinterpret_cast<void*>(&it->second);
     if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, fd, &ev) == -1) {
-        PLOG(ERROR) << "epoll_ctl failed";
+        PLOG(ERROR) << "epoll_ctl failed to add fd";
+        epoll_handlers.erase(fd);
+    }
+}
+
+void unregister_epoll_handler(int fd) {
+    if (epoll_ctl(epoll_fd, EPOLL_CTL_DEL, fd, nullptr) == -1) {
+        PLOG(ERROR) << "epoll_ctl failed to remove fd";
+    }
+    if (epoll_handlers.erase(fd) != 1) {
+        LOG(ERROR) << "Attempting to remove epoll handler for FD without an existing handler";
     }
 }
 
@@ -809,7 +829,7 @@ int main(int argc, char** argv) {
         if (nr == -1) {
             PLOG(ERROR) << "epoll_wait failed";
         } else if (nr == 1) {
-            ((void (*)()) ev.data.ptr)();
+            std::invoke(*reinterpret_cast<std::function<void()>*>(ev.data.ptr));
         }
     }