Merge branch 'main' into stabilize-queue-unit-tests

Author: jasonsandlin

Include/XTaskQueue.h

@@ -185,6 +185,27 @@ STDAPI_(void) XTaskQueueCloseHandle(
 /// preventing new items from being queued.  Once a queue is terminated
 /// its handle can be closed. New items cannot be enqueued to a task
 /// queue that has been terminated.
+///
+/// Each task queue terminates independently. For composite queues created
+/// with XTaskQueueCreateComposite, terminating a composite queue does NOT
+/// automatically terminate other queues sharing the same ports.
+///
+/// When wait=true:
+/// - Blocks until this queue's termination callback completes
+/// - Does NOT wait for other independent queues (including composite delegates)
+/// - Ensures this queue's termination callback has finished executing
+/// - Safe to unload code/modules after this returns
+///
+/// When wait=false:
+/// - Returns immediately after initiating termination
+/// - The termination callback will be invoked asynchronously when termination completes
+///
+/// The termination callback is invoked after all work and completion callbacks
+/// have been canceled or completed. After the termination callback returns, the
+/// implementation will no longer access the queue's internal state.
+///
+/// For coordinated shutdown of multiple queues sharing ports, use termination
+/// callbacks to track completion of each queue before performing final cleanup.
 /// </summary>
 /// <param name='queue'>The queue to terminate.</param>
 /// <param name='wait'>True to wait for the termination to complete.</param>

Source/Task/TaskQueue.cpp

@@ -518,7 +518,10 @@ HRESULT __stdcall TaskQueuePortImpl::PrepareTerminate(
     std::unique_ptr<TerminationEntry> term(new (std::nothrow) TerminationEntry);
     RETURN_IF_NULL_ALLOC(term);
 
-    RETURN_HR_IF(E_OUTOFMEMORY, !m_terminationList->reserve_node(term->node));
+    {
+        std::lock_guard<std::mutex> lock(m_terminationLock);
+        RETURN_HR_IF(E_OUTOFMEMORY, !m_terminationList->reserve_node(term->node));
+    }
 
     term->callbackContext = callbackContext;
     term->callback = callback;
@@ -540,6 +543,7 @@ void __stdcall TaskQueuePortImpl::CancelTermination(
 
     if (term->node != 0)
     {
+        std::lock_guard<std::mutex> lock(m_terminationLock);
         m_terminationList->free_node(term->node);
     }
 
@@ -565,6 +569,7 @@ void __stdcall TaskQueuePortImpl::Terminate(
     }
     else
     {
+        std::lock_guard<std::mutex> lock(m_terminationLock);
         m_pendingTerminationList->push_back(term, term->node);
         term->node = 0;
     }
@@ -687,7 +692,7 @@ bool TaskQueuePortImpl::Wait(
     _In_ uint32_t timeout)
 {
 #ifdef _WIN32
-    while (m_suspended || (m_queueList->empty() && m_terminationList->empty()))
+    while (m_suspended || (m_queueList->empty() && TerminationListEmpty()))
     {
         if (portContext->GetStatus() == TaskQueuePortStatus::Terminated)
         {
@@ -749,7 +754,7 @@ bool TaskQueuePortImpl::Wait(
     }
 
 #else
-    while (m_suspended || (m_queueList->empty() && m_terminationList->empty()))
+    while (m_suspended || (m_queueList->empty() && TerminationListEmpty()))
     {
         if (portContext->GetStatus() == TaskQueuePortStatus::Terminated)
         {
@@ -825,20 +830,31 @@ void __stdcall TaskQueuePortImpl::ResumeTermination(
     {
         // Removed the last external callback.  Look for
         // parked terminations and reschedule them.
-
-        m_pendingTerminationList->remove_if([&](auto& entry, auto address)
+        // Use a temporary list sharing the same heap to avoid allocation
+        LocklessQueue<TerminationEntry*> entries_to_schedule(*m_pendingTerminationList.get());
+        
         {
-            if (entry->portContext == portContext)
+            std::lock_guard<std::mutex> lock(m_terminationLock);
+            m_pendingTerminationList->remove_if([&](auto& entry, auto address)
             {
-                // This entry is for the port that's resuming,
-                // we can schedule it.
-                entry->node = address;
-                ScheduleTermination(entry);
-                return true;
-            }
+                if (entry->portContext == portContext)
+                {
+                    entries_to_schedule.push_back(entry, address);
+                    return true;
+                }
 
-            return false;
-        });
+                return false;
+            });
+        }
+
+        // Schedule entries outside the lock
+        TerminationEntry* entry;
+        uint64_t address;
+        while (entries_to_schedule.pop_front(entry, address))
+        {
+            entry->node = address;
+            ScheduleTermination(entry);
+        }
     }
 }
 
@@ -872,18 +888,21 @@ void __stdcall TaskQueuePortImpl::ResumePort()
         m_queueList->push_back(std::move(queueEntry), address);
     }
 
-    TerminationEntry* terminationEntry;
-    LocklessQueue<TerminationEntry*> retainTerminations(*(m_terminationList.get()));
-
-    while (m_terminationList->pop_front(terminationEntry, address))
     {
-        notifyCount++;
-        retainTerminations.push_back(std::move(terminationEntry), address);
-    }
+        std::lock_guard<std::mutex> lock(m_terminationLock);
+        TerminationEntry* terminationEntry;
+        LocklessQueue<TerminationEntry*> retainTerminations(*(m_terminationList.get()));
 
-    while (retainTerminations.pop_front(terminationEntry, address))
-    {
-        m_terminationList->push_back(std::move(terminationEntry), address);
+        while (m_terminationList->pop_front(terminationEntry, address))
+        {
+            notifyCount++;
+            retainTerminations.push_back(std::move(terminationEntry), address);
+        }
+
+        while (retainTerminations.pop_front(terminationEntry, address))
+        {
+            m_terminationList->push_back(std::move(terminationEntry), address);
+        }
     }
 
     m_suspended = false;
@@ -1199,19 +1218,46 @@ void TaskQueuePortImpl::NotifyItemQueued()
 
 void TaskQueuePortImpl::SignalTerminations()
 {
-    m_terminationList->remove_if([this](auto& entry, auto address)
+    // Collect entries to process outside the iteration to avoid concurrent modification races
+    // when callbacks invoke nested Terminate() calls.
+    // Use a temporary list sharing the same heap to avoid allocation.
+    LocklessQueue<TerminationEntry*> entries_to_process(*m_terminationList.get());
+
     {
-        if (entry->portContext->GetStatus() >= TaskQueuePortStatus::Terminating)
+        std::lock_guard<std::mutex> lock(m_terminationLock);
+        m_terminationList->remove_if([this, &entries_to_process](auto& entry, auto address)
         {
-            entry->portContext->SetStatus(TaskQueuePortStatus::Terminated);
-            entry->callback(entry->callbackContext);
+            if (entry->portContext->GetStatus() >= TaskQueuePortStatus::Terminating)
+            {
+                entry->portContext->SetStatus(TaskQueuePortStatus::Terminated);
+                entries_to_process.push_back(entry, address);
+                return true;
+            }
+
+            return false;
+        });
+    }
+
+    // Now process callbacks outside the remove_if iteration
+    // This prevents races when callbacks invoke nested operations like Terminate()
+    TerminationEntry* entry;
+    uint64_t address;
+    while (entries_to_process.pop_front(entry, address))
+    {
+        // AddRef portContext to prevent UAF if callback releases the queue
+        entry->portContext->AddRef();
+        
+        entry->callback(entry->callbackContext);
+        
+        // Release portContext after callback completes
+        entry->portContext->Release();
+        
+        {
+            std::lock_guard<std::mutex> lock(m_terminationLock);
             m_terminationList->free_node(address);
-            delete entry;
-            return true;
         }
-
-        return false;
-    });
+        delete entry;
+    }
 }
 
 void TaskQueuePortImpl::ScheduleTermination(
@@ -1225,8 +1271,11 @@ void TaskQueuePortImpl::ScheduleTermination(
 
     // This never fails because we preallocate the
     // list node.
-    m_terminationList->push_back(term, term->node);
-    term->node = 0; // now owned by the list
+    {
+        std::lock_guard<std::mutex> lock(m_terminationLock);
+        m_terminationList->push_back(term, term->node);
+        term->node = 0; // now owned by the list
+    }
 
     // The port should have already been marked as terminated, so now
     // we can signal it to wake up. This should drain pending calls and
@@ -1236,6 +1285,12 @@ void TaskQueuePortImpl::ScheduleTermination(
     NotifyItemQueued();
 }
 
+bool TaskQueuePortImpl::TerminationListEmpty()
+{
+    std::lock_guard<std::mutex> lock(m_terminationLock);
+    return m_terminationList->empty();
+}
+
 #ifdef _WIN32
 void CALLBACK TaskQueuePortImpl::WaitCallback(
     _In_ PTP_CALLBACK_INSTANCE instance,

Source/Task/TaskQueueImpl.h

@@ -262,6 +262,7 @@ class TaskQueuePortImpl: public Api<ApiId::TaskQueuePort, ITaskQueuePort>
     std::unique_ptr<LocklessQueue<QueueEntry>> m_pendingList;
     std::unique_ptr<LocklessQueue<TerminationEntry*>> m_terminationList;
     std::unique_ptr<LocklessQueue<TerminationEntry*>> m_pendingTerminationList;
+    std::mutex m_terminationLock;
     OS::WaitTimer m_timer;
     OS::ThreadPool m_threadPool;
     std::atomic<uint64_t> m_timerDue = { UINT64_MAX };
@@ -312,6 +313,7 @@ class TaskQueuePortImpl: public Api<ApiId::TaskQueuePort, ITaskQueuePort>
 
     void SignalTerminations();
     void ScheduleTermination(_In_ TerminationEntry* term);
+    bool TerminationListEmpty();
 
     void SignalQueue();
     void NotifyItemQueued();