[PATCH v2 1/1] rtworkq: Fix leak of thread pool work items.

[Yuxuan Shui]

From: Yuxuan Shui

---
 dlls/rtworkq/queue.c | 30 +++++++++++++++++++++++++-----
 1 file changed, 25 insertions(+), 5 deletions(-)

diff --git a/dlls/rtworkq/queue.c b/dlls/rtworkq/queue.c
index 1f82eb7ad12..57e99de4e09 100644
--- a/dlls/rtworkq/queue.c
+++ b/dlls/rtworkq/queue.c
@@ -118,6 +118,13 @@ enum system_queue_index
     SYS_QUEUE_COUNT,
 };

+enum work_item_type
+{
+    WORK_ITEM_WORK,
+    WORK_ITEM_TIMER,
+    WORK_ITEM_WAIT,
+};
+
 struct work_item
 {
     IUnknown IUnknown_iface;
@@ -129,10 +136,11 @@ struct work_item
     RTWQWORKITEM_KEY key;
     LONG priority;
     DWORD flags;
-    TP_WORK *work_object;
     PTP_SIMPLE_CALLBACK finalization_callback;
+    enum work_item_type type;
     union
     {
+        TP_WORK *work_object;
         TP_WAIT *wait_object;
         TP_TIMER *timer_object;
     } u;
@@ -387,8 +395,9 @@ static void pool_queue_submit(struct queue *queue, struct work_item *item)
        we need finalization callback. */
     if (item->finalization_callback)
         IUnknown_AddRef(&item->IUnknown_iface);
-    item->work_object = CreateThreadpoolWork(standard_queue_worker, item, (TP_CALLBACK_ENVIRON *)&env);
-    SubmitThreadpoolWork(item->work_object);
+    item->u.work_object = CreateThreadpoolWork(standard_queue_worker, item, (TP_CALLBACK_ENVIRON *)&env);
+    item->type = WORK_ITEM_WORK;
+    SubmitThreadpoolWork(item->u.work_object);

TRACE("dispatched %p.\n", item->result);
 }
@@ -549,8 +558,17 @@ static ULONG WINAPI work_item_Release(IUnknown *iface)

if (!refcount)
     {
-        if (item->work_object)
-            CloseThreadpoolWork(item->work_object);
+        switch(item->type) {
+            case WORK_ITEM_WORK:
+                if (item->u.work_object) CloseThreadpoolWork(item->u.work_object);
+                break;
+            case WORK_ITEM_WAIT:
+                if (item->u.wait_object) CloseThreadpoolWait(item->u.wait_object);
+                break;
+            case WORK_ITEM_TIMER:
+                if (item->u.timer_object) CloseThreadpoolTimer(item->u.timer_object);
+                break;
+        }
         if (item->reply_result)
             IRtwqAsyncResult_Release(item->reply_result);
         IRtwqAsyncResult_Release(item->result);
@@ -814,6 +832,7 @@ static HRESULT queue_submit_wait(struct queue *queue, HANDLE event, LONG priorit

item->u.wait_object = CreateThreadpoolWait(callback, item,
             (TP_CALLBACK_ENVIRON *)&queue->envs[TP_CALLBACK_PRIORITY_NORMAL]);
+    item->type = WORK_ITEM_WAIT;
     SetThreadpoolWait(item->u.wait_object, event, NULL);

TRACE("dispatched %p.\n", result);
@@ -848,6 +867,7 @@ static HRESULT queue_submit_timer(struct queue *queue, IRtwqAsyncResult *result,

item->u.timer_object = CreateThreadpoolTimer(callback, item,
             (TP_CALLBACK_ENVIRON *)&queue->envs[TP_CALLBACK_PRIORITY_NORMAL]);
+    item->type = WORK_ITEM_TIMER;
     SetThreadpoolTimer(item->u.timer_object, &filetime, period, 0);

TRACE("dispatched %p.\n", result);