[Runtime] Faster dynamic exclusivity checking implemented in Swift.

Replace C++ implementation of swift_beginAccess and swift_endAccess with (almost) pure Swift implementation. Helpers remain in C++ for TLS, getting return addresses, and raising a fatal error on violations.

This change also moves the exclusivity access set head from the shared SwiftTLSContext structure to a dedicated TLS key. This improves performance, which is important for exclusivity checking. This is particularly the case where we can inline TLS access with a constant key, as on Darwin ARM64.

The code that bridges exclusivity tracking into Concurrency remains in C++. The new Swift implementation exposes a few helpers for it to use as a replacement for directly manipulating the C++ implementation.

rdar://161122309

Author: mikeash

include/swift/Threading/Impl/Darwin.h

@@ -285,6 +285,8 @@ inline tls_key_t tls_get_key(tls_key k) {
     return __PTK_FRAMEWORK_SWIFT_KEY5;
   case tls_key::observation_transaction:
     return __PTK_FRAMEWORK_SWIFT_KEY6;
+  case tls_key::exclusivity:
+    return __PTK_FRAMEWORK_SWIFT_KEY7;
   }
 }
 

include/swift/Threading/TLSKeys.h

@@ -22,7 +22,8 @@ enum class tls_key {
   concurrency_task,
   concurrency_executor_tracking_info,
   concurrency_fallback,
-  observation_transaction
+  observation_transaction,
+  exclusivity,
 };
 
 } // namespace swift

stdlib/public/SwiftShims/swift/shims/AssertionReporting.h

@@ -68,6 +68,13 @@ void _swift_stdlib_reportUnimplementedInitializer(
     const unsigned char *initName, int initNameLength,
     __swift_uint32_t flags);
 
+SWIFT_RUNTIME_STDLIB_SPI
+void _swift_reportExclusivityConflict(__swift_uintptr_t oldAction,
+                                      const void *_Nullable oldPC,
+                                      __swift_uintptr_t newFlags,
+                                      const void *_Nullable newPC,
+                                      const void *_Nullable pointer);
+
 #ifdef __cplusplus
 } // extern "C"
 #endif

stdlib/public/SwiftShims/swift/shims/LibcShims.h

@@ -38,6 +38,9 @@ SWIFT_RUNTIME_STDLIB_INTERNAL
 __swift_size_t _swift_stdlib_fwrite_stdout(const void *ptr, __swift_size_t size,
                                            __swift_size_t nitems);
 
+SWIFT_RUNTIME_STDLIB_INTERNAL
+void _swift_stdlib_fputs_stderr(const char *str);
+
 // General utilities <stdlib.h>
 // Memory management functions
 static inline void _swift_stdlib_free(void *_Nullable ptr) {

stdlib/public/SwiftShims/swift/shims/RuntimeStubs.h

@@ -20,6 +20,7 @@
 #define SWIFT_STDLIB_SHIMS_RUNTIMESTUBS_H_
 
 #include "LibcShims.h"
+#include "Visibility.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -31,6 +32,25 @@ SWIFT_RUNTIME_STDLIB_API
 __swift_ssize_t
 swift_stdlib_readLine_stdin(unsigned char * _Nullable * _Nonnull LinePtr);
 
+// Pick a return-address strategy
+#if defined(__wasm__)
+// Wasm can't access call frame for security purposes
+#define get_return_address() ((void*) 0)
+#elif __GNUC__
+#define get_return_address() __builtin_return_address(0)
+#elif _MSC_VER
+#include <intrin.h>
+#define get_return_address() _ReturnAddress()
+#else
+#error missing implementation for get_return_address
+#define get_return_address() ((void*) 0)
+#endif
+
+SWIFT_ALWAYS_INLINE
+static inline const void *_Nullable _swift_stdlib_get_return_address() {
+  return get_return_address();
+}
+
 SWIFT_END_NULLABILITY_ANNOTATIONS
 
 #ifdef __cplusplus

stdlib/public/SwiftShims/swift/shims/ThreadLocalStorage.h

@@ -18,4 +18,48 @@
 SWIFT_RUNTIME_STDLIB_INTERNAL
 void * _Nonnull _swift_stdlib_threadLocalStorageGet(void);
 
+SWIFT_RUNTIME_STDLIB_INTERNAL
+void * _Nullable _swift_getExclusivityTLSImpl();
+
+SWIFT_RUNTIME_STDLIB_INTERNAL
+void _swift_setExclusivityTLSImpl(void * _Nullable newValue);
+
+#if defined(__APPLE__) && __arm64__
+
+// Use a fast path on Apple ARM64, where we have a dedicated TLS key and fast
+// access to read/write it.
+
+#ifndef __PTK_FRAMEWORK_SWIFT_KEY7
+# define __PTK_FRAMEWORK_SWIFT_KEY7 107
+#endif
+
+#define SWIFT_RUNTIME_EXCLUSIVITY_KEY __PTK_FRAMEWORK_SWIFT_KEY7
+
+static inline void * _Nullable * _Nonnull _swift_getExclusivityTLSPointer() {
+  unsigned long tsd;
+  __asm__ ("mrs %0, TPIDRRO_EL0" : "=r" (tsd));
+  void **base = (void **)tsd;
+  return &base[SWIFT_RUNTIME_EXCLUSIVITY_KEY];
+}
+
+static inline void * _Nullable _swift_getExclusivityTLS() {
+  return *_swift_getExclusivityTLSPointer();
+}
+
+static inline void _swift_setExclusivityTLS(void * _Nullable newValue) {
+  *_swift_getExclusivityTLSPointer() = newValue;
+}
+
+#else
+
+static inline void * _Nullable _swift_getExclusivityTLS() {
+  return _swift_getExclusivityTLSImpl();
+}
+
+static inline void _swift_setExclusivityTLS(void * _Nullable newValue) {
+  _swift_setExclusivityTLSImpl(newValue);
+}
+
+#endif
+
 #endif // SWIFT_STDLIB_SHIMS_THREADLOCALSTORAGE_H

stdlib/public/core/CMakeLists.txt

@@ -78,6 +78,7 @@ split_embedded_sources(
   EMBEDDED EnumeratedSequence.swift
   EMBEDDED Equatable.swift
   EMBEDDED ErrorType.swift
+    NORMAL Exclusivity.swift
   EMBEDDED ExistentialCollection.swift
   EMBEDDED Filter.swift
   EMBEDDED FlatMap.swift

stdlib/public/core/Exclusivity.swift

@@ -0,0 +1,262 @@
+import SwiftShims
+
+fileprivate let ValueBufferSize = unsafe 3 * MemoryLayout<UnsafeRawPointer>.stride
+
+fileprivate let TrackingFlag: UInt = 0x20
+fileprivate let ActionMask: UInt = 0x1
+
+fileprivate typealias AccessPointer = UnsafeMutablePointer<Access>
+@unsafe fileprivate struct Access {
+
+  enum Action: UInt {
+    case read
+    case modify
+  }
+
+  struct NextAndAction {
+    private var rawValue: UInt
+
+    var action: Action {
+      get {
+        Action(rawValue: rawValue & ActionMask)!
+      }
+      set {
+        rawValue = (rawValue & ~ActionMask) | newValue.rawValue
+      }
+    }
+
+    var next: AccessPointer? {
+      get { unsafe UnsafeMutablePointer(bitPattern: rawValue & ~ActionMask) }
+      set { rawValue = UInt(bitPattern: newValue) | (rawValue & ActionMask) }
+    }
+  }
+
+  var location: UnsafeRawPointer?
+  var pc: UnsafeRawPointer?
+  var nextAndAction: NextAndAction
+
+  var action: Action {
+    get { unsafe nextAndAction.action }
+    set { unsafe nextAndAction.action = newValue }
+  }
+
+  var next: AccessPointer? {
+    get { unsafe nextAndAction.next }
+    set { unsafe nextAndAction.next = newValue }
+  }
+
+  static func from(rawPointer: UnsafeMutableRawPointer?) -> AccessPointer? {
+    guard let rawPointer = unsafe rawPointer else { return nil }
+    return unsafe rawPointer.assumingMemoryBound(to: Access.self)
+  }
+
+  @inline(__always)
+  static func search(
+    buffer: UnsafeMutableRawPointer,
+    location: UnsafeRawPointer,
+    pc: UnsafeRawPointer?,
+    action: Action,
+    inserting: Bool,
+    head: inout AccessPointer?
+  ) {
+    var cursor = unsafe head
+    while let nextPtr = unsafe cursor {
+      if unsafe nextPtr.pointee.location == location {
+        if unsafe nextPtr.pointee.action == Action.modify
+            || action == Action.modify {
+          unsafe _swift_reportExclusivityConflict(
+            nextPtr.pointee.action.rawValue,
+            nextPtr.pointee.pc,
+            action.rawValue,
+            pc,
+            location)
+        }
+      }
+      unsafe cursor = nextPtr.pointee.next
+    }
+
+    if inserting {
+      guard let access = unsafe Access.from(rawPointer: buffer) else {
+        nullAccessBuffer()
+      }
+
+      unsafe access.pointee.location = location
+      unsafe access.pointee.pc = pc
+      unsafe access.pointee.action = action
+
+      unsafe access.pointee.next = head
+      unsafe head = access
+    }
+  }
+
+  @inline(__always)
+  static func remove(access: AccessPointer, head: inout AccessPointer?) {
+    var cursor = unsafe head
+    var previous: AccessPointer? = nil
+    while let nextPtr = unsafe cursor {
+      if unsafe nextPtr == access {
+        if let previous = unsafe previous {
+          unsafe previous.pointee.next = access.pointee.next
+        } else {
+          unsafe head = access.pointee.next
+        }
+        return
+      }
+      unsafe previous = nextPtr
+      unsafe cursor = nextPtr.pointee.next
+    }
+
+    unsafe accessNotFound(access)
+  }
+
+  @inline(__always)
+  static func findParent(
+    access: AccessPointer,
+    child: AccessPointer?
+  ) -> AccessPointer? {
+    var cursor = unsafe access
+    while let next = unsafe cursor.pointee.next {
+      if unsafe next == child {
+        return unsafe cursor
+      }
+      unsafe cursor = next
+    }
+
+    // If we were searching for nil, then we found it.
+    if unsafe child == nil {
+      return unsafe cursor
+    }
+
+    // If we were searching for a non-nil node, we didn't find it.
+    return nil
+  }
+
+  static func forEach(_ head: AccessPointer?, _ action: (Access) -> Void) {
+    var cursor = unsafe head
+    while let nextPtr = unsafe cursor {
+      unsafe action(nextPtr.pointee)
+      unsafe cursor = nextPtr.pointee.next
+    }
+  }
+}
+
+fileprivate var accessHead: AccessPointer? {
+  get { unsafe Access.from(rawPointer: _swift_getExclusivityTLS()) }
+  set { unsafe _swift_setExclusivityTLS(newValue) }
+}
+
+@_cdecl("swift_beginAccess")
+@usableFromInline
+@unsafe
+internal func swift_beginAccess(
+  pointer: UnsafeRawPointer,
+  buffer: UnsafeMutableRawPointer,
+  flags: UInt,
+  pc: UnsafeRawPointer?
+) {
+  // Make sure that Access fits in the value buffer. This should be a static
+  // assert, but we'll settle for a precondition for now. The optimizer should
+  // eliminate the check when it's true.
+  precondition(unsafe MemoryLayout<Access>.size <= ValueBufferSize)
+
+  guard let action = Access.Action(rawValue: flags & ActionMask) else {
+    invalidFlags(flags)
+  }
+
+  let isTracking = (flags & TrackingFlag) != 0
+
+  unsafe Access.search(
+    buffer: buffer,
+    location: pointer,
+    pc: pc ?? _swift_stdlib_get_return_address(),
+    action: action,
+    inserting: isTracking,
+    head: &accessHead)
+}
+
+@_cdecl("swift_endAccess")
+@usableFromInline
+@unsafe
+internal func swift_endAccess(buffer: UnsafeMutableRawPointer) {
+  guard let access = unsafe Access.from(rawPointer: buffer) else {
+    nullAccessBuffer()
+  }
+  unsafe Access.remove(access: access, head: &accessHead)
+}
+
+@_cdecl("_swift_exclusivityAccessSetNext")
+@usableFromInline
+@unsafe
+internal func _swift_exclusivityAccessSetNext(
+  access: UnsafeMutableRawPointer,
+  next: UnsafeMutableRawPointer
+) {
+  let access = unsafe Access.from(rawPointer: access)
+  let next = unsafe Access.from(rawPointer: next)
+  unsafe access?.pointee.next = next
+}
+
+@_cdecl("swift_dumpTrackedAccesses")
+@usableFromInline
+@unsafe
+internal func swift_dumpTrackedAccesses() {
+  if let head = unsafe accessHead {
+    unsafe Access.forEach(head) {
+      unsafe _swift_stdlib_fputs_stderr("        Access. " +
+          "Pointer: \($0.location, default: "<null>")." +
+          "PC: \($0.pc, default: "<null>"). " +
+          "AccessAction: \($0.action)\n")
+    }
+  } else {
+    unsafe _swift_stdlib_fputs_stderr("        No Accesses.\n")
+  }
+}
+
+/// Starting from `access`, find the access that is the parent node of `child`.
+/// If `child` is `nil`, find the last access in the list.
+@_cdecl("_swift_exclusivityAccessGetParent")
+@usableFromInline
+@unsafe
+internal func _swift_exclusivityAccessGetParent(
+  access: UnsafeMutableRawPointer?,
+  child: UnsafeMutableRawPointer?
+) -> UnsafeMutableRawPointer? {
+    if let access = unsafe Access.from(rawPointer: access) {
+      let result = unsafe Access.findParent(
+        access: access, child: Access.from(rawPointer: child))
+      return UnsafeMutableRawPointer(result)
+    }
+    return nil
+}
+
+@inline(never)
+fileprivate func invalidFlags(_ flags: UInt) -> Never {
+  reportExclusivityError("Internal exclusivity error",
+                         "unable to construct action from flags \(flags)")
+}
+
+@inline(never)
+fileprivate func accessNotFound(_ access: AccessPointer) -> Never {
+  unsafe reportExclusivityError("Internal exclusivity error",
+                                "didn't find exclusive access buffer \(access)")
+}
+
+@inline(never)
+fileprivate func nullAccessBuffer() -> Never {
+  reportExclusivityError("Internal exclusivity error", "NULL access buffer")
+}
+
+fileprivate func reportExclusivityError(
+  _ prefix: StaticString, _ message: String
+) -> Never {
+  prefix.withUTF8Buffer { prefixBuffer in
+    var message = message
+    message.withUTF8 { messageBuffer in
+      unsafe _swift_stdlib_reportFatalError(
+          prefixBuffer.baseAddress!, CInt(prefixBuffer.count),
+          messageBuffer.baseAddress!, CInt(messageBuffer.count),
+          0)
+    }
+  }
+  Builtin.int_trap()
+}