[3/3] ntdll/tests: Add tests for RtlRegisterWait.

[Sebastian Lackner]

Signed-off-by: Sebastian Lackner <sebastian at fds-team.de>
---
 dlls/ntdll/tests/threadpool.c |  253 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 253 insertions(+)

diff --git a/dlls/ntdll/tests/threadpool.c b/dlls/ntdll/tests/threadpool.c
index a262ecb..ec0a29c 100644
--- a/dlls/ntdll/tests/threadpool.c
+++ b/dlls/ntdll/tests/threadpool.c
@@ -143,6 +143,258 @@ static void test_RtlQueueWorkItem(void)
     CloseHandle(semaphore);
 }
 
+struct rtl_wait_info
+{
+    HANDLE semaphore1;
+    HANDLE semaphore2;
+    DWORD wait_result;
+    LONG userdata;
+};
+
+static void CALLBACK rtl_wait_cb(void *userdata, BOOLEAN timeout)
+{
+    struct rtl_wait_info *info = userdata;
+    DWORD result;
+
+    trace("Running rtl_wait callback\n");
+
+    if (!timeout)
+        InterlockedIncrement(&info->userdata);
+    else
+        InterlockedExchangeAdd(&info->userdata, 0x10000);
+    ReleaseSemaphore(info->semaphore1, 1, NULL);
+
+    if (info->semaphore2)
+    {
+        result = WaitForSingleObject(info->semaphore2, 200);
+        ok(result == info->wait_result, "expected %u, got %u\n", info->wait_result, result);
+        ReleaseSemaphore(info->semaphore1, 1, NULL);
+    }
+}
+
+static void test_RtlRegisterWait(void)
+{
+    struct rtl_wait_info info;
+    HANDLE semaphores[2];
+    NTSTATUS status;
+    DWORD result;
+    HANDLE wait1;
+    HANDLE event;
+
+    semaphores[0] = CreateSemaphoreW(NULL, 0, 2, NULL);
+    ok(semaphores[0] != NULL, "failed to create semaphore\n");
+    semaphores[1] = CreateSemaphoreW(NULL, 0, 1, NULL);
+    ok(semaphores[1] != NULL, "failed to create semaphore\n");
+    info.semaphore1 = semaphores[0];
+    info.semaphore2 = NULL;
+
+    event = CreateEventW(NULL, FALSE, FALSE, NULL);
+    ok(event != NULL, "failed to create event\n");
+
+    /* basic test for RtlRegisterWait and RtlDeregisterWait */
+    wait1 = NULL;
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEDEFAULT);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    ok(wait1 != NULL, "expected wait1 != NULL\n");
+    status = RtlDeregisterWait(wait1);
+    ok(!status, "RtlDeregisterWait failed with status %x\n", status);
+    ok(info.userdata == 0, "expected info.userdata = 0, got %u\n", info.userdata);
+
+    /* infinite timeout, signal the semaphore two times */
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEDEFAULT);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 1, "expected info.userdata = 1, got %u\n", info.userdata);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 2, "expected info.userdata = 2, got %u\n", info.userdata);
+    result = WaitForSingleObject(semaphores[1], 0);
+    ok(result == WAIT_TIMEOUT, "WaitForSingleObject returned %u\n", result);
+    status = RtlDeregisterWait(wait1);
+    ok(!status || broken(status == STATUS_PENDING) /* >= Win 8 */,
+       "RtlDeregisterWait failed with status %x\n", status);
+
+    /* repeat test with WT_EXECUTEONLYONCE */
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEONLYONCE);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 1, "expected info.userdata = 1, got %u\n", info.userdata);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_TIMEOUT, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 1, "expected info.userdata = 1, got %u\n", info.userdata);
+    result = WaitForSingleObject(semaphores[1], 0);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    status = RtlDeregisterWait(wait1);
+    ok(!status, "RtlDeregisterWait failed with status %x\n", status);
+
+    /* finite timeout, no event */
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, 200, WT_EXECUTEDEFAULT);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_TIMEOUT, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 0, "expected info.userdata = 0, got %u\n", info.userdata);
+    result = WaitForSingleObject(semaphores[0], 200);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 0x10000, "expected info.userdata = 0x10000, got %u\n", info.userdata);
+    result = WaitForSingleObject(semaphores[1], 0);
+    ok(result == WAIT_TIMEOUT, "WaitForSingleObject returned %u\n", result);
+    status = RtlDeregisterWait(wait1);
+    ok(!status || broken(status == STATUS_PENDING) /* >= Win 8 */,
+       "RtlDeregisterWait failed with status %x\n", status);
+
+    /* finite timeout, with event */
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, 200, WT_EXECUTEDEFAULT);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_TIMEOUT, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 0, "expected info.userdata = 0, got %u\n", info.userdata);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 1, "expected info.userdata = 1, got %u\n", info.userdata);
+    result = WaitForSingleObject(semaphores[1], 0);
+    ok(result == WAIT_TIMEOUT, "WaitForSingleObject returned %u\n", result);
+    status = RtlDeregisterWait(wait1);
+    ok(!status || broken(status == STATUS_PENDING) /* >= Win 8 */,
+       "RtlDeregisterWait failed with status %x\n", status);
+
+    /* test RtlDeregisterWaitEx before wait expired */
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEDEFAULT);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    status = RtlDeregisterWaitEx(wait1, NULL);
+    ok(!status, "RtlDeregisterWaitEx failed with status %x\n", status);
+    ok(info.userdata == 0, "expected info.userdata = 0, got %u\n", info.userdata);
+
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEDEFAULT);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    status = RtlDeregisterWaitEx(wait1, INVALID_HANDLE_VALUE);
+    ok(!status, "RtlDeregisterWaitEx failed with status %x\n", status);
+    ok(info.userdata == 0, "expected info.userdata = 0, got %u\n", info.userdata);
+
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEDEFAULT);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    status = RtlDeregisterWaitEx(wait1, event);
+    ok(!status, "RtlDeregisterWaitEx failed with status %x\n", status);
+    ok(info.userdata == 0, "expected info.userdata = 0, got %u\n", info.userdata);
+    result = WaitForSingleObject(event, 200);
+    todo_wine
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+
+    /* test RtlDeregisterWaitEx after wait expired */
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, 0, WT_EXECUTEONLYONCE);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 0x10000, "expected info.userdata = 0x10000, got %u\n", info.userdata);
+    Sleep(50);
+    status = RtlDeregisterWaitEx(wait1, NULL);
+    ok(!status, "RtlDeregisterWaitEx failed with status %x\n", status);
+    ok(info.userdata == 0x10000, "expected info.userdata = 0x10000, got %u\n", info.userdata);
+
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, 0, WT_EXECUTEONLYONCE);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 0x10000, "expected info.userdata = 0x10000, got %u\n", info.userdata);
+    Sleep(50);
+    status = RtlDeregisterWaitEx(wait1, INVALID_HANDLE_VALUE);
+    ok(!status, "RtlDeregisterWaitEx failed with status %x\n", status);
+    ok(info.userdata == 0x10000, "expected info.userdata = 0x10000, got %u\n", info.userdata);
+
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, 0, WT_EXECUTEONLYONCE);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    result = WaitForSingleObject(semaphores[0], 100);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 0x10000, "expected info.userdata = 0x10000, got %u\n", info.userdata);
+    Sleep(50);
+    status = RtlDeregisterWaitEx(wait1, event);
+    ok(!status, "RtlDeregisterWaitEx failed with status %x\n", status);
+    ok(info.userdata == 0x10000, "expected info.userdata = 0x10000, got %u\n", info.userdata);
+    result = WaitForSingleObject(event, 200);
+    todo_wine
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+
+    /* test RtlDeregisterWaitEx while callback is running */
+    info.semaphore2 = semaphores[1];
+    info.wait_result = WAIT_OBJECT_0;
+
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEONLYONCE);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 1000);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 1, "expected info.userdata = 1, got %u\n", info.userdata);
+    status = RtlDeregisterWait(wait1);
+    ok(status == STATUS_PENDING, "expected STATUS_PENDING, got %x\n", status);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 1000);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEONLYONCE);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 1000);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 1, "expected info.userdata = 1, got %u\n", info.userdata);
+    status = RtlDeregisterWaitEx(wait1, NULL);
+    ok(status == STATUS_PENDING, "expected STATUS_PENDING, got %x\n", status);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 1000);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+
+    info.wait_result = WAIT_TIMEOUT;
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEONLYONCE);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 1000);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 1, "expected info.userdata = 1, got %u\n", info.userdata);
+    status = RtlDeregisterWaitEx(wait1, INVALID_HANDLE_VALUE);
+    ok(!status, "RtlDeregisterWaitEx failed with status %x\n", status);
+    result = WaitForSingleObject(semaphores[0], 0);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+
+    info.wait_result = WAIT_OBJECT_0;
+    info.userdata = 0;
+    status = RtlRegisterWait(&wait1, semaphores[1], rtl_wait_cb, &info, INFINITE, WT_EXECUTEONLYONCE);
+    ok(!status, "RtlRegisterWait failed with status %x\n", status);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(semaphores[0], 1000);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    ok(info.userdata == 1, "expected info.userdata = 1, got %u\n", info.userdata);
+    status = RtlDeregisterWaitEx(wait1, event);
+    ok(status == STATUS_PENDING, "expected STATUS_PENDING, got %x\n", status);
+    ReleaseSemaphore(semaphores[1], 1, NULL);
+    result = WaitForSingleObject(event, 1000);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+    result = WaitForSingleObject(semaphores[0], 0);
+    ok(result == WAIT_OBJECT_0, "WaitForSingleObject returned %u\n", result);
+
+    CloseHandle(semaphores[0]);
+    CloseHandle(semaphores[1]);
+    CloseHandle(event);
+}
+
 static void CALLBACK simple_cb(TP_CALLBACK_INSTANCE *instance, void *userdata)
 {
     HANDLE semaphore = userdata;
@@ -1340,6 +1592,7 @@ static void test_tp_multi_wait(void)
 START_TEST(threadpool)
 {
     test_RtlQueueWorkItem();
+    test_RtlRegisterWait();
 
     if (!init_threadpool())
         return;
-- 
2.6.0