[ASan/Win] Fix a CHECK failure when an exception is thrown from a callback passed to BindIoCompletionCallback

This also simplifies how we handle QueueUserWorkItem

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@232499 91177308-0d34-0410-b5e6-96231b3b80d8

1 files changed, 30 insertions, 40 deletions

diff --git a/lib/asan/asan_win.cc b/lib/asan/asan_win.cc
index d80efaaae..addb3d40a 100644
--- a/lib/asan/asan_win.cc
+++ b/lib/asan/asan_win.cc
@@ -102,61 +102,51 @@ INTERCEPTOR_WINAPI(DWORD, CreateThread,
 }
 
 namespace {
-struct UserWorkItemInfo {
-  DWORD (__stdcall *function)(void *arg);
-  void *arg;
-  u32 parent_tid;
-};
-
 BlockingMutex mu_for_thread_tracking(LINKER_INITIALIZED);
 
-// QueueUserWorkItem may silently create a thread we should keep track of.
-// We achieve this by wrapping the user-supplied work items with our function.
-DWORD __stdcall QueueUserWorkItemWrapper(void *arg) {
-  UserWorkItemInfo *item = (UserWorkItemInfo *)arg;
-
-  {
-    // FIXME: GetCurrentThread relies on TSD, which might not play well with
-    // system thread pools.  We might want to use something like reference
-    // counting to zero out GetCurrentThread() underlying storage when the last
-    // work item finishes?  Or can we disable reclaiming of threads in the pool?
-    BlockingMutexLock l(&mu_for_thread_tracking);
-    AsanThread *t = __asan::GetCurrentThread();
-    if (!t) {
-      GET_STACK_TRACE_THREAD;
-      t = AsanThread::Create(/* start_routine */ nullptr, /* arg */ nullptr,
-                             item->parent_tid, &stack, /* detached */ true);
-      t->Init();
-      asanThreadRegistry().StartThread(t->tid(), 0, 0);
-      SetCurrentThread(t);
-    }
-  }
-
-  DWORD ret = item->function(item->arg);
-  delete item;
-  return ret;
+void EnsureWorkerThreadRegistered() {
+  // FIXME: GetCurrentThread relies on TSD, which might not play well with
+  // system thread pools.  We might want to use something like reference
+  // counting to zero out GetCurrentThread() underlying storage when the last
+  // work item finishes?  Or can we disable reclaiming of threads in the pool?
+  BlockingMutexLock l(&mu_for_thread_tracking);
+  if (__asan::GetCurrentThread())
+    return;
+
+  AsanThread *t = AsanThread::Create(
+      /* start_routine */ nullptr, /* arg */ nullptr,
+      /* parent_tid */ -1, /* stack */ nullptr, /* detached */ true);
+  t->Init();
+  asanThreadRegistry().StartThread(t->tid(), 0, 0);
+  SetCurrentThread(t);
 }
 }  // namespace
 
-INTERCEPTOR_WINAPI(BOOL, QueueUserWorkItem, LPTHREAD_START_ROUTINE function,
-                   PVOID arg, ULONG flags) {
-  UserWorkItemInfo *work_item_info = new UserWorkItemInfo;
-  work_item_info->function = function;
-  work_item_info->arg = arg;
-  work_item_info->parent_tid = GetCurrentTidOrInvalid();
-  return REAL(QueueUserWorkItem)(QueueUserWorkItemWrapper,
-                                 work_item_info, flags);
+INTERCEPTOR_WINAPI(DWORD, NtWaitForWorkViaWorkerFactory, DWORD a, DWORD b) {
+  // NtWaitForWorkViaWorkerFactory is called from system worker pool threads to
+  // query work scheduled by BindIoCompletionCallback, QueueUserWorkItem, etc.
+  // System worker pool threads are created at arbitraty point in time and
+  // without using CreateThread, so we wrap NtWaitForWorkViaWorkerFactory
+  // instead and don't register a specific parent_tid/stack.
+  EnsureWorkerThreadRegistered();
+  return REAL(NtWaitForWorkViaWorkerFactory)(a, b);
 }
+
 // }}}
 
 namespace __asan {
 
 void InitializePlatformInterceptors() {
   ASAN_INTERCEPT_FUNC(CreateThread);
-  ASAN_INTERCEPT_FUNC(QueueUserWorkItem);
   ASAN_INTERCEPT_FUNC(RaiseException);
   ASAN_INTERCEPT_FUNC(_except_handler3);
   ASAN_INTERCEPT_FUNC(_except_handler4);
+
+  // NtWaitForWorkViaWorkerFactory is always linked dynamically.
+  CHECK(::__interception::OverrideFunction(
+      "NtWaitForWorkViaWorkerFactory",
+      (uptr)WRAP(NtWaitForWorkViaWorkerFactory),
+      (uptr *)&REAL(NtWaitForWorkViaWorkerFactory)));
 }
 
 // ---------------------- TSD ---------------- {{{