driver: queueing: don't wait on more than 64 cores at once

It turns out we can't wait on more than 64 handles at once, which
happens on monster systems with 128 cores. So, split this into chunks.
While we're at it, make both the normal and low-mem path use the same
logic, with the low-mem path simply doing 3 threads at once.

Reported-by: Joe Mulvihill <Joe.Mulvihill@Hardsuitlabs.com>
Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>

1 files changed, 19 insertions, 34 deletions

diff --git a/driver/queueing.c b/driver/queueing.c
index 029f253..c1b7afe 100644
--- a/driver/queueing.c
+++ b/driver/queueing.c
@@ -115,47 +115,32 @@ MulticoreWorkQueueDestroy(MULTICORE_WORKQUEUE *WorkQueue)
     KeWaitForSingleObject(WorkQueue->WorkerSpawnerThread, Executive, KernelMode, FALSE, NULL);
     ObDereferenceObject(WorkQueue->WorkerSpawnerThread);
 
-    ULONG ThreadCount = 0;
-    for (MULTICORE_WORKTHREAD *Thread = WorkQueue->FirstThread; Thread; Thread = Thread->NextThread)
-    {
-        if (Thread->Thread)
-            ++ThreadCount;
-    }
-    if (!ThreadCount)
-        return;
-    PKTHREAD *Threads = MemAllocateArray(ThreadCount, sizeof(*Threads) + sizeof(KWAIT_BLOCK));
-    if (Threads)
+    PKTHREAD Threads[MAXIMUM_WAIT_OBJECTS];
+    static_assert(sizeof(Threads) <= 512, "Must move thread handle allocation to heap");
+    PKWAIT_BLOCK WaitBlock = MemAllocateArray(MAXIMUM_WAIT_OBJECTS, sizeof(KWAIT_BLOCK));
+    ULONG MaxPerWait = WaitBlock ? MAXIMUM_WAIT_OBJECTS : THREAD_WAIT_OBJECTS;
+
+    MULTICORE_WORKTHREAD *Thread = WorkQueue->FirstThread, *Next;
+    while (Thread)
     {
-        PKWAIT_BLOCK WaitBlock = (PKWAIT_BLOCK)((ULONG_PTR)Threads + (ThreadCount * sizeof(*Threads)));
-        ThreadCount = 0;
-        for (MULTICORE_WORKTHREAD *Thread = WorkQueue->FirstThread; Thread; Thread = Thread->NextThread)
+        ULONG Count = 0;
+        for (; Thread; Thread = Thread->NextThread)
         {
+            if (Count >= MaxPerWait)
+                break;
             if (Thread->Thread)
-                Threads[ThreadCount++] = Thread->Thread;
-        }
-        KeWaitForMultipleObjects(ThreadCount, Threads, WaitAll, Executive, KernelMode, FALSE, NULL, WaitBlock);
-        for (MULTICORE_WORKTHREAD *Thread = WorkQueue->FirstThread, *Next; Thread; Thread = Next)
-        {
-            Next = Thread->NextThread;
-            if (Thread->Thread)
-                ObDereferenceObject(Thread->Thread);
-            MemFree(Thread);
+                Threads[Count++] = Thread->Thread;
         }
-        MemFree(Threads);
+        KeWaitForMultipleObjects(Count, Threads, WaitAll, Executive, KernelMode, FALSE, NULL, WaitBlock);
     }
-    else
+    for (Thread = WorkQueue->FirstThread; Thread; Thread = Next)
     {
-        for (MULTICORE_WORKTHREAD *Thread = WorkQueue->FirstThread, *Next; Thread; Thread = Next)
-        {
-            Next = Thread->NextThread;
-            if (Thread->Thread)
-            {
-                KeWaitForSingleObject(Thread->Thread, Executive, KernelMode, FALSE, NULL);
-                ObDereferenceObject(Thread->Thread);
-            }
-            MemFree(Thread);
-        }
+        Next = Thread->NextThread;
+        if (Thread->Thread)
+            ObDereferenceObject(Thread->Thread);
+        MemFree(Thread);
     }
+    MemFree(WaitBlock);
 }
 
 #define NEXT(Nbl) NET_BUFFER_LIST_PER_PEER_LIST_LINK(Nbl)