Kernel/Threads: When putting a thread to wait, specify a function to execute when it is awoken.

This change makes for a clearer (less confusing) path of execution in the scheduler, now the code to execute when a thread awakes is closer to the code that puts the thread to sleep (WaitSynch1, WaitSynchN). It also allows us to implement the special wake up behavior of ReplyAndReceive without hacking up WaitObject::WakeupAllWaitingThreads.

If savestates are desired in the future, we can change this implementation to one similar to the CoreTiming event system, where we first register the callback functions at startup and assign their identifiers to the Thread callback variable instead of directly assigning a lambda to the wake up callback variable.

src/core/hle/kernel/thread.cpp

@@ -247,12 +247,15 @@ static void ThreadWakeupCallback(u64 thread_handle, int cycles_late) {
 
     if (thread->status == THREADSTATUS_WAIT_SYNCH_ANY ||
         thread->status == THREADSTATUS_WAIT_SYNCH_ALL || thread->status == THREADSTATUS_WAIT_ARB) {
-        thread->wait_set_output = false;
+
+        // 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();
-        thread->SetWaitSynchronizationResult(RESULT_TIMEOUT);
     }
 
     thread->ResumeFromWait();
@@ -278,6 +281,9 @@ void Thread::ResumeFromWait() {
         break;
 
     case THREADSTATUS_READY:
+        // The thread's wakeup callback must have already been cleared when the thread was first
+        // awoken.
+        ASSERT(wakeup_callback == nullptr);
         // If the thread is waiting on multiple wait objects, it might be awoken more than once
         // before actually resuming. We can ignore subsequent wakeups if the thread status has
         // already been set to THREADSTATUS_READY.
@@ -293,6 +299,8 @@ void Thread::ResumeFromWait() {
         return;
     }
 
+    wakeup_callback = nullptr;
+
     ready_queue.push_back(current_priority, this);
     status = THREADSTATUS_READY;
     Core::System::GetInstance().PrepareReschedule();
@@ -394,7 +402,6 @@ ResultVal<SharedPtr<Thread>> Thread::Create(std::string name, VAddr entry_point,
     thread->nominal_priority = thread->current_priority = priority;
     thread->last_running_ticks = CoreTiming::GetTicks();
     thread->processor_id = processor_id;
-    thread->wait_set_output = false;
     thread->wait_objects.clear();
     thread->wait_address = 0;
     thread->name = std::move(name);

src/core/hle/kernel/thread.h

@@ -41,6 +41,11 @@ enum ThreadStatus {
     THREADSTATUS_DEAD            ///< Run to completion, or forcefully terminated
 };
 
+enum class ThreadWakeupReason {
+    Signal, // The thread was woken up by WakeupAllWaitingThreads due to an object signal.
+    Timeout // The thread was woken up due to a wait timeout.
+};
+
 namespace Kernel {
 
 class Mutex;
@@ -197,14 +202,18 @@ public:
 
     VAddr wait_address; ///< If waiting on an AddressArbiter, this is the arbitration address
 
-    /// True if the WaitSynchronizationN output parameter should be set on thread wakeup.
-    bool wait_set_output;
-
     std::string name;
 
     /// 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);
+    // 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.
+    std::function<WakeupCallback> wakeup_callback;
+
 private:
     Thread();
     ~Thread() override;

src/core/hle/kernel/wait_object.cpp

@@ -71,23 +71,20 @@ void WaitObject::WakeupAllWaitingThreads() {
     while (auto thread = GetHighestPriorityReadyThread()) {
         if (!thread->IsSleepingOnWaitAll()) {
             Acquire(thread.get());
-            // Set the output index of the WaitSynchronizationN call to the index of this object.
-            if (thread->wait_set_output) {
-                thread->SetWaitSynchronizationOutput(thread->GetWaitObjectIndex(this));
-                thread->wait_set_output = false;
-            }
         } else {
             for (auto& object : thread->wait_objects) {
                 object->Acquire(thread.get());
             }
-            // Note: This case doesn't update the output index of WaitSynchronizationN.
         }
 
+        // Invoke the wakeup callback before clearing the wait objects
+        if (thread->wakeup_callback)
+            thread->wakeup_callback(ThreadWakeupReason::Signal, thread, this);
+
         for (auto& object : thread->wait_objects)
             object->RemoveWaitingThread(thread.get());
         thread->wait_objects.clear();
 
-        thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
         thread->ResumeFromWait();
     }
 }

src/core/hle/svc.cpp

@@ -271,6 +271,24 @@ static ResultCode WaitSynchronization1(Kernel::Handle handle, s64 nano_seconds)
         // Create an event to wake the thread up after the specified nanosecond delay has passed
         thread->WakeAfterDelay(nano_seconds);
 
+        thread->wakeup_callback = [](ThreadWakeupReason reason,
+                                     Kernel::SharedPtr<Kernel::Thread> thread,
+                                     Kernel::SharedPtr<Kernel::WaitObject> object) {
+
+            ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ANY);
+
+            if (reason == ThreadWakeupReason::Timeout) {
+                thread->SetWaitSynchronizationResult(Kernel::RESULT_TIMEOUT);
+                return;
+            }
+
+            ASSERT(reason == ThreadWakeupReason::Signal);
+            thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+
+            // WaitSynchronization1 doesn't have an output index like WaitSynchronizationN, so we
+            // don't have to do anything else here.
+        };
+
         Core::System::GetInstance().PrepareReschedule();
 
         // Note: The output of this SVC will be set to RESULT_SUCCESS if the thread
@@ -344,6 +362,23 @@ static ResultCode WaitSynchronizationN(s32* out, Kernel::Handle* handles, s32 ha
         // Create an event to wake the thread up after the specified nanosecond delay has passed
         thread->WakeAfterDelay(nano_seconds);
 
+        thread->wakeup_callback = [](ThreadWakeupReason reason,
+                                     Kernel::SharedPtr<Kernel::Thread> thread,
+                                     Kernel::SharedPtr<Kernel::WaitObject> object) {
+
+            ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ALL);
+
+            if (reason == ThreadWakeupReason::Timeout) {
+                thread->SetWaitSynchronizationResult(Kernel::RESULT_TIMEOUT);
+                return;
+            }
+
+            ASSERT(reason == ThreadWakeupReason::Signal);
+
+            thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+            // The wait_all case does not update the output index.
+        };
+
         Core::System::GetInstance().PrepareReschedule();
 
         // This value gets set to -1 by default in this case, it is not modified after this.
@@ -389,12 +424,28 @@ static ResultCode WaitSynchronizationN(s32* out, Kernel::Handle* handles, s32 ha
         // Create an event to wake the thread up after the specified nanosecond delay has passed
         thread->WakeAfterDelay(nano_seconds);
 
+        thread->wakeup_callback = [](ThreadWakeupReason reason,
+                                     Kernel::SharedPtr<Kernel::Thread> thread,
+                                     Kernel::SharedPtr<Kernel::WaitObject> object) {
+
+            ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ANY);
+
+            if (reason == ThreadWakeupReason::Timeout) {
+                thread->SetWaitSynchronizationResult(Kernel::RESULT_TIMEOUT);
+                return;
+            }
+
+            ASSERT(reason == ThreadWakeupReason::Signal);
+
+            thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+            thread->SetWaitSynchronizationOutput(thread->GetWaitObjectIndex(object.get()));
+        };
+
         Core::System::GetInstance().PrepareReschedule();
 
         // Note: The output of this SVC will be set to RESULT_SUCCESS if the thread resumes due to a
         // signal in one of its wait objects.
         // Otherwise we retain the default value of timeout, and -1 in the out parameter
-        thread->wait_set_output = true;
         *out = -1;
         return Kernel::RESULT_TIMEOUT;
     }
@@ -483,8 +534,6 @@ static ResultCode ReplyAndReceive(s32* index, Kernel::Handle* handles, s32 handl
 
     // No objects were ready to be acquired, prepare to suspend the thread.
 
-    // TODO(Subv): Perform IPC translation upon wakeup.
-
     // Put the thread to sleep
     thread->status = THREADSTATUS_WAIT_SYNCH_ANY;
 
@@ -496,12 +545,24 @@ static ResultCode ReplyAndReceive(s32* index, Kernel::Handle* handles, s32 handl
 
     thread->wait_objects = std::move(objects);
 
+    thread->wakeup_callback = [](ThreadWakeupReason reason,
+                                 Kernel::SharedPtr<Kernel::Thread> thread,
+                                 Kernel::SharedPtr<Kernel::WaitObject> object) {
+
+        ASSERT(thread->status == THREADSTATUS_WAIT_SYNCH_ANY);
+        ASSERT(reason == ThreadWakeupReason::Signal);
+
+        thread->SetWaitSynchronizationResult(RESULT_SUCCESS);
+        thread->SetWaitSynchronizationOutput(thread->GetWaitObjectIndex(object.get()));
+
+        // TODO(Subv): Perform IPC translation upon wakeup.
+    };
+
     Core::System::GetInstance().PrepareReschedule();
 
     // Note: The output of this SVC will be set to RESULT_SUCCESS if the thread resumes due to a
     // signal in one of its wait objects, or to 0xC8A01836 if there was a translation error.
     // By default the index is set to -1.
-    thread->wait_set_output = true;
     *index = -1;
     return RESULT_SUCCESS;
 }