Skip to content

Commit

Permalink
rtworkq: Release cancelled work items.
Browse files Browse the repository at this point in the history
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, 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 we would double release the reference.

We have to stop any further callbacks to be queued, wait for any currently running callbacks to
finish, then finally we can close the threadpool object. After that, if the callback wasn't called,
we can safely release the work_item reference.
  • Loading branch information
Yuxuan Shui committed Nov 6, 2023
1 parent 5efa97e commit 3f088c2
Showing 1 changed file with 64 additions and 9 deletions.
73 changes: 64 additions & 9 deletions dlls/rtworkq/queue.c
Original file line number Diff line number Diff line change
Expand Up @@ -880,32 +880,87 @@ static HRESULT queue_submit_timer(struct queue *queue, IRtwqAsyncResult *result,
return S_OK;
}

static HRESULT queue_cancel_item(struct queue *queue, RTWQWORKITEM_KEY key)
static HRESULT queue_cancel_item(struct queue *queue, const RTWQWORKITEM_KEY key)
{
HRESULT hr = RTWQ_E_NOT_FOUND;
union { TP_WAIT *wait_object; TP_TIMER *timer_object; } work_object;
enum work_item_type work_object_type;
struct work_item *item;
const UINT64 mask = key >> 32;

EnterCriticalSection(&queue->cs);
LIST_FOR_EACH_ENTRY(item, &queue->pending_items, struct work_item, entry)
{
if (item->key == key)
{
key >>= 32;
if ((key & WAIT_ITEM_KEY_MASK) == WAIT_ITEM_KEY_MASK)
hr = S_OK;
if ((mask & WAIT_ITEM_KEY_MASK) == WAIT_ITEM_KEY_MASK)
{
if (item->type != WORK_ITEM_WAIT)
WARN("Item %p is not a wait item, but its key has wait item mask.\n", item);

work_object_type = WORK_ITEM_WAIT;
work_object.wait_object = item->u.wait_object;
item->u.wait_object = NULL;
}
else if ((mask & SCHEDULED_ITEM_KEY_MASK) == SCHEDULED_ITEM_KEY_MASK)
{
if (item->type != WORK_ITEM_TIMER)
WARN("Item %p is not a timer item, but its key has timer item mask.\n", item);

work_object_type = WORK_ITEM_TIMER;
work_object.timer_object = item->u.timer_object;
item->u.timer_object = NULL;
}
else
{
WARN("Unknown item key mask %#I64x.\n", mask);
queue_release_pending_item(item);
goto out;
}
break;
}
}

if (FAILED(hr))
goto out;

LeaveCriticalSection(&queue->cs);

// Safely either stop the thread pool object, or if the callback is already running, wait for it to finish.
// This way, we can safely release the reference to the work item.
if (work_object_type == WORK_ITEM_WAIT)
{
SetThreadpoolWait(work_object.wait_object, NULL, NULL);
WaitForThreadpoolWaitCallbacks(work_object.wait_object, TRUE);
CloseThreadpoolWait(work_object.wait_object);
}
else if (work_object_type == WORK_ITEM_TIMER)
{
SetThreadpoolTimer(work_object.timer_object, NULL, 0, 0);
WaitForThreadpoolTimerCallbacks(work_object.timer_object, TRUE);
CloseThreadpoolTimer(work_object.timer_object);
}

// If the work item is still in pending items, its callback hasn't been invoked yet;
// we remove it. Otherwise its callback would have already released it.
EnterCriticalSection(&queue->cs);
LIST_FOR_EACH_ENTRY(item, &queue->pending_items, struct work_item, entry)
{
if (item->key == key)
{
if (work_object_type == WORK_ITEM_WAIT)
{
IRtwqAsyncResult_SetStatus(item->result, RTWQ_E_OPERATION_CANCELLED);
invoke_async_callback(item->result);
CloseThreadpoolWait(item->u.wait_object);
}
else if ((key & SCHEDULED_ITEM_KEY_MASK) == SCHEDULED_ITEM_KEY_MASK)
CloseThreadpoolTimer(item->u.timer_object);
else
WARN("Unknown item key mask %#I64x.\n", key);
queue_release_pending_item(item);
hr = S_OK;
IUnknown_Release(&item->IUnknown_iface);
break;
}
}

out:
LeaveCriticalSection(&queue->cs);

return hr;
Expand Down

0 comments on commit 3f088c2

Please sign in to comment.