4 files changed, 78 insertions, 30 deletions

diff --git a/src/core/hle/kernel/svc.cpp b/src/core/hle/kernel/svc.cpp
index 74643f598..73793955a 100644
--- a/src/core/hle/kernel/svc.cpp
+++ b/src/core/hle/kernel/svc.cpp
@@ -120,17 +120,19 @@ static ResultCode GetProcessId(u32* process_id, Handle process_handle) {
 }
 
 /// Default thread wakeup callback for WaitSynchronization
-static void DefaultThreadWakeupCallback(ThreadWakeupReason reason, SharedPtr<Thread> thread,
-                                        SharedPtr<WaitObject> object) {
+static bool DefaultThreadWakeupCallback(ThreadWakeupReason reason, SharedPtr<Thread> thread,
+                                        SharedPtr<WaitObject> object, size_t index) {
     ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ANY);
 
     if (reason == ThreadWakeupReason::Timeout) {
         thread->SetWaitSynchronizationResult(RESULT_TIMEOUT);
-        return;
+        return true;
     }
 
     ASSERT(reason == ThreadWakeupReason::Signal);
     thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+
+    return true;
 };
 
 /// Wait for a kernel object to synchronize, timeout after the specified nanoseconds
@@ -499,20 +501,44 @@ static ResultCode WaitProcessWideKeyAtomic(VAddr mutex_addr, VAddr semaphore_add
     ASSERT(semaphore->available_count == 0);
     ASSERT(semaphore->mutex_addr == mutex_addr);
 
-    CASCADE_CODE(WaitSynchronization1(
-        semaphore, thread.get(), nano_seconds,
-        [mutex](ThreadWakeupReason reason, SharedPtr<Thread> thread, SharedPtr<WaitObject> object) {
-            ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ANY);
+    auto wakeup_callback = [mutex, nano_seconds](ThreadWakeupReason reason,
+                                                 SharedPtr<Thread> thread,
+                                                 SharedPtr<WaitObject> object, size_t index) {
+        ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ANY);
+
+        if (reason == ThreadWakeupReason::Timeout) {
+            thread->SetWaitSynchronizationResult(RESULT_TIMEOUT);
+            return true;
+        }
+
+        ASSERT(reason == ThreadWakeupReason::Signal);
+
+        // Now try to acquire the mutex and don't resume if it's not available.
+        if (!mutex->ShouldWait(thread.get())) {
+            mutex->Acquire(thread.get());
+            thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+            return true;
+        }
 
-            if (reason == ThreadWakeupReason::Timeout) {
-                thread->SetWaitSynchronizationResult(RESULT_TIMEOUT);
-                return;
-            }
+        if (nano_seconds == 0) {
+            thread->SetWaitSynchronizationResult(RESULT_TIMEOUT);
+            return true;
+        }
+
+        thread->wait_objects = {mutex};
+        mutex->AddWaitingThread(thread);
+        thread->status = THREADSTATUS_WAIT_SYNCH_ANY;
+
+        // Create an event to wake the thread up after the
+        // specified nanosecond delay has passed
+        thread->WakeAfterDelay(nano_seconds);
+        thread->wakeup_callback = DefaultThreadWakeupCallback;
+
+        Core::System::GetInstance().PrepareReschedule();
 
-            ASSERT(reason == ThreadWakeupReason::Signal);
-            thread->SetWaitSynchronizationResult(WaitSynchronization1(mutex, thread.get()));
-            thread->SetWaitSynchronizationOutput(thread->GetWaitObjectIndex(object.get()));
-        }));
+        return false;
+    };
+    CASCADE_CODE(WaitSynchronization1(semaphore, thread.get(), nano_seconds, wakeup_callback));
 
     mutex->Release(thread.get());
 
diff --git a/src/core/hle/kernel/thread.cpp b/src/core/hle/kernel/thread.cpp
index f9d821a80..111c496b9 100644
--- a/src/core/hle/kernel/thread.cpp
+++ b/src/core/hle/kernel/thread.cpp
@@ -244,20 +244,23 @@ static void ThreadWakeupCallback(u64 thread_handle, int cycles_late) {
         return;
     }
 
+    bool resume = true;
+
     if (thread->status == THREADSTATUS_WAIT_SYNCH_ANY ||
         thread->status == THREADSTATUS_WAIT_SYNCH_ALL || thread->status == THREADSTATUS_WAIT_ARB) {
 
-        // Invoke the wakeup callback before clearing the wait objects
-        if (thread->wakeup_callback)
-            thread->wakeup_callback(ThreadWakeupReason::Timeout, thread, nullptr);
-
         // Remove the thread from each of its waiting objects' waitlists
         for (auto& object : thread->wait_objects)
             object->RemoveWaitingThread(thread.get());
         thread->wait_objects.clear();
+
+        // Invoke the wakeup callback before clearing the wait objects
+        if (thread->wakeup_callback)
+            resume = thread->wakeup_callback(ThreadWakeupReason::Timeout, thread, nullptr, 0);
     }
 
-    thread->ResumeFromWait();
+    if (resume)
+        thread->ResumeFromWait();
 }
 
 void Thread::WakeAfterDelay(s64 nanoseconds) {
@@ -268,6 +271,10 @@ void Thread::WakeAfterDelay(s64 nanoseconds) {
     CoreTiming::ScheduleEvent(nsToCycles(nanoseconds), ThreadWakeupEventType, callback_handle);
 }
 
+void Thread::CancelWakeupTimer() {
+    CoreTiming::UnscheduleEvent(ThreadWakeupEventType, callback_handle);
+}
+
 void Thread::ResumeFromWait() {
     ASSERT_MSG(wait_objects.empty(), "Thread is waking up while waiting for objects");
 
@@ -444,7 +451,8 @@ ResultVal<SharedPtr<Thread>> Thread::Create(std::string name, VAddr entry_point,
         // Map the page to the current process' address space.
         // TODO(Subv): Find the correct MemoryState for this region.
         vm_manager.MapMemoryBlock(Memory::TLS_AREA_VADDR + available_page * Memory::PAGE_SIZE,
-                                  linheap_memory, offset, Memory::PAGE_SIZE, MemoryState::ThreadLocalStorage);
+                                  linheap_memory, offset, Memory::PAGE_SIZE,
+                                  MemoryState::ThreadLocalStorage);
     }
 
     // Mark the slot as used
@@ -501,7 +509,8 @@ SharedPtr<Thread> SetupMainThread(VAddr entry_point, u32 priority,
     SharedPtr<Thread> thread = std::move(thread_res).Unwrap();
 
     // Register 1 must be a handle to the main thread
-    thread->guest_handle = Kernel::g_handle_table.Create(thread).Unwrap();;
+    thread->guest_handle = Kernel::g_handle_table.Create(thread).Unwrap();
+
     thread->context.cpu_registers[1] = thread->guest_handle;
 
     // Threads by default are dormant, wake up the main thread so it runs when the scheduler fires
@@ -572,4 +581,4 @@ const std::vector<SharedPtr<Thread>>& GetThreadList() {
     return thread_list;
 }
 
-} // namespace
+} // namespace Kernel
diff --git a/src/core/hle/kernel/thread.h b/src/core/hle/kernel/thread.h
index ed44ee933..19ba6e0af 100644
--- a/src/core/hle/kernel/thread.h
+++ b/src/core/hle/kernel/thread.h
@@ -128,6 +128,9 @@ public:
      */
     void WakeAfterDelay(s64 nanoseconds);
 
+    /// Cancel any outstanding wakeup events for this thread
+    void CancelWakeupTimer();
+
     /**
      * Sets the result after the thread awakens (from either WaitSynchronization SVC)
      * @param result Value to set to the returned result
@@ -218,8 +221,8 @@ public:
     /// Handle used as userdata to reference this object when inserting into the CoreTiming queue.
     Handle callback_handle;
 
-    using WakeupCallback = void(ThreadWakeupReason reason, SharedPtr<Thread> thread,
-                                SharedPtr<WaitObject> object);
+    using WakeupCallback = bool(ThreadWakeupReason reason, SharedPtr<Thread> thread,
+                                SharedPtr<WaitObject> object, size_t index);
     // Callback that will be invoked when the thread is resumed from a waiting state. If the thread
     // was waiting via WaitSynchronizationN then the object will be the last object that became
     // available. In case of a timeout, the object will be nullptr.
@@ -237,7 +240,8 @@ private:
  * @param owner_process The parent process for the main thread
  * @return A shared pointer to the main thread
  */
-SharedPtr<Thread> SetupMainThread(VAddr entry_point, u32 priority, SharedPtr<Process> owner_process);
+SharedPtr<Thread> SetupMainThread(VAddr entry_point, u32 priority,
+                                  SharedPtr<Process> owner_process);
 
 /**
  * Returns whether there are any threads that are ready to run.
diff --git a/src/core/hle/kernel/wait_object.cpp b/src/core/hle/kernel/wait_object.cpp
index c942a40fa..ec147b84c 100644
--- a/src/core/hle/kernel/wait_object.cpp
+++ b/src/core/hle/kernel/wait_object.cpp
@@ -68,6 +68,8 @@ SharedPtr<Thread> WaitObject::GetHighestPriorityReadyThread() {
 }
 
 void WaitObject::WakeupWaitingThread(SharedPtr<Thread> thread) {
+    ASSERT(!ShouldWait(thread.get()));
+
     if (!thread)
         return;
 
@@ -75,19 +77,26 @@ void WaitObject::WakeupWaitingThread(SharedPtr<Thread> thread) {
         Acquire(thread.get());
     } else {
         for (auto& object : thread->wait_objects) {
+            ASSERT(!object->ShouldWait(thread.get()));
             object->Acquire(thread.get());
         }
     }
 
-    // Invoke the wakeup callback before clearing the wait objects
-    if (thread->wakeup_callback)
-        thread->wakeup_callback(ThreadWakeupReason::Signal, thread, this);
+    size_t index = thread->GetWaitObjectIndex(this);
 
     for (auto& object : thread->wait_objects)
         object->RemoveWaitingThread(thread.get());
     thread->wait_objects.clear();
 
-    thread->ResumeFromWait();
+    thread->CancelWakeupTimer();
+
+    bool resume = true;
+
+    if (thread->wakeup_callback)
+        resume = thread->wakeup_callback(ThreadWakeupReason::Signal, thread, this, index);
+
+    if (resume)
+        thread->ResumeFromWait();
 }
 
 void WaitObject::WakeupAllWaitingThreads() {