[Threading][TSan] Fix TSan errors from lazy init on Linux.

Move the TSan functionality from Concurrency into Threading. Use it
in the Linux ulock implementation so that TSan knows about ulock
and will tolerate the newer swift_once implementation that uses it.

This should also slightly improve the performance of Concurrency,
since the tsan tests will be inlined rather than it always doing a
subroutine call.

rdar://110665213

Author: al45tair

cmake/modules/AddSwiftUnittests.cmake

@@ -58,7 +58,8 @@ function(add_swift_unittest test_dirname)
 
   string(TOUPPER "${SWIFT_SDK_${SWIFT_HOST_VARIANT_SDK}_THREADING_PACKAGE}" _threading_package)
   target_compile_definitions("${test_dirname}" PRIVATE
-    "SWIFT_THREADING_${_threading_package}")
+    "SWIFT_THREADING_${_threading_package}"
+    "SWIFT_THREADING_STATIC")
 
   if(NOT SWIFT_COMPILER_IS_MSVC_LIKE)
     if(SWIFT_USE_LINKER)

include/swift/Threading/Impl/Linux/ulock.h

@@ -34,6 +34,8 @@
 #include <atomic>
 #include <cstdint>
 
+#include "swift/Threading/ThreadSanitizer.h"
+
 namespace swift {
 namespace threading_impl {
 namespace linux {
@@ -59,31 +61,38 @@ inline void ulock_lock(ulock_t *lock) {
   const ulock_t tid = ulock_get_tid();
   do {
     ulock_t zero = 0;
-    if (ulock_fastpath(__atomic_compare_exchange_n(
-            lock, &zero, tid, true, __ATOMIC_ACQUIRE, __ATOMIC_RELAXED)))
-      return;
-
+    if (ulock_fastpath(__atomic_compare_exchange_n(lock, &zero, tid,
+                                                   true, __ATOMIC_ACQUIRE,
+                                                   __ATOMIC_RELAXED)))
+      break;
   } while (ulock_futex(lock, FUTEX_LOCK_PI) != 0);
+
+  tsan::acquire(lock);
 }
 
 inline bool ulock_trylock(ulock_t *lock) {
   ulock_t zero = 0;
   if (ulock_fastpath(__atomic_compare_exchange_n(lock, &zero, ulock_get_tid(),
                                                  true, __ATOMIC_ACQUIRE,
-                                                 __ATOMIC_RELAXED)))
+                                                 __ATOMIC_RELAXED))
+      || ulock_futex(lock, FUTEX_TRYLOCK_PI) == 0) {
+    tsan::acquire(lock);
     return true;
+  }
 
-  return ulock_futex(lock, FUTEX_TRYLOCK_PI) == 0;
+  return false;
 }
 
 inline void ulock_unlock(ulock_t *lock) {
+  tsan::release(lock);
+
   const ulock_t tid = ulock_get_tid();
   do {
     ulock_t expected = tid;
-    if (ulock_fastpath(__atomic_compare_exchange_n(
-            lock, &expected, 0, true, __ATOMIC_RELEASE, __ATOMIC_RELAXED)))
-      return;
-
+    if (ulock_fastpath(__atomic_compare_exchange_n(lock, &expected, 0,
+                                                   true, __ATOMIC_RELEASE,
+                                                   __ATOMIC_RELAXED)))
+      break;
   } while (ulock_futex(lock, FUTEX_UNLOCK_PI) != 0);
 }
 

include/swift/Threading/ThreadSanitizer.h

@@ -0,0 +1,120 @@
+//===--- ThreadSanitizer.h - Thread Sanitizer support --------- -*- C++ -*-===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// Helper functions for code that needs to integrate with the thread
+// sanitizer.  In particular, TSan can't see inside the runtime libraries,
+// so we occasionally need to give it a hint that we're doing synchronization
+// in order to avoid false positives.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SWIFT_THREADING_THREAD_SANITIZER_H
+#define SWIFT_THREADING_THREAD_SANITIZER_H
+
+#include "swift/shims/Visibility.h"
+
+namespace swift {
+
+#if defined(_WIN32) || defined(__wasi__) || !__has_include(<dlfcn.h>)
+
+#define SWIFT_THREADING_TSAN_SUPPORT 0
+
+namespace tsan {
+
+inline bool enabled() { return false; }
+template <typename T> T *acquire(T *ptr) { return ptr; }
+template <typename T> T *release(T *ptr) { return ptr; }
+
+} // namespace tsan
+#else
+
+#define SWIFT_THREADING_TSAN_SUPPORT 1
+
+// If we're static linking to libswiftThreading.a, these symbols can come
+// from there.  If, on the other hand, we're dynamically linked, we want
+// to get them from libswiftCore.dylib instead.
+#if SWIFT_THREADING_STATIC
+#define SWIFT_THREADING_EXPORT extern "C"
+#else
+#define SWIFT_THREADING_EXPORT SWIFT_RUNTIME_EXPORT
+#endif
+
+namespace threading_impl {
+
+SWIFT_THREADING_EXPORT bool _swift_tsan_enabled;
+SWIFT_THREADING_EXPORT void (*_swift_tsan_acquire)(const void *ptr);
+SWIFT_THREADING_EXPORT void (*_swift_tsan_release)(const void *ptr);
+
+} // namespace threading_impl
+
+namespace tsan {
+
+/// Returns true if TSan is enabled
+inline bool enabled() {
+  return threading_impl::_swift_tsan_enabled;
+}
+
+/// Inform TSan about a synchronization operation.
+///
+/// This is used when TSan cannot see the synchronization operation, for
+/// example, if it is using a custom primitive for which TSan doesn't have
+/// a built-in interceptor.  This does not necessarily mean a lock or a C(++)
+/// atomic operation - it could be any kind of synchronization mechanism.
+///
+/// An acquire-release pair using the same address establishes an ordering
+/// constraint in TSan's happens-before graph, which TSan uses to determine
+/// whether two memory accesses from different threads have a well-defined
+/// order.
+///
+/// For instance, in
+///
+///     Thread 1                                Thread 2
+///
+///     access to y
+///     tsan::release(x)
+///     lock given away
+///
+///                      --> sync point -->
+///
+///                                             lock taken
+///                                             tsan::acquire(x)
+///                                             access to y
+///
+/// the access to y from Thread 2 is safe relative to the preceding access to
+/// y on Thread 1 because it is preceded by an acquire of x that was itself
+/// preceded by a release of x.
+template <typename T>
+T *acquire(T *ptr) {
+  if (threading_impl::_swift_tsan_acquire) {
+    threading_impl::_swift_tsan_acquire(ptr);
+  }
+  return ptr;
+}
+
+/// Inform TSan about a synchronization operation.
+///
+/// This is the counterpart to tsan::acquire.
+template <typename T>
+T *release(T *ptr) {
+  if (threading_impl::_swift_tsan_release) {
+    threading_impl::_swift_tsan_release(ptr);
+  }
+  return ptr;
+}
+
+} // namespace tsan
+
+#endif
+
+} // namespace swift
+
+#endif

lib/Threading/CMakeLists.txt

@@ -10,4 +10,5 @@ add_swift_host_library(swiftThreading STATIC
   Linux.cpp
   Pthreads.cpp
   Win32.cpp
-  Errors.cpp)
+  Errors.cpp
+  ThreadSanitizer.cpp)

lib/Threading/Linux.cpp

@@ -18,6 +18,7 @@
 
 #include "swift/Threading/Impl.h"
 #include "swift/Threading/Errors.h"
+#include "swift/Threading/ThreadSanitizer.h"
 
 namespace {
 
@@ -61,7 +62,7 @@ void swift::threading_impl::once_slow(once_t &predicate, void (*fn)(void *),
 #endif
   if (predicate.flag.load(std::memory_order_acquire) == 0) {
     fn(context);
-    predicate.flag.store(-1, std::memory_order_release);
+    predicate.flag.store(tsan::enabled() ? 1 : -1, std::memory_order_release);
   }
 #if defined(__LP64__) || defined(_LP64)
   linux::ulock_unlock(&predicate.lock);

lib/Threading/ThreadSanitizer.cpp

@@ -0,0 +1,61 @@
+//===--- ThreadSanitizer.cpp - Thread Sanitizer support -------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2023 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+//
+// Helper functions for code that needs to integrate with the thread
+// sanitizer.  In particular, TSan can't see inside the runtime libraries,
+// so we occasionally need to give it a hint that we're doing synchronization
+// in order to avoid false positives.
+//
+//===----------------------------------------------------------------------===//
+
+#include "swift/Threading/ThreadSanitizer.h"
+
+#if SWIFT_THREADING_TSAN_SUPPORT
+
+#include "swift/shims/Visibility.h"
+
+#include <cstdio>
+
+namespace swift {
+namespace threading_impl {
+
+SWIFT_THREADING_EXPORT bool _swift_tsan_enabled = false;
+SWIFT_THREADING_EXPORT void (*_swift_tsan_acquire)(const void *) = nullptr;
+SWIFT_THREADING_EXPORT void (*_swift_tsan_release)(const void *) = nullptr;
+
+#if __has_include(<dlfcn.h>)
+#include <dlfcn.h>
+
+// The TSan library code will call this function when it starts up
+extern "C" SWIFT_ATTRIBUTE_FOR_EXPORTS
+void __tsan_on_initialize() {
+  _swift_tsan_enabled = true;
+  _swift_tsan_acquire = (void (*)(const void *))dlsym(RTLD_DEFAULT,
+                                                      "__tsan_acquire");
+  _swift_tsan_release = (void (*)(const void *))dlsym(RTLD_DEFAULT,
+                                                      "__tsan_release");
+
+  // Always call through to the next image; this won't work on macOS, but it's
+  // important on Linux to allow others to hook into the thread sanitizer if
+  // they wish.
+  void (*next_init)(void);
+  next_init = (void (*)(void))dlsym(RTLD_NEXT, "__tsan_on_initialize");
+  if (next_init) {
+    next_init();
+  }
+}
+#endif // __has_include(<dlfcn.h>)
+
+} // namespace threading_impl
+} // namespace swift
+
+#endif // SWIFT_THREADING_TSAN_SUPPORT

stdlib/public/Concurrency/CMakeLists.txt

@@ -120,7 +120,6 @@ add_swift_target_library(swift_Concurrency ${SWIFT_STDLIB_LIBRARY_BUILD_TYPES} I
   TaskLocal.cpp
   TaskLocal.swift
   TaskSleep.swift
-  ThreadSanitizer.cpp
   ThreadingError.cpp
   TracingSignpost.cpp
   AsyncStreamBuffer.swift

stdlib/public/Concurrency/TaskPrivate.h

@@ -29,6 +29,7 @@
 #include "swift/Runtime/Exclusivity.h"
 #include "swift/Runtime/HeapObject.h"
 #include "swift/Threading/Thread.h"
+#include "swift/Threading/ThreadSanitizer.h"
 #include <atomic>
 #include <new>
 
@@ -99,15 +100,23 @@ void _swift_taskGroup_cancelAllChildren(TaskGroup *group);
 /// should generally use a higher-level function.
 void _swift_taskGroup_detachChild(TaskGroup *group, AsyncTask *child);
 
-/// release() establishes a happens-before relation with a preceding acquire()
-/// on the same address.
-void _swift_tsan_acquire(void *addr);
-void _swift_tsan_release(void *addr);
-/// Technically, this consume relies on implicit HW address dependency ordering
-/// and is paired with a corresponding release. Since TSAN doesn't know how to
-/// reason about this, we tell TSAN it's an acquire instead. See also
-/// SWIFT_MEMORY_ORDER_CONSUME definition.
-#define _swift_tsan_consume _swift_tsan_acquire
+/// Tell TSan about an acquiring load
+inline void _swift_tsan_acquire(void *addr) {
+  swift::tsan::acquire(addr);
+}
+/// Tell TSan about a releasing store
+inline void _swift_tsan_release(void *addr) {
+  swift::tsan::release(addr);
+}
+/// Tell TSan about a consuming load
+inline void _swift_tsan_consume(void *addr) {
+  // TSan doesn't support consume, so pretend it's an acquire.
+  //
+  // Note that that means that TSan won't generate errors for non-dependent
+  // reads, so this isn't entirely safe if you're relying solely on TSan to
+  // spot bugs.
+  swift::tsan::acquire(addr);
+}
 
 /// Special values used with DispatchQueueIndex to indicate the global and main
 /// executors.

stdlib/public/Concurrency/ThreadSanitizer.cpp

@@ -4,6 +4,7 @@
 list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_SOURCE_DIR}/../../cmake/modules")
 include(AddSwiftStdlib)
 
+# This should *not* include ThreadSanitizer.cpp, as that is part of libswiftCore
 add_swift_target_library(swiftThreading OBJECT_LIBRARY
   "${SWIFT_SOURCE_DIR}/lib/Threading/C11.cpp"
   "${SWIFT_SOURCE_DIR}/lib/Threading/Linux.cpp"