From cf243818996ca12304cad2ba5cee4ec3495d8878 Mon Sep 17 00:00:00 2001 From: Yuxuan Shui Date: Wed, 1 Nov 2023 02:05:38 +0000 Subject: [PATCH] rtworkq: Release cancelled work items. Usually the threadpool holds a reference to the `work_item`, which is released when the `work_item`'s callback is invoked. On the other hand, `queue_cancel_item` closes the threadpool object without releasing the `work_item`. So if the callbacks don't get a chance to run - which is not guaranteed - the `work_item` will be leaked. 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/mf/tests/mf.c | 1 - dlls/mfplat/tests/mfplat.c | 4 ++-- dlls/rtworkq/queue.c | 48 +++++++++++++++++++++++++++++++------- 3 files changed, 41 insertions(+), 12 deletions(-) diff --git a/dlls/mf/tests/mf.c b/dlls/mf/tests/mf.c index 8ba199557a1..fa5b0f55391 100644 --- a/dlls/mf/tests/mf.c +++ b/dlls/mf/tests/mf.c @@ -4562,7 +4562,6 @@ if (SUCCEEDED(hr)) check_sar_rate_support(sink); ref = IMFMediaSink_Release(sink); - todo_wine ok(ref == 0, "Release returned %ld\n", ref); /* Activation */ diff --git a/dlls/mfplat/tests/mfplat.c b/dlls/mfplat/tests/mfplat.c index 5cd9d28f074..a4ca567367c 100644 --- a/dlls/mfplat/tests/mfplat.c +++ b/dlls/mfplat/tests/mfplat.c @@ -3613,7 +3613,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); @@ -3760,7 +3760,7 @@ static void test_periodic_callback(void) hr = pMFRemovePeriodicCallback(key); ok(hr == S_OK, "Failed to remove callback, hr %#lx.\n", hr); Sleep(500); - todo_wine EXPECT_REF(&context->IMFAsyncCallback_iface, 1); + EXPECT_REF(&context->IMFAsyncCallback_iface, 1); IMFAsyncCallback_Release(&context->IMFAsyncCallback_iface); hr = MFShutdown(); diff --git a/dlls/rtworkq/queue.c b/dlls/rtworkq/queue.c index 00b77bf6953..9046a70b359 100644 --- a/dlls/rtworkq/queue.c +++ b/dlls/rtworkq/queue.c @@ -883,7 +883,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,29 +892,58 @@ 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); } 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); } else + { WARN("Unknown item key mask %#I64x.\n", key); - queue_release_pending_item(item); - hr = S_OK; - break; + LeaveCriticalSection(&queue->cs); + } + + 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); + return S_OK; } } LeaveCriticalSection(&queue->cs); - return hr; + return RTWQ_E_NOT_FOUND; } static HRESULT alloc_user_queue(const struct queue_desc *desc, DWORD *queue_id)