init/epoll: Make Epoll::Wait() easier to use

Invoke the callback functions from inside Epoll::Wait() instead of
returning a vector with pointers to callback functions. Remove handlers
after handler invocation finished to prevent that self-removal triggers
a use-after-free.

The CL that made Epoll::Wait() return a vector is available at
https://android-review.googlesource.com/c/platform/system/core/+/1112042.

Bug: 213617178
Change-Id: I52c6ade5746a911510746f83802684f2d9cfb429
Signed-off-by: Bart Van Assche <bvanassche@google.com>
diff --git a/init/epoll.cpp b/init/epoll.cpp
index 3a830ce..f814c65 100644
--- a/init/epoll.cpp
+++ b/init/epoll.cpp
@@ -69,9 +69,11 @@
     if (epoll_ctl(epoll_fd_, EPOLL_CTL_DEL, fd, nullptr) == -1) {
         return ErrnoError() << "epoll_ctl failed to remove fd";
     }
-    if (epoll_handlers_.erase(fd) != 1) {
+    auto it = epoll_handlers_.find(fd);
+    if (it == epoll_handlers_.end()) {
         return Error() << "Attempting to remove epoll handler for FD without an existing handler";
     }
+    to_remove_.insert(it->first);
     return {};
 }
 
@@ -79,8 +81,7 @@
     first_callback_ = std::move(first_callback);
 }
 
-Result<std::vector<std::shared_ptr<Epoll::Handler>>> Epoll::Wait(
-        std::optional<std::chrono::milliseconds> timeout) {
+Result<int> Epoll::Wait(std::optional<std::chrono::milliseconds> timeout) {
     int timeout_ms = -1;
     if (timeout && timeout->count() < INT_MAX) {
         timeout_ms = timeout->count();
@@ -94,7 +95,6 @@
     if (num_events > 0 && first_callback_) {
         first_callback_();
     }
-    std::vector<std::shared_ptr<Handler>> pending_functions;
     for (int i = 0; i < num_events; ++i) {
         const auto it = epoll_handlers_.find(ev[i].data.fd);
         if (it == epoll_handlers_.end()) {
@@ -107,10 +107,13 @@
             // Log something informational.
             LOG(ERROR) << "Received unexpected epoll event set: " << ev[i].events;
         }
-        pending_functions.emplace_back(info.handler);
+        (*info.handler)();
+        for (auto fd : to_remove_) {
+            epoll_handlers_.erase(fd);
+        }
+        to_remove_.clear();
     }
-
-    return pending_functions;
+    return num_events;
 }
 
 }  // namespace init