From 646b4f2b46765993d56cc4b7252667acba237ebe Mon Sep 17 00:00:00 2001
From: Ivan Chikish <nekotekina@gmail.com>
Date: Tue, 11 Jul 2023 07:35:55 +0300
Subject: [PATCH] [orbis-kernel] Fix umtx_wait

---
 orbis-kernel/src/sys/sys_umtx.cpp | 13 +++++++++----
 orbis-kernel/src/umtx.cpp         | 20 ++++++++++++++++----
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/orbis-kernel/src/sys/sys_umtx.cpp b/orbis-kernel/src/sys/sys_umtx.cpp
index 9514ea45e..4a7895596 100644
--- a/orbis-kernel/src/sys/sys_umtx.cpp
+++ b/orbis-kernel/src/sys/sys_umtx.cpp
@@ -32,7 +32,7 @@ orbis::SysResult orbis::sys__umtx_op(Thread *thread, ptr<void> obj, sint op,
   ORBIS_LOG_TRACE(__FUNCTION__, obj, op, val, uaddr1, uaddr2);
   if (reinterpret_cast<std::uintptr_t>(obj) - 0x10000 > 0xff'fffe'ffff)
     return ErrorCode::FAULT;
-  auto with_timeout = [&](auto op) -> orbis::ErrorCode {
+  auto with_timeout = [&](auto op, bool loop = true) -> orbis::ErrorCode {
     timespec *ts = nullptr;
     timespec timeout{};
     if (uaddr2 != nullptr) {
@@ -43,19 +43,23 @@ orbis::SysResult orbis::sys__umtx_op(Thread *thread, ptr<void> obj, sint op,
 
       ts = &timeout;
     }
-    __uint128_t usec = timeout.sec;
-    auto start = std::chrono::steady_clock::now();
 
     if (!ts) {
+      if (!loop)
+        return op(-1);
       while (true) {
         if (auto r = op(-1); r != ErrorCode::TIMEDOUT)
           return r;
       }
     } else {
+      __uint128_t usec = timeout.sec;
       usec *= 1000'000;
       usec += (timeout.nsec + 999) / 1000;
       if (usec >= UINT64_MAX)
         usec = -2;
+      if (!loop)
+        return op(usec);
+      auto start = std::chrono::steady_clock::now();
       std::uint64_t udiff = 0;
       while (true) {
         if (auto r = op(usec - udiff); r != ErrorCode::TIMEDOUT)
@@ -77,7 +81,8 @@ orbis::SysResult orbis::sys__umtx_op(Thread *thread, ptr<void> obj, sint op,
     return umtx_unlock_umtx(thread, (ptr<umtx>)obj, val);
   case 2: {
     return with_timeout(
-        [&](std::uint64_t ut) { return umtx_wait(thread, obj, val, ut); });
+        [&](std::uint64_t ut) { return umtx_wait(thread, obj, val, ut); },
+        false);
   }
   case 3:
     return umtx_wake(thread, obj, val);
diff --git a/orbis-kernel/src/umtx.cpp b/orbis-kernel/src/umtx.cpp
index af9ee37cc..46de7796b 100644
--- a/orbis-kernel/src/umtx.cpp
+++ b/orbis-kernel/src/umtx.cpp
@@ -56,11 +56,23 @@ orbis::ErrorCode orbis::umtx_wait(Thread *thread, ptr<void> addr, ulong id,
   auto [chain, key, lock] = g_context.getUmtxChain0(thread->tproc->pid, addr);
   auto node = chain.enqueue(key, thread);
   ErrorCode result = {};
-  // TODO: this is inaccurate with timeout as FreeBsd 9.1 only checks id once
   if (reinterpret_cast<ptr<std::atomic<ulong>>>(addr)->load() == id) {
-    node->second.cv.wait(chain.mtx, ut);
-    if (node->second.thr == thread)
-      result = ErrorCode::TIMEDOUT;
+    if (ut + 1 == 0) {
+      node->second.cv.wait(chain.mtx);
+    } else {
+      auto start = std::chrono::steady_clock::now();
+      std::uint64_t udiff = 0;
+      while (true) {
+        node->second.cv.wait(chain.mtx, ut - udiff);
+        if (node->second.thr != thread)
+          break;
+        udiff = (std::chrono::steady_clock::now() - start).count() / 1000;
+        if (udiff >= ut) {
+          result = ErrorCode::TIMEDOUT;
+          break;
+        }
+      }
+    }
   }
   if (node->second.thr == thread)
     chain.erase(node);