kernel: support timed umtx waits through shared address deadlines

Co-authored-by: BlackCAT <52720489+psnwd@users.noreply.github.com>

src/kernel/syncOnAddress.cpp

@@ -47,11 +47,17 @@ struct WaitDeadline {
 	Clock::time_point end {};
 };
 
+[[nodiscard]] WaitDeadline MakeDeadline(std::chrono::nanoseconds timeout) {
+	const auto now = Clock::now();
+	const auto remaining = Clock::time_point::max() - now;
+	return {true, timeout >= remaining ? Clock::time_point::max() : now + timeout};
+}
+
 [[nodiscard]] WaitDeadline MakeDeadline(const uint32_t* timeout_micros) {
 	if (timeout_micros == nullptr) {
 		return {};
 	}
-	return {true, Clock::now() + std::chrono::microseconds(*timeout_micros)};
+	return MakeDeadline(std::chrono::microseconds(*timeout_micros));
 }
 
 [[nodiscard]] uint32_t GetWaitSliceMicros(const WaitDeadline& deadline, bool first_wait) {
@@ -65,7 +71,7 @@ struct WaitDeadline {
 	}
 
 	const auto remaining =
-	    std::chrono::duration_cast<std::chrono::microseconds>(deadline.end - now).count();
+	    std::chrono::ceil<std::chrono::microseconds>(deadline.end - now).count();
 	return static_cast<uint32_t>(std::min<int64_t>(remaining, SIGNAL_POLL_MICROS));
 }
 
@@ -78,9 +84,8 @@ void PollSignals(signal_poll_func_t signal_poll) {
 #if KYTY_PLATFORM == KYTY_PLATFORM_LINUX && !defined(__APPLE__)
 
 template <typename T>
-int WaitLinux(volatile T* address, T expected, const uint32_t* timeout_micros,
+int WaitLinux(volatile T* address, T expected, const WaitDeadline& deadline,
               signal_poll_func_t signal_poll) {
-	const auto deadline   = MakeDeadline(timeout_micros);
 	bool       first_wait = true;
 
 	for (;;) {
@@ -185,11 +190,10 @@ void UnregisterPortableWaiter(volatile void*                               addre
 }
 
 template <typename T>
-int WaitPortable(volatile T* address, T expected, const uint32_t* timeout_micros,
+int WaitPortable(volatile T* address, T expected, const WaitDeadline& deadline,
                  signal_poll_func_t signal_poll) {
 	PortableWaiter waiter;
 	auto           entry      = RegisterPortableWaiter(address, &waiter);
-	const auto     deadline   = MakeDeadline(timeout_micros);
 	bool           first_wait = true;
 	int            result     = OK;
 
@@ -250,7 +254,7 @@ int WakePortable(volatile void* address, int32_t count) {
 }
 
 template <typename T>
-int WaitImpl(volatile T* address, T expected, const uint32_t* timeout_micros,
+int WaitImpl(volatile T* address, T expected, const WaitDeadline& deadline,
              signal_poll_func_t signal_poll) {
 	if (!IsValidWaitAddress(address)) {
 		return KERNEL_ERROR_EINVAL;
@@ -258,9 +262,9 @@ int WaitImpl(volatile T* address, T expected, const uint32_t* timeout_micros,
 
 	int result = OK;
 #if KYTY_PLATFORM == KYTY_PLATFORM_LINUX && !defined(__APPLE__)
-	result = WaitLinux(address, expected, timeout_micros, signal_poll);
+	result = WaitLinux(address, expected, deadline, signal_poll);
 #else
-	result = WaitPortable(address, expected, timeout_micros, signal_poll);
+	result = WaitPortable(address, expected, deadline, signal_poll);
 #endif
 	PollSignals(signal_poll);
 	return result;
@@ -270,12 +274,17 @@ int WaitImpl(volatile T* address, T expected, const uint32_t* timeout_micros,
 
 int Wait32(volatile uint32_t* address, uint32_t expected, const uint32_t* timeout_micros,
            signal_poll_func_t signal_poll) {
-	return WaitImpl(address, expected, timeout_micros, signal_poll);
+	return WaitImpl(address, expected, MakeDeadline(timeout_micros), signal_poll);
 }
 
 int Wait64(volatile uint64_t* address, uint64_t expected, const uint32_t* timeout_micros,
            signal_poll_func_t signal_poll) {
-	return WaitImpl(address, expected, timeout_micros, signal_poll);
+	return WaitImpl(address, expected, MakeDeadline(timeout_micros), signal_poll);
+}
+
+int Wait64(volatile uint64_t* address, uint64_t expected, std::chrono::nanoseconds timeout,
+           signal_poll_func_t signal_poll) {
+	return WaitImpl(address, expected, MakeDeadline(timeout), signal_poll);
 }
 
 int Wake(volatile void* address, int32_t count) {

src/kernel/syncOnAddress.h

@@ -4,6 +4,8 @@
 #include "common/abi.h"
 #include "common/common.h"
 
+#include <chrono>
+
 namespace Libs::LibKernel::SyncOnAddress {
 
 using signal_poll_func_t = void (*)();
@@ -12,6 +14,8 @@ int Wait32(volatile uint32_t* address, uint32_t expected, const uint32_t* timeou
            signal_poll_func_t signal_poll = nullptr);
 int Wait64(volatile uint64_t* address, uint64_t expected, const uint32_t* timeout_micros,
            signal_poll_func_t signal_poll = nullptr);
+int Wait64(volatile uint64_t* address, uint64_t expected, std::chrono::nanoseconds timeout,
+           signal_poll_func_t signal_poll = nullptr);
 int Wake(volatile void* address, int32_t count);
 
 } // namespace Libs::LibKernel::SyncOnAddress

src/libs/libKernel.cpp

@@ -2128,7 +2128,13 @@ int KYTY_SYSV_ABI UmtxOp(volatile void* address, int operation, uint64_t value,
 					*GetErrorAddr() = POSIX_EINVAL;
 					return -1;
 				}
-				EXIT("Valid _umtx_op timed waits are not implemented\n");
+				const auto max_ns = std::chrono::nanoseconds::max().count();
+				const auto timeout_ns = duration.tv_sec > (max_ns - duration.tv_nsec) / 1000000000
+				                            ? max_ns
+				                            : duration.tv_sec * 1000000000 + duration.tv_nsec;
+				return POSIX_CALL(LibKernel::SyncOnAddress::Wait64(
+				    static_cast<volatile uint64_t*>(address), value, std::chrono::nanoseconds(timeout_ns),
+				    LibKernel::KernelDispatchPendingSignalForCurrentThread));
 			}
 			return POSIX_CALL(LibKernel::SyncOnAddress::Wait64(
 			    static_cast<volatile uint64_t*>(address), value, nullptr,

tests/SyncOnAddressTests.cpp

@@ -110,6 +110,38 @@ void TestValueChangeAndWake() {
   Check(result == OK, "a value change plus wake releases the waiter");
 }
 
+void TestNanosecondTimeout() {
+  uint64_t word = 0;
+  const auto start = std::chrono::steady_clock::now();
+  Check(Wait64(&word, 0, std::chrono::milliseconds(1)) ==
+            Libs::LibKernel::KERNEL_ERROR_ETIMEDOUT,
+        "64-bit timespec wait expires");
+  Check(std::chrono::steady_clock::now() - start >= std::chrono::milliseconds(1),
+        "timespec wait preserves its deadline");
+  Check(Wait64(&word, 0, std::chrono::nanoseconds(0)) ==
+            Libs::LibKernel::KERNEL_ERROR_ETIMEDOUT,
+        "zero timespec polls a matching value");
+  Check(Wait64(&word, 1, std::chrono::nanoseconds(0)) == OK,
+        "zero timespec succeeds on mismatch");
+
+  g_signal_poll_count.store(0, std::memory_order_relaxed);
+  std::atomic<bool> returned{false};
+  std::thread waker([&] {
+    while (g_signal_poll_count.load(std::memory_order_relaxed) < 2 &&
+           !returned.load(std::memory_order_acquire)) {
+      std::this_thread::yield();
+    }
+    Store(&word, UINT64_C(0x100000000));
+    Check(Wake(&word, 1) == OK, "wake releases a long timespec wait");
+  });
+  const int result = Wait64(&word, 0, std::chrono::nanoseconds::max(), CountSignalPoll);
+  returned.store(true, std::memory_order_release);
+  waker.join();
+  Check(result == OK, "long timespec wait does not overflow its deadline");
+  Check(g_signal_poll_count.load(std::memory_order_relaxed) > 1,
+        "timespec wait continues polling guest signals");
+}
+
 void TestWakeOneThenAll() {
   constexpr int WAITER_COUNT = 4;
   uint32_t word = 0;
@@ -260,6 +292,7 @@ int main() {
   TestMismatchReturnsImmediately();
   TestTimeout();
   TestValueChangeAndWake();
+  TestNanosecondTimeout();
   TestWakeOneThenAll();
   TestAddressesAreIsolated();
   TestCompareRegisterWakeRace();