[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.

rdar://110665213

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

@@ -34,6 +34,8 @@
 #include <atomic>
 #include <cstdint>
 
+#include "swift/Threading/TSan.h"
+
 namespace swift {
 namespace threading_impl {
 namespace linux {
@@ -59,32 +61,39 @@ 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) {
   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);
+
+  tsan::release(lock);
 }
 
 } // namespace linux

include/swift/Threading/TSan.h

@@ -0,0 +1,95 @@
+//===--- TSan.h - TSan support functions ---------------------- -*- 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_TSAN_H
+#define SWIFT_THREADING_TSAN_H
+
+namespace swift {
+
+#if defined(_WIN32) || defined(__wasi__) || !__has_include(<dlfcn.h>)
+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; }
+template <typename T> T *consume(T *ptr) { return ptr; }
+
+} // namespace tsan
+#else
+
+namespace threading_impl {
+
+extern bool tsan_enabled;
+extern void (*tsan_acquire)(const void *ptr);
+extern void (*tsan_release)(const void *ptr);
+
+} // namespace threading_impl
+
+namespace tsan {
+
+/// Returns true if TSan is enabled
+inline bool enabled() {
+  return threading_impl::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::tsan_acquire) {
+    threading_impl::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::tsan_release) {
+    threading_impl::tsan_release(ptr);
+  }
+  return ptr;
+}
+
+/// Indicate to TSan that a consuming 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 those operations that carry a
+/// dependency on the loaded value*.
+///
+/// TSan doesn't currently know about consume, so we lie and say it's an
+/// acquire instead.
+template <typename T>
+T *consume(T *ptr) {
+  return acquire(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
+  TSan.cpp)

lib/Threading/Linux.cpp

@@ -18,6 +18,7 @@
 
 #include "swift/Threading/Impl.h"
 #include "swift/Threading/Errors.h"
+#include "swift/Threading/TSan.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/TSan.cpp

@@ -0,0 +1,49 @@
+//===--- TSan.h - TSan support functions ----------------------------------===//
+//
+// 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/TSan.h"
+#include "swift/shims/Visibility.h"
+
+namespace swift {
+namespace threading_impl {
+
+bool tsan_enabled = false;
+void (*tsan_acquire)(const void *) = nullptr;
+void (*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() {
+  tsan_enabled = true;
+  tsan_acquire = (void (*)(const void *))dlsym(RTLD_DEFAULT, "__tsan_acquire");
+  tsan_release = (void (*)(const void *))dlsym(RTLD_DEFAULT, "__tsan_release");
+
+  // Always call through to the next image
+  void (*next_init)(void);
+  next_init = (void (*)(void))dlsym(RTLD_NEXT, "__tsan_on_initialize");
+  if (next_init)
+    next_init();
+}
+#endif
+
+} // namespace threading_impl
+} // namespace swift

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/TSan.h"
 #include <atomic>
 #include <new>
 
@@ -99,15 +100,18 @@ 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) {
+  swift::tsan::consume(addr);
+}
 
 /// Special values used with DispatchQueueIndex to indicate the global and main
 /// executors.

stdlib/public/Concurrency/ThreadSanitizer.cpp

@@ -1,50 +0,0 @@
-//===--- ThreadSanitizer.cpp - Thread Sanitizer support -------------------===//
-//
-// This source file is part of the Swift.org open source project
-//
-// Copyright (c) 2014 - 2021 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
-//
-//===----------------------------------------------------------------------===//
-//
-// Thread Sanitizer support for the Swift Task runtime.
-//
-//===----------------------------------------------------------------------===//
-
-#include "TaskPrivate.h"
-
-// Thread Sanitizer is not supported on Windows or WASI.
-#if defined(_WIN32) || defined(__wasi__) || !__has_include(<dlfcn.h>)
-void swift::_swift_tsan_acquire(void *addr) {}
-void swift::_swift_tsan_release(void *addr) {}
-#else
-#include <dlfcn.h>
-
-namespace {
-using TSanFunc = void(void *);
-TSanFunc *tsan_acquire, *tsan_release;
-} // anonymous namespace
-
-void swift::_swift_tsan_acquire(void *addr) {
-  if (tsan_acquire) {
-    tsan_acquire(addr);
-    SWIFT_TASK_DEBUG_LOG("tsan_acquire on %p", addr);
-  }
-}
-
-void swift::_swift_tsan_release(void *addr) {
-  if (tsan_release) {
-    tsan_release(addr);
-    SWIFT_TASK_DEBUG_LOG("tsan_release on %p", addr);
-  }
-}
-
-SWIFT_EXPORT_FROM(swift_Concurrency) SWIFT_CC(c)
-void __tsan_on_initialize() {
-  tsan_acquire = (TSanFunc *)dlsym(RTLD_DEFAULT, "__tsan_acquire");
-  tsan_release = (TSanFunc *)dlsym(RTLD_DEFAULT, "__tsan_release");
-}
-#endif

stdlib/public/Threading/CMakeLists.txt

@@ -9,4 +9,5 @@ add_swift_target_library(swiftThreading OBJECT_LIBRARY
   "${SWIFT_SOURCE_DIR}/lib/Threading/Linux.cpp"
   "${SWIFT_SOURCE_DIR}/lib/Threading/Pthreads.cpp"
   "${SWIFT_SOURCE_DIR}/lib/Threading/Win32.cpp"
+  "${SWIFT_SOURCE_DIR}/lib/Threading/TSan.cpp"
   INSTALL_IN_COMPONENT never_install)

test/Sanitizers/tsan/once.swift

@@ -0,0 +1,50 @@
+// RUN: %target-swiftc_driver %s -Xfrontend -parse-as-library -g -sanitize=thread %import-libdispatch -o %t_tsan-binary
+// RUN: %target-codesign %t_tsan-binary
+// RUN: env %env-TSAN_OPTIONS=abort_on_error=0 %target-run %t_tsan-binary 2>&1 | %FileCheck %s --implicit-check-not='ThreadSanitizer'
+// REQUIRES: executable_test
+// REQUIRES: tsan_runtime
+
+// rdar://101876380
+// UNSUPPORTED: OS=ios
+
+// FIXME: This should be covered by "tsan_runtime"; older versions of Apple OSs
+// don't support TSan.
+// UNSUPPORTED: remote_run
+
+// Test that we do not report a race on initialization; Swift doesn't initialize
+// globals at start-up, but rather uses swift_once().  This is thread safe, but
+// on some platforms TSan wasn't seeing the synchronization, so would report
+// a false positive.
+
+import Dispatch
+
+var count = 0
+let foo = {
+  count += 1
+  return count
+}()
+
+@main
+struct Main {
+  static func main() {
+    let q = DispatchQueue(label: "q", attributes: .concurrent)
+    let finished = DispatchSemaphore(value: 0)
+    let count = 100
+
+    for _ in 0..<count {
+      q.async {
+        print(foo)
+        finished.signal()
+      }
+    }
+
+    for _ in 0..<count {
+      finished.wait()
+    }
+
+    print("Done!")
+  }
+}
+
+// CHECK-NOT: ThreadSanitizer: data race
+// CHECK: Done!