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

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,97 @@
+//===--- 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;
+}
+
+/// Indicate to TSan that an acquiring load has occurred on the current
+/// thread.  If some other thread does a releasing store with the same
+/// pointer, we are indicating to TSan that all writes that happened
+/// before that store will be visible to the current thread after the
+/// `acquire()`.
+template <typename T>
+T *acquire(T *ptr) {
+  if (threading_impl::_swift_tsan_acquire) {
+    threading_impl::_swift_tsan_acquire(ptr);
+  }
+  return ptr;
+}
+
+/// Indicate to TSan that a releasing store has occurred on the current
+/// thread.  If some other thread does an acquiring load with the same
+/// pointer, we are indicating to TSan that that thread will be able to
+/// see all writes that happened before the `release()`.
+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/Threading/CMakeLists.txt

@@ -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"

stdlib/public/runtime/CMakeLists.txt

@@ -80,6 +80,12 @@ set(swift_runtime_sources
     AccessibleFunction.cpp
     Win32.cpp)
 
+# We pull this in separately here because other dylibs will need it, but only
+# will have the __tsan_on_initialize called, and on Darwin, RTLD_NEXT can't be
+# used to call subsequence dylibs' copies of that.
+set(swift_runtime_threading_sources
+  ${SWIFT_SOURCE_DIR}/lib/Threading/ThreadSanitizer.cpp)
+
 set(swift_runtime_backtracing_sources
   Backtrace.cpp
   CrashHandlerMacOS.cpp
@@ -132,6 +138,7 @@ add_swift_target_library(swiftRuntime OBJECT_LIBRARY
   ${swift_runtime_sources}
   ${swift_runtime_objc_sources}
   ${swift_runtime_leaks_sources}
+  ${swift_runtime_threading_sources}
   C_COMPILE_FLAGS
     ${swift_runtime_library_compile_flags}
   LINK_FLAGS ${swift_runtime_linker_flags}