Use RegisterWaitForSingleObject() in MultiHandleIOWait (#40658)

* Save state

* Add test coverage

* Cleanup for PR

* Remove stale command

* Apply PR feedback

* Create explicit move ctor

Author: OneBlue

src/windows/common/HandleIO.cpp

@@ -61,6 +61,14 @@ void CancelPendingIo(auto Handle, OVERLAPPED& Overlapped)
     }
 }
 
+inline void UnregisterWait(HANDLE waitHandle) noexcept
+{
+    // INVALID_HANDLE_VALUE makes UnregisterWaitEx block until any in-flight wait callback returns.
+    LOG_LAST_ERROR_IF(!UnregisterWaitEx(waitHandle, INVALID_HANDLE_VALUE));
+}
+
+using unique_registered_wait = wil::unique_any_handle_null<decltype(&UnregisterWait), &UnregisterWait>;
+
 } // namespace
 
 // HandleWrapper
@@ -914,133 +922,158 @@ void DockerIORelayHandle::OnRead(const gsl::span<char>& Buffer)
 
 // MultiHandleWait
 
+MultiHandleWait::MultiHandleWait(MultiHandleWait&& other) noexcept
+{
+    *this = std::move(other);
+}
+
+MultiHandleWait& MultiHandleWait::operator=(MultiHandleWait&& other) noexcept
+{
+    if (this != &other)
+    {
+        m_handles = std::move(other.m_handles);
+        m_handleSignaledEvent = std::move(other.m_handleSignaledEvent);
+        m_cancel = other.m_cancel;
+
+        for (auto& entry : m_handles)
+        {
+            entry->self = this;
+        }
+
+        // N.B. moving a MultiHandleWait() while running is not supported
+        WI_ASSERT(m_signaledHandles.empty());
+    }
+
+    return *this;
+}
+
 void MultiHandleWait::AddHandle(std::unique_ptr<OverlappedIOHandle>&& handle, Flags flags)
 {
-    m_handles.emplace_back(flags, std::move(handle));
+    auto entry = std::make_unique<Entry>();
+    entry->HandleFlags = flags;
+    entry->Handle = std::move(handle);
+    entry->self = this;
+    m_handles.emplace_back(std::move(entry));
 }
 
 void MultiHandleWait::Cancel()
 {
     m_cancel = true;
 }
 
+void NTAPI MultiHandleWait::WaitCallback(PVOID Context, BOOLEAN /*TimerOrWaitFired*/)
+{
+    auto* entry = static_cast<Entry*>(Context);
+
+    entry->self->m_signaledHandles.push(entry);
+    entry->self->m_handleSignaledEvent.SetEvent();
+}
+
 bool MultiHandleWait::Run(std::optional<std::chrono::milliseconds> Timeout)
 {
     m_cancel = false; // Run may be called multiple times.
 
     std::optional<std::chrono::steady_clock::time_point> deadline;
-
     if (Timeout.has_value())
     {
         deadline = std::chrono::steady_clock::now() + Timeout.value();
     }
 
-    // Run until all handles are completed.
+    std::vector<unique_registered_wait> callbacks;
 
-    while (!m_handles.empty() && !m_cancel)
+    while (!m_cancel)
     {
-        // Schedule IO on each handle until all are either pending, or completed.
-        for (size_t i = 0; i < m_handles.size() && !m_cancel; i++)
+        // Cancel any pending callback.
+        callbacks.clear();
+
+        Entry* signaledEntry = nullptr;
+        while (m_signaledHandles.try_pop(signaledEntry))
+        {
+            try
+            {
+                signaledEntry->Handle->Collect();
+            }
+            catch (...)
+            {
+                if (WI_IsFlagSet(signaledEntry->HandleFlags, Flags::IgnoreErrors))
+                {
+                    signaledEntry->Handle.reset();
+                    continue;
+                }
+
+                throw;
+            }
+        }
+
+        m_handleSignaledEvent.ResetEvent();
+
+        bool hasHandleToWaitFor = false;
+        for (auto it = m_handles.begin(); it != m_handles.end();)
         {
-            while (m_handles[i].second->GetState() == IOHandleStatus::Standby && !m_cancel)
+            auto& entry = **it;
+
+            while (entry.Handle && entry.Handle->GetState() == IOHandleStatus::Standby && !m_cancel)
             {
                 try
                 {
-                    m_handles[i].second->Schedule();
+                    entry.Handle->Schedule();
                 }
                 catch (...)
                 {
-                    if (WI_IsFlagSet(m_handles[i].first, Flags::IgnoreErrors))
+                    if (WI_IsFlagSet(entry.HandleFlags, Flags::IgnoreErrors))
                     {
-                        m_handles[i].second.reset(); // Reset the handle so it can be deleted.
+                        entry.Handle.reset();
                         break;
                     }
-                    else
-                    {
-                        throw;
-                    }
+
+                    throw;
                 }
             }
-        }
 
-        // Remove completed handles from m_handles.
-        bool hasHandleToWaitFor = false;
-        for (auto it = m_handles.begin(); it != m_handles.end();)
-        {
-            if (!it->second)
-            {
-                it = m_handles.erase(it);
-            }
-            else if (it->second->GetState() == IOHandleStatus::Completed)
+            if (!entry.Handle || entry.Handle->GetState() == IOHandleStatus::Completed)
             {
-                if (WI_IsFlagSet(it->first, Flags::CancelOnCompleted))
+                if (entry.Handle && WI_IsFlagSet(entry.HandleFlags, Flags::CancelOnCompleted))
                 {
-                    m_cancel = true; // Cancel the IO if a handle with CancelOnCompleted is in the completed state.
+                    m_cancel = true;
                 }
 
                 it = m_handles.erase(it);
+                continue;
             }
-            else
+
+            auto& callback = callbacks.emplace_back();
+
+            THROW_IF_WIN32_BOOL_FALSE(RegisterWaitForSingleObject(
+                &callback, entry.Handle->GetHandle(), &WaitCallback, &entry, INFINITE, WT_EXECUTEINWAITTHREAD | WT_EXECUTEONLYONCE));
+
+            if (WI_IsFlagClear(entry.HandleFlags, Flags::NeedNotComplete))
             {
-                // If only NeedNotComplete handles are left, we want to exit Run.
-                if (WI_IsFlagClear(it->first, Flags::NeedNotComplete))
-                {
-                    hasHandleToWaitFor = true;
-                }
-                ++it;
+                hasHandleToWaitFor = true;
             }
-        }
 
-        if (!hasHandleToWaitFor || m_cancel)
-        {
-            break;
+            ++it;
         }
 
-        // Wait for the next operation to complete.
-        std::vector<HANDLE> waitHandles;
-        for (const auto& e : m_handles)
+        if (m_handles.empty() || !hasHandleToWaitFor || m_cancel)
         {
-            waitHandles.emplace_back(e.second->GetHandle());
+            break;
         }
 
         DWORD waitTimeout = INFINITE;
         if (deadline.has_value())
         {
-            auto miliseconds =
+            auto milliseconds =
                 std::chrono::duration_cast<std::chrono::milliseconds>(deadline.value() - std::chrono::steady_clock::now()).count();
 
-            waitTimeout = static_cast<DWORD>(std::max(0LL, miliseconds));
+            waitTimeout = static_cast<DWORD>(std::max<long long>(0, milliseconds));
         }
 
-        auto result = WaitForMultipleObjects(static_cast<DWORD>(waitHandles.size()), waitHandles.data(), false, waitTimeout);
-        if (result == WAIT_TIMEOUT)
-        {
-            THROW_WIN32(ERROR_TIMEOUT);
-        }
-        else if (result >= WAIT_OBJECT_0 && result < WAIT_OBJECT_0 + m_handles.size())
-        {
-            auto index = result - WAIT_OBJECT_0;
-
-            try
-            {
-                m_handles[index].second->Collect();
-            }
-            catch (...)
-            {
-                if (WI_IsFlagSet(m_handles[index].first, Flags::IgnoreErrors))
-                {
-                    m_handles.erase(m_handles.begin() + index);
-                }
-                else
-                {
-                    throw;
-                }
-            }
-        }
-        else
-        {
-            THROW_LAST_ERROR_MSG("Timeout: %lu, Count: %llu", waitTimeout, waitHandles.size());
-        }
+        THROW_HR_IF_MSG(
+            HRESULT_FROM_WIN32(ERROR_TIMEOUT),
+            !m_handleSignaledEvent.wait(waitTimeout),
+            "Timed out waiting for %llu handles. Timeout: %lu",
+            m_handles.size(),
+            waitTimeout);
     }
 
     return !m_cancel;

src/windows/common/HandleIO.h

@@ -2,6 +2,8 @@
 
 #pragma once
 
+#include <concurrent_queue.h>
+
 #define LX_RELAY_BUFFER_SIZE 0x1000
 
 namespace wsl::windows::common::io {
@@ -349,12 +351,10 @@ class DockerIORelayHandle : public OverlappedIOHandle
     WriteHandle* ActiveHandle = nullptr;
     size_t RemainingBytes = 0;
 };
-
 class MultiHandleWait
 {
 public:
     NON_COPYABLE(MultiHandleWait);
-    DEFAULT_MOVABLE(MultiHandleWait);
 
     enum Flags
     {
@@ -365,13 +365,27 @@ class MultiHandleWait
     };
 
     MultiHandleWait() = default;
+    MultiHandleWait(MultiHandleWait&&) noexcept;
+    MultiHandleWait& operator=(MultiHandleWait&&) noexcept;
 
     void AddHandle(std::unique_ptr<OverlappedIOHandle>&& handle, Flags flags = Flags::None);
     bool Run(std::optional<std::chrono::milliseconds> Timeout);
     void Cancel();
 
 private:
-    std::vector<std::pair<Flags, std::unique_ptr<OverlappedIOHandle>>> m_handles;
+    struct Entry
+    {
+        Flags HandleFlags{};
+        std::unique_ptr<OverlappedIOHandle> Handle;
+        MultiHandleWait* self;
+    };
+
+    static void NTAPI WaitCallback(PVOID Context, BOOLEAN TimerOrWaitFired);
+
+    concurrency::concurrent_queue<Entry*> m_signaledHandles;
+    wil::unique_event m_handleSignaledEvent{wil::EventOptions::ManualReset};
+
+    std::vector<std::unique_ptr<Entry>> m_handles;
     bool m_cancel = false;
 };
 

test/windows/UnitTests.cpp

@@ -6946,6 +6946,94 @@ Error code: Wsl/InstallDistro/WSL_E_INVALID_JSON\r\n",
         }
     }
 
+    TEST_METHOD(MultiHandleWaitAboveMaximumWaitObjects)
+    {
+        // Validate that MultiHandleWait can wait on more than MAXIMUM_WAIT_OBJECTS (64) handles.
+        constexpr size_t handleCount = 100;
+        static_assert(handleCount > MAXIMUM_WAIT_OBJECTS);
+
+        // Scenario 1: signal every event before Run(); all callbacks must fire and Run() must return.
+        {
+            std::vector<wil::unique_event> events;
+            events.reserve(handleCount);
+            for (size_t i = 0; i < handleCount; ++i)
+            {
+                events.emplace_back(wil::EventOptions::ManualReset);
+            }
+
+            std::vector<bool> fired(handleCount, false);
+            std::atomic<size_t> firedCount{0};
+            std::mutex firedLock;
+
+            wsl::windows::common::io::MultiHandleWait io;
+            for (size_t i = 0; i < handleCount; ++i)
+            {
+                io.AddHandle(std::make_unique<wsl::windows::common::io::EventHandle>(
+                    wsl::windows::common::io::HandleWrapper{events[i].get()}, [&fired, &firedCount, &firedLock, i]() {
+                        std::lock_guard lock{firedLock};
+                        VERIFY_IS_FALSE(fired[i]);
+                        fired[i] = true;
+                        firedCount.fetch_add(1);
+                    }));
+            }
+
+            for (auto& e : events)
+            {
+                e.SetEvent();
+            }
+
+            VERIFY_IS_TRUE(io.Run(std::chrono::seconds(60)));
+            VERIFY_ARE_EQUAL(firedCount.load(), handleCount);
+            for (size_t i = 0; i < handleCount; ++i)
+            {
+                VERIFY_IS_TRUE(fired[i]);
+            }
+        }
+
+        // Scenario 2: signal events one at a time from another thread while Run() processes them.
+        {
+            std::vector<wil::unique_event> events;
+            events.reserve(handleCount);
+            for (size_t i = 0; i < handleCount; ++i)
+            {
+                events.emplace_back(wil::EventOptions::ManualReset);
+            }
+
+            std::vector<bool> fired(handleCount, false);
+            std::atomic<size_t> firedCount{0};
+            std::mutex firedLock;
+
+            wsl::windows::common::io::MultiHandleWait io;
+            for (size_t i = 0; i < handleCount; ++i)
+            {
+                io.AddHandle(std::make_unique<wsl::windows::common::io::EventHandle>(
+                    wsl::windows::common::io::HandleWrapper{events[i].get()}, [&fired, &firedCount, &firedLock, i]() {
+                        std::lock_guard lock{firedLock};
+                        VERIFY_IS_FALSE(fired[i]);
+                        fired[i] = true;
+                        firedCount.fetch_add(1);
+                    }));
+            }
+
+            std::thread signaller([&events]() {
+                for (auto& e : events)
+                {
+                    e.SetEvent();
+                    std::this_thread::sleep_for(std::chrono::milliseconds(1));
+                }
+            });
+
+            VERIFY_IS_TRUE(io.Run(std::chrono::seconds(60)));
+            signaller.join();
+
+            VERIFY_ARE_EQUAL(firedCount.load(), handleCount);
+            for (size_t i = 0; i < handleCount; ++i)
+            {
+                VERIFY_IS_TRUE(fired[i]);
+            }
+        }
+    }
+
     TEST_METHOD(SocketChannel)
     {
         // Read exactly `size` bytes from a raw socket into the destination buffer.