When rm_interest needs to re-arm other interests, add EPOLLONESHOT

Author: beef9999

io/epoll.cpp

@@ -156,18 +156,20 @@ class EventEngineEPoll : public MasterEventEngine, public CascadingEventEngine,
     virtual int rm_interest(Event e) override {
         if (e.fd < 0 || (size_t)e.fd >= _inflight_events.size())
             LOG_ERROR_RETURN(EINVAL, -1, "invalid file descriptor ", e.fd);
-        if (unlikely(!e.interests)) return 0;
+        if (unlikely(e.interests == 0)) return 0;
         auto& entry = _inflight_events[e.fd];
         auto intersection = e.interests & entry.interests & EVENT_RWE;
         if (intersection == 0) return 0;
 
         int ret, op = 0;    // ^ is to flip intersected bits
         auto x = (entry.interests ^ intersection) & EVENT_RWE;
         if (likely(e.interests & ONE_SHOT)) {
+            // If e is ONE_SHOT, the entry must be ONE_SHOT as well
             if (!x) {
                 ret = 0; // no need to epoll_ctl()
             } else {
                 auto events = evmap.translate_bitwisely(x);
+                events |= EPOLLONESHOT;
                 ret = ctl(e.fd, EPOLL_CTL_MOD, events); // re-arm other interests
             }
         } else {
@@ -286,7 +288,7 @@ class EventEngineEPoll : public MasterEventEngine, public CascadingEventEngine,
             LOG_ERROR_RETURN(EINVAL, -1, "invalid fd");
         if (interest & (interest-1))
             LOG_ERROR_RETURN(EINVAL, -1, "can not wait for multiple interests");
-        if (unlikely(!interest))
+        if (unlikely(interest == 0))
             return rm_interest({fd, EVENT_RWE, 0}); // remove fd from epoll
         int ret = add_interest({fd, interest | ONE_SHOT, CURRENT});
         if (ret < 0) LOG_ERROR_RETURN(0, -1, "failed to add event interest");

io/iouring-wrapper.cpp

@@ -224,6 +224,8 @@ class iouringEngine : public MasterEventEngine, public CascadingEventEngine, pub
     }
 
     int wait_for_fd(int fd, uint32_t interests, uint64_t timeout) override {
+        if (unlikely(interests == 0))
+            return 0;
         unsigned poll_mask = evmap.translate_bitwisely(interests);
         // The io_uring_prep_poll_add's return value is the same as poll(2)'s revents.
         int ret = async_io(&io_uring_prep_poll_add, timeout, 0, fd, poll_mask);

io/kqueue.cpp

@@ -88,6 +88,8 @@ class KQueue : public MasterEventEngine, public CascadingEventEngine, public Res
     }
 
     int wait_for_fd(int fd, uint32_t interests, uint64_t timeout) override {
+        if (unlikely(interests == 0))
+            return 0;
         short ev = (interests == EVENT_READ) ? EVFILT_READ : EVFILT_WRITE;
         enqueue(fd, ev, EV_ADD | EV_ONESHOT, 0, CURRENT);
         int ret = thread_usleep(timeout);

net/kernel_socket.cpp

@@ -135,6 +135,7 @@ class KernelSocketStream : public SocketStreamBase {
         return (Object*) (uint64_t) fd;
     }
     int close() final {
+        get_vcpu()->master_event_engine->wait_for_fd(fd, 0, -1UL);
         auto ret = ::close(fd);
         fd = -1;
         return ret;