@samitouri / QOSAMI-WSL-1 / commits / 0b402cd8

Use RegisterWaitForSingleObject() in MultiHandleIOWait (#40658)

* Save state * Add test coverage * Cleanup for PR * Remove stale command * Apply PR feedback * Create explicit move ctor

Blue committed May 28, 2026 at 12:48 UTC 0b402cd8ff7510d4bc5ba1b1f0e57811d929e471

4 files changed +261 -75

src/windows/common/HandleIO.cpp

+105 -72

@@ -61,6 +61,14 @@ void CancelPendingIo(auto Handle, OVERLAPPED& Overlapped)

61 }

62 }

64 +inline void UnregisterWait(HANDLE waitHandle) noexcept

65 +{

66 + // INVALID_HANDLE_VALUE makes UnregisterWaitEx block until any in-flight wait callback returns.

67 + LOG_LAST_ERROR_IF(!UnregisterWaitEx(waitHandle, INVALID_HANDLE_VALUE));

68 +}

69 +

70 +using unique_registered_wait = wil::unique_any_handle_null<decltype(&UnregisterWait), &UnregisterWait>;

71 +

72 } // namespace

74 // HandleWrapper

@@ -914,9 +922,38 @@ void DockerIORelayHandle::OnRead(const gsl::span<char>& Buffer)

922

923 // MultiHandleWait

924

925 +MultiHandleWait::MultiHandleWait(MultiHandleWait&& other) noexcept

926 +{

927 + *this = std::move(other);

928 +}

929 +

930 +MultiHandleWait& MultiHandleWait::operator=(MultiHandleWait&& other) noexcept

931 +{

932 + if (this != &other)

933 + {

934 + m_handles = std::move(other.m_handles);

935 + m_handleSignaledEvent = std::move(other.m_handleSignaledEvent);

936 + m_cancel = other.m_cancel;

937 +

938 + for (auto& entry : m_handles)

939 + {

940 + entry->self = this;

941 + }

942 +

943 + // N.B. moving a MultiHandleWait() while running is not supported

944 + WI_ASSERT(m_signaledHandles.empty());

945 + }

946 +

947 + return *this;

948 +}

949 +

950 void MultiHandleWait::AddHandle(std::unique_ptr<OverlappedIOHandle>&& handle, Flags flags)

951 {

919 - m_handles.emplace_back(flags, std::move(handle));

952 + auto entry = std::make_unique<Entry>();

953 + entry->HandleFlags = flags;

954 + entry->Handle = std::move(handle);

955 + entry->self = this;

956 + m_handles.emplace_back(std::move(entry));

957 }

958

959 void MultiHandleWait::Cancel()

@@ -924,123 +961,119 @@ void MultiHandleWait::Cancel()

961 m_cancel = true;

962 }

963

964 +void NTAPI MultiHandleWait::WaitCallback(PVOID Context, BOOLEAN /*TimerOrWaitFired*/)

965 +{

966 + auto* entry = static_cast<Entry*>(Context);

967 +

968 + entry->self->m_signaledHandles.push(entry);

969 + entry->self->m_handleSignaledEvent.SetEvent();

970 +}

971 +

972 bool MultiHandleWait::Run(std::optional<std::chrono::milliseconds> Timeout)

973 {

974 m_cancel = false; // Run may be called multiple times.

975

976 std::optional<std::chrono::steady_clock::time_point> deadline;

932 -

977 if (Timeout.has_value())

978 {

979 deadline = std::chrono::steady_clock::now() + Timeout.value();

980 }

981

938 - // Run until all handles are completed.

982 + std::vector<unique_registered_wait> callbacks;

983

940 - while (!m_handles.empty() && !m_cancel)

984 + while (!m_cancel)

985 {

942 - // Schedule IO on each handle until all are either pending, or completed.

943 - for (size_t i = 0; i < m_handles.size() && !m_cancel; i++)

986 + // Cancel any pending callback.

987 + callbacks.clear();

988 +

989 + Entry* signaledEntry = nullptr;

990 + while (m_signaledHandles.try_pop(signaledEntry))

991 + {

992 + try

993 + {

994 + signaledEntry->Handle->Collect();

995 + }

996 + catch (...)

997 + {

998 + if (WI_IsFlagSet(signaledEntry->HandleFlags, Flags::IgnoreErrors))

999 + {

1000 + signaledEntry->Handle.reset();

1001 + continue;

1002 + }

1003 +

1004 + throw;

1005 + }

1006 + }

1007 +

1008 + m_handleSignaledEvent.ResetEvent();

1009 +

1010 + bool hasHandleToWaitFor = false;

1011 + for (auto it = m_handles.begin(); it != m_handles.end();)

1012 {

945 - while (m_handles[i].second->GetState() == IOHandleStatus::Standby && !m_cancel)

1013 + auto& entry = **it;

1014 +

1015 + while (entry.Handle && entry.Handle->GetState() == IOHandleStatus::Standby && !m_cancel)

1016 {

1017 try

1018 {

949 - m_handles[i].second->Schedule();

1019 + entry.Handle->Schedule();

1020 }

1021 catch (...)

1022 {

953 - if (WI_IsFlagSet(m_handles[i].first, Flags::IgnoreErrors))

1023 + if (WI_IsFlagSet(entry.HandleFlags, Flags::IgnoreErrors))

1024 {

955 - m_handles[i].second.reset(); // Reset the handle so it can be deleted.

1025 + entry.Handle.reset();

1026 break;

1027 }

958 - else

959 - {

960 - throw;

961 - }

1028 +

1029 + throw;

1030 }

1031 }

964 - }

1032

966 - // Remove completed handles from m_handles.

967 - bool hasHandleToWaitFor = false;

968 - for (auto it = m_handles.begin(); it != m_handles.end();)

969 - {

970 - if (!it->second)

971 - {

972 - it = m_handles.erase(it);

973 - }

974 - else if (it->second->GetState() == IOHandleStatus::Completed)

1033 + if (!entry.Handle || entry.Handle->GetState() == IOHandleStatus::Completed)

1034 {

976 - if (WI_IsFlagSet(it->first, Flags::CancelOnCompleted))

1035 + if (entry.Handle && WI_IsFlagSet(entry.HandleFlags, Flags::CancelOnCompleted))

1036 {

978 - m_cancel = true; // Cancel the IO if a handle with CancelOnCompleted is in the completed state.

1037 + m_cancel = true;

1038 }

1039

1040 it = m_handles.erase(it);

1041 + continue;

1042 }

983 - else

1043 +

1044 + auto& callback = callbacks.emplace_back();

1045 +

1046 + THROW_IF_WIN32_BOOL_FALSE(RegisterWaitForSingleObject(

1047 + &callback, entry.Handle->GetHandle(), &WaitCallback, &entry, INFINITE, WT_EXECUTEINWAITTHREAD | WT_EXECUTEONLYONCE));

1048 +

1049 + if (WI_IsFlagClear(entry.HandleFlags, Flags::NeedNotComplete))

1050 {

985 - // If only NeedNotComplete handles are left, we want to exit Run.

986 - if (WI_IsFlagClear(it->first, Flags::NeedNotComplete))

987 - {

988 - hasHandleToWaitFor = true;

989 - }

990 - ++it;

1051 + hasHandleToWaitFor = true;

1052 }

992 - }

1053

994 - if (!hasHandleToWaitFor || m_cancel)

995 - {

996 - break;

1054 + ++it;

1055 }

1056

999 - // Wait for the next operation to complete.

1000 - std::vector<HANDLE> waitHandles;

1001 - for (const auto& e : m_handles)

1057 + if (m_handles.empty() || !hasHandleToWaitFor || m_cancel)

1058 {

1003 - waitHandles.emplace_back(e.second->GetHandle());

1059 + break;

1060 }

1061

1062 DWORD waitTimeout = INFINITE;

1063 if (deadline.has_value())

1064 {

1009 - auto miliseconds =

1065 + auto milliseconds =

1066 std::chrono::duration_cast<std::chrono::milliseconds>(deadline.value() - std::chrono::steady_clock::now()).count();

1067

1012 - waitTimeout = static_cast<DWORD>(std::max(0LL, miliseconds));

1068 + waitTimeout = static_cast<DWORD>(std::max<long long>(0, milliseconds));

1069 }

1070

1015 - auto result = WaitForMultipleObjects(static_cast<DWORD>(waitHandles.size()), waitHandles.data(), false, waitTimeout);

1016 - if (result == WAIT_TIMEOUT)

1017 - {

1018 - THROW_WIN32(ERROR_TIMEOUT);

1019 - }

1020 - else if (result >= WAIT_OBJECT_0 && result < WAIT_OBJECT_0 + m_handles.size())

1021 - {

1022 - auto index = result - WAIT_OBJECT_0;

1023 -

1024 - try

1025 - {

1026 - m_handles[index].second->Collect();

1027 - }

1028 - catch (...)

1029 - {

1030 - if (WI_IsFlagSet(m_handles[index].first, Flags::IgnoreErrors))

1031 - {

1032 - m_handles.erase(m_handles.begin() + index);

1033 - }

1034 - else

1035 - {

1036 - throw;

1037 - }

1038 - }

1039 - }

1040 - else

1041 - {

1042 - THROW_LAST_ERROR_MSG("Timeout: %lu, Count: %llu", waitTimeout, waitHandles.size());

1043 - }

1071 + THROW_HR_IF_MSG(

1072 + HRESULT_FROM_WIN32(ERROR_TIMEOUT),

1073 + !m_handleSignaledEvent.wait(waitTimeout),

1074 + "Timed out waiting for %llu handles. Timeout: %lu",

1075 + m_handles.size(),

1076 + waitTimeout);

1077 }

1078

1079 return !m_cancel;

src/windows/common/HandleIO.h

+17 -3

@@ -2,6 +2,8 @@

3 #pragma once

5 +#include <concurrent_queue.h>

6 +

7 #define LX_RELAY_BUFFER_SIZE 0x1000

9 namespace wsl::windows::common::io {

@@ -349,12 +351,10 @@ private:

351 WriteHandle* ActiveHandle = nullptr;

352 size_t RemainingBytes = 0;

353 };

352 -

354 class MultiHandleWait

355 {

356 public:

357 NON_COPYABLE(MultiHandleWait);

357 - DEFAULT_MOVABLE(MultiHandleWait);

358

359 enum Flags

360 {

@@ -365,13 +365,27 @@ public:

365 };

366

367 MultiHandleWait() = default;

368 + MultiHandleWait(MultiHandleWait&&) noexcept;

369 + MultiHandleWait& operator=(MultiHandleWait&&) noexcept;

370

371 void AddHandle(std::unique_ptr<OverlappedIOHandle>&& handle, Flags flags = Flags::None);

372 bool Run(std::optional<std::chrono::milliseconds> Timeout);

373 void Cancel();

374

375 private:

374 - std::vector<std::pair<Flags, std::unique_ptr<OverlappedIOHandle>>> m_handles;

376 + struct Entry

377 + {

378 + Flags HandleFlags{};

379 + std::unique_ptr<OverlappedIOHandle> Handle;

380 + MultiHandleWait* self;

381 + };

382 +

383 + static void NTAPI WaitCallback(PVOID Context, BOOLEAN TimerOrWaitFired);

384 +

385 + concurrency::concurrent_queue<Entry*> m_signaledHandles;

386 + wil::unique_event m_handleSignaledEvent{wil::EventOptions::ManualReset};

387 +

388 + std::vector<std::unique_ptr<Entry>> m_handles;

389 bool m_cancel = false;

390 };

391

test/windows/UnitTests.cpp

+88

@@ -6946,6 +6946,94 @@ Error code: Wsl/InstallDistro/WSL_E_INVALID_JSON\r\n",

6946 }

6947 }

6948

6949 + TEST_METHOD(MultiHandleWaitAboveMaximumWaitObjects)

6950 + {

6951 + // Validate that MultiHandleWait can wait on more than MAXIMUM_WAIT_OBJECTS (64) handles.

6952 + constexpr size_t handleCount = 100;

6953 + static_assert(handleCount > MAXIMUM_WAIT_OBJECTS);

6954 +

6955 + // Scenario 1: signal every event before Run(); all callbacks must fire and Run() must return.

6956 + {

6957 + std::vector<wil::unique_event> events;

6958 + events.reserve(handleCount);

6959 + for (size_t i = 0; i < handleCount; ++i)

6960 + {

6961 + events.emplace_back(wil::EventOptions::ManualReset);

6962 + }

6963 +

6964 + std::vector<bool> fired(handleCount, false);

6965 + std::atomic<size_t> firedCount{0};

6966 + std::mutex firedLock;

6967 +

6968 + wsl::windows::common::io::MultiHandleWait io;

6969 + for (size_t i = 0; i < handleCount; ++i)

6970 + {

6971 + io.AddHandle(std::make_unique<wsl::windows::common::io::EventHandle>(

6972 + wsl::windows::common::io::HandleWrapper{events[i].get()}, [&fired, &firedCount, &firedLock, i]() {

6973 + std::lock_guard lock{firedLock};

6974 + VERIFY_IS_FALSE(fired[i]);

6975 + fired[i] = true;

6976 + firedCount.fetch_add(1);

6977 + }));

6978 + }

6979 +

6980 + for (auto& e : events)

6981 + {

6982 + e.SetEvent();

6983 + }

6984 +

6985 + VERIFY_IS_TRUE(io.Run(std::chrono::seconds(60)));

6986 + VERIFY_ARE_EQUAL(firedCount.load(), handleCount);

6987 + for (size_t i = 0; i < handleCount; ++i)

6988 + {

6989 + VERIFY_IS_TRUE(fired[i]);

6990 + }

6991 + }

6992 +

6993 + // Scenario 2: signal events one at a time from another thread while Run() processes them.

6994 + {

6995 + std::vector<wil::unique_event> events;

6996 + events.reserve(handleCount);

6997 + for (size_t i = 0; i < handleCount; ++i)

6998 + {

6999 + events.emplace_back(wil::EventOptions::ManualReset);

7000 + }

7001 +

7002 + std::vector<bool> fired(handleCount, false);

7003 + std::atomic<size_t> firedCount{0};

7004 + std::mutex firedLock;

7005 +

7006 + wsl::windows::common::io::MultiHandleWait io;

7007 + for (size_t i = 0; i < handleCount; ++i)

7008 + {

7009 + io.AddHandle(std::make_unique<wsl::windows::common::io::EventHandle>(

7010 + wsl::windows::common::io::HandleWrapper{events[i].get()}, [&fired, &firedCount, &firedLock, i]() {

7011 + std::lock_guard lock{firedLock};

7012 + VERIFY_IS_FALSE(fired[i]);

7013 + fired[i] = true;

7014 + firedCount.fetch_add(1);

7015 + }));

7016 + }

7017 +

7018 + std::thread signaller([&events]() {

7019 + for (auto& e : events)

7020 + {

7021 + e.SetEvent();

7022 + std::this_thread::sleep_for(std::chrono::milliseconds(1));

7023 + }

7024 + });

7025 +

7026 + VERIFY_IS_TRUE(io.Run(std::chrono::seconds(60)));

7027 + signaller.join();

7028 +

7029 + VERIFY_ARE_EQUAL(firedCount.load(), handleCount);

7030 + for (size_t i = 0; i < handleCount; ++i)

7031 + {

7032 + VERIFY_IS_TRUE(fired[i]);

7033 + }

7034 + }

7035 + }

7036 +

7037 TEST_METHOD(SocketChannel)

7038 {

7039 // Read exactly `size` bytes from a raw socket into the destination buffer.

test/windows/WSLCTests.cpp

+51

@@ -8333,6 +8333,57 @@ class WSLCTests

8333 }

8334 }

8335

8336 + WSLC_TEST_METHOD(ContainerLogsManyConcurrentFollowers)

8337 + {

8338 + constexpr size_t followerCount = 100;

8339 + static_assert(followerCount > MAXIMUM_WAIT_OBJECTS);

8340 +

8341 + WSLCContainerLauncher launcher("debian:latest", "logs-test-many-followers", {"/bin/cat"}, {}, {}, WSLCProcessFlagsStdin);

8342 + auto container = launcher.Launch(*m_defaultSession);

8343 + auto initProcess = container.GetInitProcess();

8344 +

8345 + auto containerStdin = initProcess.GetStdHandle(0);

8346 + VERIFY_WIN32_BOOL_SUCCEEDED(WriteFile(containerStdin.get(), "OK\n", 3, nullptr, nullptr));

8347 +

8348 + std::atomic<size_t> readersReady{0};

8349 + std::atomic<size_t> readersSucceeded{0};

8350 + std::vector<std::thread> threads;

8351 + threads.reserve(followerCount);

8352 +

8353 + for (size_t i = 0; i < followerCount; ++i)

8354 + {

8355 + threads.emplace_back([&]() {

8356 + try

8357 + {

8358 + COMOutputHandle stdoutHandle{};

8359 + COMOutputHandle stderrHandle{};

8360 + VERIFY_SUCCEEDED(container.Get().Logs(WSLCLogsFlagsFollow, &stdoutHandle, &stderrHandle, 0, 0, 0));

8361 +

8362 + PartialHandleRead reader(stdoutHandle.Get());

8363 + reader.Expect("OK\n");

8364 + readersReady.fetch_add(1);

8365 + reader.ExpectClosed();

8366 + readersSucceeded.fetch_add(1);

8367 + }

8368 + CATCH_LOG();

8369 + });

8370 + }

8371 +

8372 + // Wait until every follower has observed the marker before killing the container.

8373 + wsl::shared::retry::RetryWithTimeout<void>(

8374 + [&]() { THROW_HR_IF(E_ABORT, readersReady.load() < followerCount); }, std::chrono::milliseconds(100), std::chrono::seconds(120));

8375 +

8376 + // Kill the container so all follow handles are closed.

8377 + VERIFY_SUCCEEDED(container.Get().Stop(WSLCSignalSIGKILL, 0));

8378 +

8379 + for (auto& t : threads)

8380 + {

8381 + t.join();

8382 + }

8383 +

8384 + VERIFY_ARE_EQUAL(readersSucceeded.load(), followerCount);

8385 + }

8386 +

8387 WSLC_TEST_METHOD(ContainerLabels)

8388 {

8389 // Docker labels do not have a size limit, so test with a very large label value to validate that the API can handle it.