[NFC][sanitizer] Add class to track thread arg and retval

We need something to keep arg and retval pointers for leak checking.
Pointers should keept alive even after thread exited, until the thread
is detached or joined.
We should not put this logic into ThreadRegistry as we need the the
same for the ThreadList of HWASAN.

Reviewed By: thurston

Differential Revision: https://reviews.llvm.org/D150104

5 files changed, 369 insertions, 0 deletions

diff --git a/compiler-rt/lib/sanitizer_common/CMakeLists.txt b/compiler-rt/lib/sanitizer_common/CMakeLists.txt
index 9bc93abc553a..614e63d5f9eb 100644
--- a/compiler-rt/lib/sanitizer_common/CMakeLists.txt
+++ b/compiler-rt/lib/sanitizer_common/CMakeLists.txt
@@ -37,6 +37,7 @@ set(SANITIZER_SOURCES_NOTERMINATION
   sanitizer_stoptheworld_win.cpp
   sanitizer_suppressions.cpp
   sanitizer_tls_get_addr.cpp
+  sanitizer_thread_arg_retval.cpp
   sanitizer_thread_registry.cpp
   sanitizer_type_traits.cpp
   sanitizer_win.cpp
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp b/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp
new file mode 100644
index 000000000000..fd53208bd27a
--- /dev/null
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp
@@ -0,0 +1,93 @@
+//===-- sanitizer_thread_arg_retval.cpp -------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is shared between sanitizer tools.
+//
+// Tracks thread arguments and return value for leak checking.
+//===----------------------------------------------------------------------===//
+
+#include "sanitizer_thread_arg_retval.h"
+
+#include "sanitizer_placement_new.h"
+
+namespace __sanitizer {
+
+void ThreadArgRetval::CreateLocked(uptr thread, bool detached,
+                                   const Args& args) {
+  CheckLocked();
+  Data& t = data_[thread];
+  t = {};
+  t.gen = gen_++;
+  t.detached = detached;
+  t.args = args;
+}
+
+ThreadArgRetval::Args ThreadArgRetval::GetArgs(uptr thread) const {
+  __sanitizer::Lock lock(&mtx_);
+  auto t = data_.find(thread);
+  CHECK(t);
+  if (t->second.done)
+    return {};
+  return t->second.args;
+}
+
+void ThreadArgRetval::Finish(uptr thread, void* retval) {
+  __sanitizer::Lock lock(&mtx_);
+  auto t = data_.find(thread);
+  CHECK(t);
+  if (t->second.detached) {
+    // Retval of detached thread connot be retrieved.
+    data_.erase(t);
+    return;
+  }
+  t->second.done = true;
+  t->second.args.arg_retval = retval;
+}
+
+u32 ThreadArgRetval::BeforeJoin(uptr thread) const {
+  __sanitizer::Lock lock(&mtx_);
+  auto t = data_.find(thread);
+  CHECK(t);
+  CHECK(!t->second.detached);
+  return t->second.gen;
+}
+
+void ThreadArgRetval::AfterJoin(uptr thread, u32 gen) {
+  __sanitizer::Lock lock(&mtx_);
+  auto t = data_.find(thread);
+  if (!t || gen != t->second.gen) {
+    // Thread was reused and erased by any other event.
+    return;
+  }
+  CHECK(!t->second.detached);
+  data_.erase(t);
+}
+
+void ThreadArgRetval::DetachLocked(uptr thread) {
+  CheckLocked();
+  auto t = data_.find(thread);
+  CHECK(t);
+  CHECK(!t->second.detached);
+  if (t->second.done) {
+    // Detached thread has no use after it started and returned args.
+    data_.erase(t);
+    return;
+  }
+  t->second.detached = true;
+}
+
+void ThreadArgRetval::GetAllPtrsLocked(InternalMmapVector<uptr>* ptrs) {
+  CheckLocked();
+  CHECK(ptrs);
+  data_.forEach([&](DenseMap<uptr, Data>::value_type& kv) -> bool {
+    ptrs->push_back((uptr)kv.second.args.arg_retval);
+    return true;
+  });
+}
+
+}  // namespace __sanitizer
diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h b/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h
new file mode 100644
index 000000000000..f2e5af208474
--- /dev/null
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h
@@ -0,0 +1,116 @@
+//===-- sanitizer_thread_arg_retval.h ---------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is shared between sanitizer tools.
+//
+// Tracks thread arguments and return value for leak checking.
+//===----------------------------------------------------------------------===//
+
+#ifndef SANITIZER_THREAD_ARG_RETVAL_H
+#define SANITIZER_THREAD_ARG_RETVAL_H
+
+#include "sanitizer_common.h"
+#include "sanitizer_dense_map.h"
+#include "sanitizer_list.h"
+#include "sanitizer_mutex.h"
+
+namespace __sanitizer {
+
+// Primary goal of the class is to keep alive arg and retval pointer for leak
+// checking. However it can be used to pass those pointer into wrappers used by
+// interceptors. The difference from ThreadRegistry/ThreadList is that this
+// class keeps data up to the detach or join, as exited thread still can be
+// joined to retrive retval. ThreadRegistry/ThreadList can discard exited
+// threads immediately.
+class SANITIZER_MUTEX ThreadArgRetval {
+ public:
+  struct Args {
+    void* (*routine)(void*);
+    void* arg_retval;  // Either arg or retval.
+  };
+  void Lock() SANITIZER_ACQUIRE() { mtx_.Lock(); }
+  void CheckLocked() const SANITIZER_CHECK_LOCKED() { mtx_.CheckLocked(); }
+  void Unlock() SANITIZER_RELEASE() { mtx_.Unlock(); }
+
+  // Wraps pthread_create or similar. We need to keep object locked, to
+  // prevent child thread from proceeding without thread handle.
+  template <typename CreateFn /* returns thread id on success, or 0 */>
+  void Create(bool detached, const Args& args, const CreateFn& fn) {
+    // No need to track detached threads with no args, but we will to do as it's
+    // not expensive and less edge-cases.
+    __sanitizer::Lock lock(&mtx_);
+    if (uptr thread = fn())
+      CreateLocked(thread, detached, args);
+  }
+
+  // Returns thread arg and routine.
+  Args GetArgs(uptr thread) const;
+
+  // Mark thread as done and stores retval or remove if detached. Should be
+  // called by the thread.
+  void Finish(uptr thread, void* retval);
+
+  // Mark thread as detached or remove if done.
+  template <typename DetachFn /* returns true on success */>
+  void Detach(uptr thread, const DetachFn& fn) {
+    // Lock to prevent re-use of the thread between fn() and DetachLocked()
+    // calls.
+    __sanitizer::Lock lock(&mtx_);
+    if (fn())
+      DetachLocked(thread);
+  }
+
+  // Joins the thread.
+  template <typename JoinFn /* returns true on success */>
+  void Join(uptr thread, const JoinFn& fn) {
+    // Remember internal id of the thread to prevent re-use of the thread
+    // between fn() and DetachLocked() calls. We can't just lock JoinFn
+    // like in Detach() implementation.
+    auto gen = BeforeJoin(thread);
+    if (fn())
+      AfterJoin(thread, gen);
+  }
+
+  // Returns all arg and retval which are considered alive.
+  void GetAllPtrsLocked(InternalMmapVector<uptr>* ptrs);
+
+  uptr size() const {
+    __sanitizer::Lock lock(&mtx_);
+    return data_.size();
+  }
+
+  // FIXME: Add fork support. Expected users of the class are sloppy with forks
+  // anyway. We likely should lock/unlock the object to avoid deadlocks, and
+  // erase all but the current threads, so we can detect leaked arg or retval in
+  // child process.
+
+  // FIXME: Add cancelation support. Now if a thread was canceled, the class
+  // will keep pointers alive forever, missing leaks caused by cancelation.
+
+ private:
+  struct Data {
+    Args args;
+    u32 gen;  // Avoid collision if thread id re-used.
+    bool detached;
+    bool done;
+  };
+
+  void CreateLocked(uptr thread, bool detached, const Args& args);
+  u32 BeforeJoin(uptr thread) const;
+  void AfterJoin(uptr thread, u32 gen);
+  void DetachLocked(uptr thread);
+
+  mutable Mutex mtx_;
+
+  DenseMap<uptr, Data> data_;
+  u32 gen_ = 0;
+};
+
+}  // namespace __sanitizer
+
+#endif  // SANITIZER_THREAD_ARG_RETVAL_H
diff --git a/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt b/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt
index 22bad618a8ff..40aa8e703b6c 100644
--- a/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt
+++ b/compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt
@@ -46,6 +46,7 @@ set(SANITIZER_UNITTESTS
   sanitizer_suppressions_test.cpp
   sanitizer_symbolizer_test.cpp
   sanitizer_test_main.cpp
+  sanitizer_thread_arg_retval_test.cpp
   sanitizer_thread_registry_test.cpp
   sanitizer_type_traits_test.cpp
   sanitizer_vector_test.cpp
diff --git a/compiler-rt/lib/sanitizer_common/tests/sanitizer_thread_arg_retval_test.cpp b/compiler-rt/lib/sanitizer_common/tests/sanitizer_thread_arg_retval_test.cpp
new file mode 100644
index 000000000000..31507c8df19c
--- /dev/null
+++ b/compiler-rt/lib/sanitizer_common/tests/sanitizer_thread_arg_retval_test.cpp
@@ -0,0 +1,158 @@
+//===-- sanitizer_thread_registry_test.cpp --------------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of shared sanitizer runtime.
+//
+//===----------------------------------------------------------------------===//
+#include "sanitizer_common/sanitizer_thread_arg_retval.h"
+
+#include "gtest/gtest.h"
+#include "sanitizer_mutex.h"
+
+namespace __sanitizer {
+
+static int t;
+static void* arg = &t;
+static void* (*routine)(void*) = [](void*) { return arg; };
+static void* retval = (&t) + 1;
+
+TEST(ThreadArgRetvalTest, CreateFail) {
+  ThreadArgRetval td;
+  td.Create(false, {routine, arg}, []() { return 0; });
+  EXPECT_EQ(0u, td.size());
+}
+
+TEST(ThreadArgRetvalTest, CreateRunJoin) {
+  ThreadArgRetval td;
+  td.Create(false, {routine, arg}, []() { return 1; });
+  EXPECT_EQ(1u, td.size());
+
+  EXPECT_EQ(arg, td.GetArgs(1).arg_retval);
+  EXPECT_EQ(routine, td.GetArgs(1).routine);
+  EXPECT_EQ(1u, td.size());
+
+  td.Finish(1, retval);
+  EXPECT_EQ(1u, td.size());
+
+  td.Join(1, []() { return false; });
+  EXPECT_EQ(1u, td.size());
+
+  td.Join(1, []() { return true; });
+  EXPECT_EQ(0u, td.size());
+}
+
+TEST(ThreadArgRetvalTest, CreateJoinRun) {
+  ThreadArgRetval td;
+  td.Create(false, {routine, arg}, []() { return 1; });
+  EXPECT_EQ(1u, td.size());
+
+  td.Join(1, []() { return false; });
+  EXPECT_EQ(1u, td.size());
+
+  td.Join(1, [&]() {
+    // Expected to happen on another thread.
+    EXPECT_EQ(1u, td.size());
+
+    EXPECT_EQ(arg, td.GetArgs(1).arg_retval);
+    EXPECT_EQ(routine, td.GetArgs(1).routine);
+    EXPECT_EQ(1u, td.size());
+
+    td.Finish(1, retval);
+    EXPECT_EQ(1u, td.size());
+    return true;
+  });
+  EXPECT_EQ(0u, td.size());
+}
+
+TEST(ThreadArgRetvalTest, CreateRunDetach) {
+  ThreadArgRetval td;
+  td.Create(false, {routine, arg}, []() { return 1; });
+  EXPECT_EQ(1u, td.size());
+
+  EXPECT_EQ(arg, td.GetArgs(1).arg_retval);
+  EXPECT_EQ(routine, td.GetArgs(1).routine);
+  EXPECT_EQ(1u, td.size());
+
+  td.Finish(1, retval);
+  EXPECT_EQ(1u, td.size());
+
+  td.Detach(1, []() { return false; });
+  EXPECT_EQ(1u, td.size());
+
+  td.Detach(1, []() { return true; });
+  EXPECT_EQ(0u, td.size());
+}
+
+TEST(ThreadArgRetvalTest, CreateDetachRun) {
+  ThreadArgRetval td;
+  td.Create(false, {routine, arg}, []() { return 1; });
+  EXPECT_EQ(1u, td.size());
+
+  td.Detach(1, []() { return true; });
+  EXPECT_EQ(1u, td.size());
+
+  EXPECT_EQ(arg, td.GetArgs(1).arg_retval);
+  EXPECT_EQ(routine, td.GetArgs(1).routine);
+  EXPECT_EQ(1u, td.size());
+
+  td.Finish(1, retval);
+  EXPECT_EQ(0u, td.size());
+}
+
+TEST(ThreadArgRetvalTest, CreateRunJoinReuse) {
+  ThreadArgRetval td;
+  td.Create(false, {routine, arg}, []() { return 1; });
+  EXPECT_EQ(1u, td.size());
+
+  td.Finish(1, retval);
+  EXPECT_EQ(1u, td.size());
+
+  td.Join(1, [&]() {
+    // Reuse thread id.
+    td.Create(false, {routine, arg}, []() { return 1; });
+    EXPECT_EQ(1u, td.size());
+    return true;
+  });
+  // Even if JoinFn succeeded, we can't erase mismatching thread.
+  EXPECT_EQ(1u, td.size());
+
+  // Now we can join another one.
+  td.Join(1, []() { return true; });
+  EXPECT_EQ(0u, td.size());
+}
+
+TEST(ThreadArgRetvalTest, GetAllPtrsLocked) {
+  ThreadArgRetval td;
+  td.Create(false, {routine, arg}, []() { return 1; });
+  {
+    GenericScopedLock<ThreadArgRetval> lock(&td);
+    InternalMmapVector<uptr> ptrs;
+    td.GetAllPtrsLocked(&ptrs);
+    EXPECT_EQ(1u, ptrs.size());
+    EXPECT_EQ((uptr)arg, ptrs[0]);
+  }
+
+  td.Finish(1, retval);
+  {
+    GenericScopedLock<ThreadArgRetval> lock(&td);
+    InternalMmapVector<uptr> ptrs;
+    td.GetAllPtrsLocked(&ptrs);
+    EXPECT_EQ(1u, ptrs.size());
+    EXPECT_EQ((uptr)retval, ptrs[0]);
+  }
+
+  td.Join(1, []() { return true; });
+  {
+    GenericScopedLock<ThreadArgRetval> lock(&td);
+    InternalMmapVector<uptr> ptrs;
+    td.GetAllPtrsLocked(&ptrs);
+    EXPECT_TRUE(ptrs.empty());
+  }
+}
+
+}  // namespace __sanitizer