10 files changed, 128 insertions, 51 deletions

diff --git a/src/core/hle/kernel/hle_ipc.cpp b/src/core/hle/kernel/hle_ipc.cpp
index d9faf4b53..293756790 100644
--- a/src/core/hle/kernel/hle_ipc.cpp
+++ b/src/core/hle/kernel/hle_ipc.cpp
@@ -7,6 +7,7 @@
 #include "common/common_funcs.h"
 #include "common/common_types.h"
 #include "core/hle/ipc_helpers.h"
+#include "core/hle/kernel/event.h"
 #include "core/hle/kernel/handle_table.h"
 #include "core/hle/kernel/hle_ipc.h"
 #include "core/hle/kernel/kernel.h"
@@ -26,6 +27,32 @@ void SessionRequestHandler::ClientDisconnected(SharedPtr<ServerSession> server_s
     boost::range::remove_erase(connected_sessions, server_session);
 }
 
+SharedPtr<Event> HLERequestContext::SleepClientThread(SharedPtr<Thread> thread,
+                                                      const std::string& reason, u64 timeout,
+                                                      WakeupCallback&& callback) {
+
+    // Put the client thread to sleep until the wait event is signaled or the timeout expires.
+    thread->wakeup_callback =
+        [context = *this, callback](ThreadWakeupReason reason, SharedPtr<Thread> thread,
+                                    SharedPtr<WaitObject> object, size_t index) mutable -> bool {
+        ASSERT(thread->status == THREADSTATUS_WAIT_HLE_EVENT);
+        callback(thread, context, reason);
+        context.WriteToOutgoingCommandBuffer(*thread);
+        return true;
+    };
+
+    auto event = Kernel::Event::Create(Kernel::ResetType::OneShot, "HLE Pause Event: " + reason);
+    thread->status = THREADSTATUS_WAIT_HLE_EVENT;
+    thread->wait_objects = {event};
+    event->AddWaitingThread(thread);
+
+    if (timeout > 0) {
+        thread->WakeAfterDelay(timeout);
+    }
+
+    return event;
+}
+
 HLERequestContext::HLERequestContext(SharedPtr<Kernel::ServerSession> server_session)
     : server_session(std::move(server_session)) {
     cmd_buf[0] = 0;
@@ -35,7 +62,7 @@ HLERequestContext::~HLERequestContext() = default;
 
 void HLERequestContext::ParseCommandBuffer(u32_le* src_cmdbuf, bool incoming) {
     IPC::RequestParser rp(src_cmdbuf);
-    command_header = std::make_unique<IPC::CommandHeader>(rp.PopRaw<IPC::CommandHeader>());
+    command_header = std::make_shared<IPC::CommandHeader>(rp.PopRaw<IPC::CommandHeader>());
 
     if (command_header->type == IPC::CommandType::Close) {
         // Close does not populate the rest of the IPC header
@@ -45,7 +72,7 @@ void HLERequestContext::ParseCommandBuffer(u32_le* src_cmdbuf, bool incoming) {
     // If handle descriptor is present, add size of it
     if (command_header->enable_handle_descriptor) {
         handle_descriptor_header =
-            std::make_unique<IPC::HandleDescriptorHeader>(rp.PopRaw<IPC::HandleDescriptorHeader>());
+            std::make_shared<IPC::HandleDescriptorHeader>(rp.PopRaw<IPC::HandleDescriptorHeader>());
         if (handle_descriptor_header->send_current_pid) {
             rp.Skip(2, false);
         }
@@ -88,7 +115,7 @@ void HLERequestContext::ParseCommandBuffer(u32_le* src_cmdbuf, bool incoming) {
         // All outgoing domain messages have the domain header, if only incoming has it
         if (incoming || domain_message_header) {
             domain_message_header =
-                std::make_unique<IPC::DomainMessageHeader>(rp.PopRaw<IPC::DomainMessageHeader>());
+                std::make_shared<IPC::DomainMessageHeader>(rp.PopRaw<IPC::DomainMessageHeader>());
         } else {
             if (Session()->IsDomain())
                 LOG_WARNING(IPC, "Domain request has no DomainMessageHeader!");
@@ -96,7 +123,7 @@ void HLERequestContext::ParseCommandBuffer(u32_le* src_cmdbuf, bool incoming) {
     }
 
     data_payload_header =
-        std::make_unique<IPC::DataPayloadHeader>(rp.PopRaw<IPC::DataPayloadHeader>());
+        std::make_shared<IPC::DataPayloadHeader>(rp.PopRaw<IPC::DataPayloadHeader>());
 
     data_payload_offset = rp.GetCurrentOffset();
 
@@ -159,8 +186,11 @@ ResultCode HLERequestContext::PopulateFromIncomingCommandBuffer(u32_le* src_cmdb
     return RESULT_SUCCESS;
 }
 
-ResultCode HLERequestContext::WriteToOutgoingCommandBuffer(u32_le* dst_cmdbuf, Process& dst_process,
-                                                           HandleTable& dst_table) {
+ResultCode HLERequestContext::WriteToOutgoingCommandBuffer(Thread& thread) {
+    std::array<u32, IPC::COMMAND_BUFFER_LENGTH> dst_cmdbuf;
+    Memory::ReadBlock(*thread.owner_process, thread.GetTLSAddress(), dst_cmdbuf.data(),
+                      dst_cmdbuf.size() * sizeof(u32));
+
     // The header was already built in the internal command buffer. Attempt to parse it to verify
     // the integrity and then copy it over to the target command buffer.
     ParseCommandBuffer(cmd_buf.data(), false);
@@ -171,7 +201,7 @@ ResultCode HLERequestContext::WriteToOutgoingCommandBuffer(u32_le* dst_cmdbuf, P
     if (domain_message_header)
         size -= sizeof(IPC::DomainMessageHeader) / sizeof(u32);
 
-    std::copy_n(cmd_buf.begin(), size, dst_cmdbuf);
+    std::copy_n(cmd_buf.begin(), size, dst_cmdbuf.data());
 
     if (command_header->enable_handle_descriptor) {
         ASSERT_MSG(!move_objects.empty() || !copy_objects.empty(),
@@ -213,6 +243,11 @@ ResultCode HLERequestContext::WriteToOutgoingCommandBuffer(u32_le* dst_cmdbuf, P
             dst_cmdbuf[domain_offset++] = static_cast<u32_le>(request_handlers.size());
         }
     }
+
+    // Copy the translated command buffer back into the thread's command buffer area.
+    Memory::WriteBlock(*thread.owner_process, thread.GetTLSAddress(), dst_cmdbuf.data(),
+                       dst_cmdbuf.size() * sizeof(u32));
+
     return RESULT_SUCCESS;
 }
 
diff --git a/src/core/hle/kernel/hle_ipc.h b/src/core/hle/kernel/hle_ipc.h
index b5631b773..8b35da4c9 100644
--- a/src/core/hle/kernel/hle_ipc.h
+++ b/src/core/hle/kernel/hle_ipc.h
@@ -6,6 +6,7 @@
 
 #include <array>
 #include <memory>
+#include <string>
 #include <vector>
 #include <boost/container/small_vector.hpp>
 #include "common/common_types.h"
@@ -13,6 +14,7 @@
 #include "core/hle/ipc.h"
 #include "core/hle/kernel/kernel.h"
 #include "core/hle/kernel/server_session.h"
+#include "core/hle/kernel/thread.h"
 
 namespace Service {
 class ServiceFrameworkBase;
@@ -24,6 +26,7 @@ class Domain;
 class HandleTable;
 class HLERequestContext;
 class Process;
+class Event;
 
 /**
  * Interface implemented by HLE Session handlers.
@@ -102,14 +105,31 @@ public:
         return server_session;
     }
 
+    using WakeupCallback = std::function<void(SharedPtr<Thread> thread, HLERequestContext& context,
+                                              ThreadWakeupReason reason)>;
+
+    /**
+     * Puts the specified guest thread to sleep until the returned event is signaled or until the
+     * specified timeout expires.
+     * @param thread Thread to be put to sleep.
+     * @param reason Reason for pausing the thread, to be used for debugging purposes.
+     * @param timeout Timeout in nanoseconds after which the thread will be awoken and the callback
+     * invoked with a Timeout reason.
+     * @param callback Callback to be invoked when the thread is resumed. This callback must write
+     * the entire command response once again, regardless of the state of it before this function
+     * was called.
+     * @returns Event that when signaled will resume the thread and call the callback function.
+     */
+    SharedPtr<Event> SleepClientThread(SharedPtr<Thread> thread, const std::string& reason,
+                                       u64 timeout, WakeupCallback&& callback);
+
     void ParseCommandBuffer(u32_le* src_cmdbuf, bool incoming);
 
     /// Populates this context with data from the requesting process/thread.
     ResultCode PopulateFromIncomingCommandBuffer(u32_le* src_cmdbuf, Process& src_process,
                                                  HandleTable& src_table);
     /// Writes data from this context back to the requesting process/thread.
-    ResultCode WriteToOutgoingCommandBuffer(u32_le* dst_cmdbuf, Process& dst_process,
-                                            HandleTable& dst_table);
+    ResultCode WriteToOutgoingCommandBuffer(Thread& thread);
 
     u32_le GetCommand() const {
         return command;
@@ -139,7 +159,7 @@ public:
         return buffer_c_desciptors;
     }
 
-    const std::unique_ptr<IPC::DomainMessageHeader>& GetDomainMessageHeader() const {
+    const std::shared_ptr<IPC::DomainMessageHeader>& GetDomainMessageHeader() const {
         return domain_message_header;
     }
 
@@ -212,10 +232,10 @@ private:
     boost::container::small_vector<SharedPtr<Object>, 8> copy_objects;
     boost::container::small_vector<std::shared_ptr<SessionRequestHandler>, 8> domain_objects;
 
-    std::unique_ptr<IPC::CommandHeader> command_header;
-    std::unique_ptr<IPC::HandleDescriptorHeader> handle_descriptor_header;
-    std::unique_ptr<IPC::DataPayloadHeader> data_payload_header;
-    std::unique_ptr<IPC::DomainMessageHeader> domain_message_header;
+    std::shared_ptr<IPC::CommandHeader> command_header;
+    std::shared_ptr<IPC::HandleDescriptorHeader> handle_descriptor_header;
+    std::shared_ptr<IPC::DataPayloadHeader> data_payload_header;
+    std::shared_ptr<IPC::DomainMessageHeader> domain_message_header;
     std::vector<IPC::BufferDescriptorX> buffer_x_desciptors;
     std::vector<IPC::BufferDescriptorABW> buffer_a_desciptors;
     std::vector<IPC::BufferDescriptorABW> buffer_b_desciptors;
diff --git a/src/core/hle/kernel/thread.cpp b/src/core/hle/kernel/thread.cpp
index a39c53db5..145f50887 100644
--- a/src/core/hle/kernel/thread.cpp
+++ b/src/core/hle/kernel/thread.cpp
@@ -55,16 +55,6 @@ inline static u32 const NewThreadId() {
 Thread::Thread() {}
 Thread::~Thread() {}
 
-/**
- * Check if the specified thread is waiting on the specified address to be arbitrated
- * @param thread The thread to test
- * @param wait_address The address to test against
- * @return True if the thread is waiting, false otherwise
- */
-static bool CheckWait_AddressArbiter(const Thread* thread, VAddr wait_address) {
-    return thread->status == THREADSTATUS_WAIT_ARB && wait_address == thread->wait_address;
-}
-
 void Thread::Stop() {
     // Cancel any outstanding wakeup events for this thread
     CoreTiming::UnscheduleEvent(ThreadWakeupEventType, callback_handle);
@@ -102,12 +92,6 @@ void WaitCurrentThread_Sleep() {
     thread->status = THREADSTATUS_WAIT_SLEEP;
 }
 
-void WaitCurrentThread_ArbitrateAddress(VAddr wait_address) {
-    Thread* thread = GetCurrentThread();
-    thread->wait_address = wait_address;
-    thread->status = THREADSTATUS_WAIT_ARB;
-}
-
 void ExitCurrentThread() {
     Thread* thread = GetCurrentThread();
     thread->Stop();
@@ -129,7 +113,8 @@ static void ThreadWakeupCallback(u64 thread_handle, int cycles_late) {
     bool resume = true;
 
     if (thread->status == THREADSTATUS_WAIT_SYNCH_ANY ||
-        thread->status == THREADSTATUS_WAIT_SYNCH_ALL || thread->status == THREADSTATUS_WAIT_ARB) {
+        thread->status == THREADSTATUS_WAIT_SYNCH_ALL ||
+        thread->status == THREADSTATUS_WAIT_HLE_EVENT) {
 
         // Remove the thread from each of its waiting objects' waitlists
         for (auto& object : thread->wait_objects)
@@ -163,7 +148,7 @@ void Thread::ResumeFromWait() {
     switch (status) {
     case THREADSTATUS_WAIT_SYNCH_ALL:
     case THREADSTATUS_WAIT_SYNCH_ANY:
-    case THREADSTATUS_WAIT_ARB:
+    case THREADSTATUS_WAIT_HLE_EVENT:
     case THREADSTATUS_WAIT_SLEEP:
     case THREADSTATUS_WAIT_IPC:
         break;
diff --git a/src/core/hle/kernel/thread.h b/src/core/hle/kernel/thread.h
index 4fd2fc2f8..dbf47e269 100644
--- a/src/core/hle/kernel/thread.h
+++ b/src/core/hle/kernel/thread.h
@@ -38,7 +38,7 @@ enum ThreadProcessorId : s32 {
 enum ThreadStatus {
     THREADSTATUS_RUNNING,        ///< Currently running
     THREADSTATUS_READY,          ///< Ready to run
-    THREADSTATUS_WAIT_ARB,       ///< Waiting on an address arbiter
+    THREADSTATUS_WAIT_HLE_EVENT, ///< Waiting for hle event to finish
     THREADSTATUS_WAIT_SLEEP,     ///< Waiting due to a SleepThread SVC
     THREADSTATUS_WAIT_IPC,       ///< Waiting for the reply from an IPC request
     THREADSTATUS_WAIT_SYNCH_ANY, ///< Waiting due to WaitSynch1 or WaitSynchN with wait_all = false
diff --git a/src/core/hle/kernel/wait_object.cpp b/src/core/hle/kernel/wait_object.cpp
index ec147b84c..b08ac72c1 100644
--- a/src/core/hle/kernel/wait_object.cpp
+++ b/src/core/hle/kernel/wait_object.cpp
@@ -39,7 +39,8 @@ SharedPtr<Thread> WaitObject::GetHighestPriorityReadyThread() {
     for (const auto& thread : waiting_threads) {
         // The list of waiting threads must not contain threads that are not waiting to be awakened.
         ASSERT_MSG(thread->status == THREADSTATUS_WAIT_SYNCH_ANY ||
-                       thread->status == THREADSTATUS_WAIT_SYNCH_ALL,
+                       thread->status == THREADSTATUS_WAIT_SYNCH_ALL ||
+                       thread->status == THREADSTATUS_WAIT_HLE_EVENT,
                    "Inconsistent thread statuses in waiting_threads");
 
         if (thread->current_priority >= candidate_priority)
diff --git a/src/core/hle/service/nvflinger/buffer_queue.cpp b/src/core/hle/service/nvflinger/buffer_queue.cpp
index e2c25048b..e4ff2e267 100644
--- a/src/core/hle/service/nvflinger/buffer_queue.cpp
+++ b/src/core/hle/service/nvflinger/buffer_queue.cpp
@@ -26,24 +26,30 @@ void BufferQueue::SetPreallocatedBuffer(u32 slot, IGBPBuffer& igbp_buffer) {
     LOG_WARNING(Service, "Adding graphics buffer %u", slot);
 
     queue.emplace_back(buffer);
+
+    if (buffer_wait_event) {
+        buffer_wait_event->Signal();
+    }
 }
 
-u32 BufferQueue::DequeueBuffer(u32 pixel_format, u32 width, u32 height) {
+boost::optional<u32> BufferQueue::DequeueBuffer(u32 width, u32 height) {
     auto itr = std::find_if(queue.begin(), queue.end(), [&](const Buffer& buffer) {
         // Only consider free buffers. Buffers become free once again after they've been Acquired
         // and Released by the compositor, see the NVFlinger::Compose method.
-        if (buffer.status != Buffer::Status::Free)
+        if (buffer.status != Buffer::Status::Free) {
             return false;
+        }
 
         // Make sure that the parameters match.
         return buffer.igbp_buffer.width == width && buffer.igbp_buffer.height == height;
     });
+
     if (itr == queue.end()) {
-        LOG_CRITICAL(Service_NVDRV, "no free buffers for pixel_format=%d, width=%d, height=%d",
-                     pixel_format, width, height);
-        itr = queue.begin();
+        return boost::none;
     }
 
+    buffer_wait_event = nullptr;
+
     itr->status = Buffer::Status::Dequeued;
     return itr->slot;
 }
@@ -81,6 +87,10 @@ void BufferQueue::ReleaseBuffer(u32 slot) {
     ASSERT(itr != queue.end());
     ASSERT(itr->status == Buffer::Status::Acquired);
     itr->status = Buffer::Status::Free;
+
+    if (buffer_wait_event) {
+        buffer_wait_event->Signal();
+    }
 }
 
 u32 BufferQueue::Query(QueryType type) {
@@ -96,5 +106,10 @@ u32 BufferQueue::Query(QueryType type) {
     return 0;
 }
 
+void BufferQueue::SetBufferWaitEvent(Kernel::SharedPtr<Kernel::Event>&& wait_event) {
+    ASSERT_MSG(!buffer_wait_event, "buffer_wait_event only supports a single waiting thread!");
+    buffer_wait_event = std::move(wait_event);
+}
+
 } // namespace NVFlinger
 } // namespace Service
diff --git a/src/core/hle/service/nvflinger/buffer_queue.h b/src/core/hle/service/nvflinger/buffer_queue.h
index ef9732769..686eadca7 100644
--- a/src/core/hle/service/nvflinger/buffer_queue.h
+++ b/src/core/hle/service/nvflinger/buffer_queue.h
@@ -69,12 +69,13 @@ public:
     };
 
     void SetPreallocatedBuffer(u32 slot, IGBPBuffer& buffer);
-    u32 DequeueBuffer(u32 pixel_format, u32 width, u32 height);
+    boost::optional<u32> DequeueBuffer(u32 width, u32 height);
     const IGBPBuffer& RequestBuffer(u32 slot) const;
     void QueueBuffer(u32 slot, BufferTransformFlags transform);
     boost::optional<const Buffer&> AcquireBuffer();
     void ReleaseBuffer(u32 slot);
     u32 Query(QueryType type);
+    void SetBufferWaitEvent(Kernel::SharedPtr<Kernel::Event>&& wait_event);
 
     u32 GetId() const {
         return id;
@@ -90,6 +91,9 @@ private:
 
     std::vector<Buffer> queue;
     Kernel::SharedPtr<Kernel::Event> native_handle;
+
+    /// Used to signal waiting thread when no buffers are available
+    Kernel::SharedPtr<Kernel::Event> buffer_wait_event;
 };
 
 } // namespace NVFlinger
diff --git a/src/core/hle/service/service.cpp b/src/core/hle/service/service.cpp
index 8818b0f0f..a1ca8a033 100644
--- a/src/core/hle/service/service.cpp
+++ b/src/core/hle/service/service.cpp
@@ -152,8 +152,7 @@ ResultCode ServiceFrameworkBase::HandleSyncRequest(Kernel::HLERequestContext& co
         UNIMPLEMENTED_MSG("command_type=%d", context.GetCommandType());
     }
 
-    u32* cmd_buf = (u32*)Memory::GetPointer(Kernel::GetCurrentThread()->GetTLSAddress());
-    context.WriteToOutgoingCommandBuffer(cmd_buf, *Core::CurrentProcess(), Kernel::g_handle_table);
+    context.WriteToOutgoingCommandBuffer(*Kernel::GetCurrentThread());
 
     return RESULT_SUCCESS;
 }
diff --git a/src/core/hle/service/vi/vi.cpp b/src/core/hle/service/vi/vi.cpp
index 0aa621dfe..7b6453447 100644
--- a/src/core/hle/service/vi/vi.cpp
+++ b/src/core/hle/service/vi/vi.cpp
@@ -486,12 +486,30 @@ private:
             ctx.WriteBuffer(response.Serialize());
         } else if (transaction == TransactionId::DequeueBuffer) {
             IGBPDequeueBufferRequestParcel request{ctx.ReadBuffer()};
-
-            u32 slot = buffer_queue->DequeueBuffer(request.data.pixel_format, request.data.width,
-                                                   request.data.height);
-
-            IGBPDequeueBufferResponseParcel response{slot};
-            ctx.WriteBuffer(response.Serialize());
+            const u32 width{request.data.width};
+            const u32 height{request.data.height};
+            boost::optional<u32> slot = buffer_queue->DequeueBuffer(width, height);
+
+            if (slot != boost::none) {
+                // Buffer is available
+                IGBPDequeueBufferResponseParcel response{*slot};
+                ctx.WriteBuffer(response.Serialize());
+            } else {
+                // Wait the current thread until a buffer becomes available
+                auto wait_event = ctx.SleepClientThread(
+                    Kernel::GetCurrentThread(), "IHOSBinderDriver::DequeueBuffer", -1,
+                    [=](Kernel::SharedPtr<Kernel::Thread> thread, Kernel::HLERequestContext& ctx,
+                        ThreadWakeupReason reason) {
+                        // Repeat TransactParcel DequeueBuffer when a buffer is available
+                        auto buffer_queue = nv_flinger->GetBufferQueue(id);
+                        boost::optional<u32> slot = buffer_queue->DequeueBuffer(width, height);
+                        IGBPDequeueBufferResponseParcel response{*slot};
+                        ctx.WriteBuffer(response.Serialize());
+                        IPC::ResponseBuilder rb{ctx, 2};
+                        rb.Push(RESULT_SUCCESS);
+                    });
+                buffer_queue->SetBufferWaitEvent(std::move(wait_event));
+            }
         } else if (transaction == TransactionId::RequestBuffer) {
             IGBPRequestBufferRequestParcel request{ctx.ReadBuffer()};
 
diff --git a/src/yuzu/debugger/wait_tree.cpp b/src/yuzu/debugger/wait_tree.cpp
index 7a62f57b5..cae2864e5 100644
--- a/src/yuzu/debugger/wait_tree.cpp
+++ b/src/yuzu/debugger/wait_tree.cpp
@@ -150,8 +150,8 @@ QString WaitTreeThread::GetText() const {
     case THREADSTATUS_READY:
         status = tr("ready");
         break;
-    case THREADSTATUS_WAIT_ARB:
-        status = tr("waiting for address 0x%1").arg(thread.wait_address, 8, 16, QLatin1Char('0'));
+    case THREADSTATUS_WAIT_HLE_EVENT:
+        status = tr("waiting for HLE return");
         break;
     case THREADSTATUS_WAIT_SLEEP:
         status = tr("sleeping");
@@ -180,7 +180,7 @@ QColor WaitTreeThread::GetColor() const {
         return QColor(Qt::GlobalColor::darkGreen);
     case THREADSTATUS_READY:
         return QColor(Qt::GlobalColor::darkBlue);
-    case THREADSTATUS_WAIT_ARB:
+    case THREADSTATUS_WAIT_HLE_EVENT:
         return QColor(Qt::GlobalColor::darkRed);
     case THREADSTATUS_WAIT_SLEEP:
         return QColor(Qt::GlobalColor::darkYellow);