From: Yuxuan Shui yshui@codeweavers.com
And avoiding race condition while doing so.
Usually the threadpool holds a reference to the `work_item`, which is released when the `work_item`'s callback is invoked. However `queue_cancel_item` closes the threadpool object without releasing it, leading to a leak. The fix is not as simple as adding a `IUnknown_Release` to `queue_cancel_item`, because the `work_item`'s callback can still be called after CloseThreadpoolTimer/Wait has returned. In fact its callback might currently be running. In which case the callback will access freed memory if `queue_cancel_item` frees the `work_item`.
We have to stop any further callbacks to be queued, wait for any currently running callbacks to finish, then finally we can release the `work_item` if it hasn't already been freed during the wait. --- dlls/mfplat/tests/mfplat.c | 2 +- dlls/rtworkq/queue.c | 53 +++++++++++++++++++++++++++++++------- 2 files changed, 45 insertions(+), 10 deletions(-)
diff --git a/dlls/mfplat/tests/mfplat.c b/dlls/mfplat/tests/mfplat.c index 363c240d13d..1e751b425fa 100644 --- a/dlls/mfplat/tests/mfplat.c +++ b/dlls/mfplat/tests/mfplat.c @@ -3612,7 +3612,7 @@ static void test_scheduled_items(void) ok(hr == S_OK, "Failed to cancel item, hr %#lx.\n", hr);
refcount = IMFAsyncCallback_Release(&callback->IMFAsyncCallback_iface); - todo_wine ok(refcount == 0, "Unexpected refcount %lu.\n", refcount); + ok(refcount == 0, "Unexpected refcount %lu.\n", refcount);
hr = MFCancelWorkItem(key); ok(hr == MF_E_NOT_FOUND || broken(hr == S_OK) /* < win10 */, "Unexpected hr %#lx.\n", hr); diff --git a/dlls/rtworkq/queue.c b/dlls/rtworkq/queue.c index 00b77bf6953..a774ee4d58e 100644 --- a/dlls/rtworkq/queue.c +++ b/dlls/rtworkq/queue.c @@ -768,8 +768,9 @@ static void CALLBACK waiting_item_cancelable_callback(TP_CALLBACK_INSTANCE *inst
TRACE("result object %p.\n", item->result);
- if (queue_release_pending_item(item)) - invoke_async_callback(item->result); + queue_release_pending_item(item); + + invoke_async_callback(item->result);
IUnknown_Release(&item->IUnknown_iface); } @@ -791,8 +792,9 @@ static void CALLBACK scheduled_item_cancelable_callback(TP_CALLBACK_INSTANCE *in
TRACE("result object %p.\n", item->result);
- if (queue_release_pending_item(item)) - invoke_async_callback(item->result); + queue_release_pending_item(item); + + invoke_async_callback(item->result);
IUnknown_Release(&item->IUnknown_iface); } @@ -884,6 +886,8 @@ static HRESULT queue_submit_timer(struct queue *queue, IRtwqAsyncResult *result, static HRESULT queue_cancel_item(struct queue *queue, RTWQWORKITEM_KEY key) { HRESULT hr = RTWQ_E_NOT_FOUND; + TP_WAIT *wait_object; + TP_TIMER *timer_object; struct work_item *item;
EnterCriticalSection(&queue->cs); @@ -891,22 +895,53 @@ static HRESULT queue_cancel_item(struct queue *queue, RTWQWORKITEM_KEY key) { if (item->key == key) { + /* We can't immediately release the item here, because the callback could already be + * running somewhere else. And if we release it here, the callback will access freed memory. + * So instead we have to make sure the callback is really stopped, or has really finished + * running before we do that. And we can't do that in this critical section, which would be a + * deadlock. So we first keep an extra reference to it, then leave the critical section to + * wait for the thread-pool objects, finally we re-enter critical section to release it. */ key >>= 32; + IUnknown_AddRef(&item->IUnknown_iface); if ((key & WAIT_ITEM_KEY_MASK) == WAIT_ITEM_KEY_MASK) { - IRtwqAsyncResult_SetStatus(item->result, RTWQ_E_OPERATION_CANCELLED); - invoke_async_callback(item->result); - CloseThreadpoolWait(item->u.wait_object); + wait_object = item->u.wait_object; item->u.wait_object = NULL; + LeaveCriticalSection(&queue->cs); + + SetThreadpoolWait(wait_object, NULL, NULL); + WaitForThreadpoolWaitCallbacks(wait_object, TRUE); + CloseThreadpoolWait(wait_object); + + EnterCriticalSection(&queue->cs); } else if ((key & SCHEDULED_ITEM_KEY_MASK) == SCHEDULED_ITEM_KEY_MASK) { - CloseThreadpoolTimer(item->u.timer_object); + timer_object = item->u.timer_object; item->u.timer_object = NULL; + LeaveCriticalSection(&queue->cs); + + SetThreadpoolTimer(timer_object, NULL, 0, 0); + WaitForThreadpoolTimerCallbacks(timer_object, TRUE); + CloseThreadpoolTimer(timer_object); + + EnterCriticalSection(&queue->cs); } else WARN("Unknown item key mask %#I64x.\n", key); - queue_release_pending_item(item); + + if (queue_release_pending_item(item)) + { + // This means the callback wasn't run during our wait, so we can invoke the + // callback with a canceled status, and release the work item. + if ((key & WAIT_ITEM_KEY_MASK) == WAIT_ITEM_KEY_MASK) + { + IRtwqAsyncResult_SetStatus(item->result, RTWQ_E_OPERATION_CANCELLED); + invoke_async_callback(item->result); + } + IUnknown_Release(&item->IUnknown_iface); + } + IUnknown_Release(&item->IUnknown_iface); hr = S_OK; break; }