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

[vitalybuka]

Summary

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.

[vitalybuka]

rebase

[vitalybuka]

@dvyukov ThreadArgRetval::Join is alternative to ThreadRegistry::ConsumeThreadUserId+ThreadRegistry::SetThreadUserId

BTW. I tried to extend ThreadRegistry::live_ but it didn't look good.
Also we need these functionality for HWASAN without Registry anyway.
We could make ThreadArgRetval track pthread -> Tids mapping as well.
However live_ is tsan only and ThreadArgRetval is asan/lsan/hwasan only, so I decided to keep them separated for now.

[vitalybuka]

rebase

[Enna1]

compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp

21

call CheckLocked(); in CreateLocked ?

66

typo: Thread

compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h

55

Please keep singular and plural consistent, "Mark", "stores"

69

typo: thread

compiler-rt/lib/sanitizer_common/tests/sanitizer_thread_arg_retval_test.cpp

122

typo: thread

[thurston]

compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp

42–43

If it can't find the thread, why does it quietly return instead of aborting?
Other functions are stricter with 'CHECK(t);'.

46

Nit: 'retrieved'

[vitalybuka]

update

[eugenis]

compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h

87

isn't that already done, in the form of Lock/Unlock?

compiler-rt/lib/sanitizer_common/CMakeLists.txt

set(SANITIZER_SOURCES_NOTERMINATION
   sanitizer_procmaps_mac.cpp
   sanitizer_procmaps_solaris.cpp
   sanitizer_solaris.cpp
   sanitizer_stoptheworld_fuchsia.cpp
   sanitizer_stoptheworld_mac.cpp
   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
   )
 
 set(SANITIZER_SOURCES
   ${SANITIZER_SOURCES_NOTERMINATION}
   sanitizer_termination.cpp

compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.h

+//===-- 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

compiler-rt/lib/sanitizer_common/sanitizer_thread_arg_retval.cpp

+//===-- 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

compiler-rt/lib/sanitizer_common/tests/CMakeLists.txt

set(SANITIZER_UNITTESTS
   sanitizer_stack_store_test.cpp
   sanitizer_stackdepot_test.cpp
   sanitizer_stacktrace_printer_test.cpp
   sanitizer_stacktrace_test.cpp
   sanitizer_stoptheworld_test.cpp
   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
   )
 
 set(SANITIZER_TEST_HEADERS
   sanitizer_pthread_wrappers.h
   sanitizer_test_config.h

compiler-rt/lib/sanitizer_common/tests/sanitizer_thread_arg_retval_test.cpp

+//===-- 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