[runtime] thread safety for weak references

It has been fairly easy to cause the runtime to crash on multithreaded read-read access to weak references (e.g. https://bugs.swift.org/browse/SR-192). Although weak references are value types, they can get elevated to the heap in multiple ways, such as when captured by a closure or when used as a property in a class object instance. In such cases, race conditions involving weak references could cause the runtime to perform to multiple decrement operations of the unowned reference count for a single increment; this eventually causes early deallocation, leading to use-after-free, modify-after-free and double-free errors.

This commit changes the weak reference operations to use a spinlock rather than assuming thread-exclusive access, when appropriate.
With this change, the crasher discussed in SR-192 no longer encounters crashes due to modify-after-free or double-free errors.

Author: glessard

include/swift/Runtime/HeapObject.h

@@ -634,9 +634,15 @@ static inline void swift_unownedTakeAssign(UnownedReference *dest,
 
 /// A weak reference value object.  This is ABI.
 struct WeakReference {
-  HeapObject *Value;
+  uintptr_t Value;
 };
 
+/// Return true if this is a native weak reference
+///
+/// \param ref - never null
+/// \return true if ref is a native weak reference
+bool isNativeSwiftWeakReference(WeakReference *ref);
+
 /// Initialize a weak reference.
 ///
 /// \param ref - never null

stdlib/public/runtime/HeapObject.cpp

@@ -779,73 +779,142 @@ void swift::swift_deallocObject(HeapObject *object,
   }
 }
 
+enum: uintptr_t {
+  WR_NATIVE = 1<<(swift::heap_object_abi::ObjCReservedLowBits),
+  WR_READING = 1<<(swift::heap_object_abi::ObjCReservedLowBits+1),
+
+  WR_NATIVEMASK = WR_NATIVE | swift::heap_object_abi::ObjCReservedBitsMask,
+};
+
+static_assert(WR_READING < alignof(void*),
+              "weakref lock bit mustn't interfere with real pointer bits");
+
+enum: short {
+  WR_SPINLIMIT = 64,
+};
+
+bool swift::isNativeSwiftWeakReference(WeakReference *ref) {
+  return (ref->Value & WR_NATIVEMASK) == WR_NATIVE;
+}
+
 void swift::swift_weakInit(WeakReference *ref, HeapObject *value) {
-  ref->Value = value;
+  ref->Value = (uintptr_t)value | WR_NATIVE;
   SWIFT_RT_ENTRY_CALL(swift_unownedRetain)(value);
 }
 
 void swift::swift_weakAssign(WeakReference *ref, HeapObject *newValue) {
   SWIFT_RT_ENTRY_CALL(swift_unownedRetain)(newValue);
-  auto oldValue = ref->Value;
-  ref->Value = newValue;
+  auto oldValue = (HeapObject*) (ref->Value & ~WR_NATIVE);
+  ref->Value = (uintptr_t)newValue | WR_NATIVE;
   SWIFT_RT_ENTRY_CALL(swift_unownedRelease)(oldValue);
 }
 
 HeapObject *swift::swift_weakLoadStrong(WeakReference *ref) {
-  auto object = ref->Value;
-  if (object == nullptr) return nullptr;
+  if (ref->Value == (uintptr_t)nullptr) {
+    return nullptr;
+  }
+
+  // ref might be visible to other threads
+  auto ptr = __atomic_fetch_or(&ref->Value, WR_READING, __ATOMIC_RELAXED);
+  while (ptr & WR_READING) {
+    short c = 0;
+    while (__atomic_load_n(&ref->Value, __ATOMIC_RELAXED) & WR_READING) {
+      if (++c == WR_SPINLIMIT) {
+        sched_yield();
+        c -= 1;
+      }
+    }
+    ptr = __atomic_fetch_or(&ref->Value, WR_READING, __ATOMIC_RELAXED);
+  }
+
+  auto object = (HeapObject*)(ptr & ~WR_NATIVE);
+  if (object == nullptr) {
+    __atomic_store_n(&ref->Value, (uintptr_t)nullptr, __ATOMIC_RELAXED);
+    return nullptr;
+  }
   if (object->refCount.isDeallocating()) {
+    __atomic_store_n(&ref->Value, (uintptr_t)nullptr, __ATOMIC_RELAXED);
     SWIFT_RT_ENTRY_CALL(swift_unownedRelease)(object);
-    ref->Value = nullptr;
     return nullptr;
   }
-  return swift_tryRetain(object);
+  auto result = swift_tryRetain(object);
+  __atomic_store_n(&ref->Value, ptr, __ATOMIC_RELAXED);
+  return result;
 }
 
 HeapObject *swift::swift_weakTakeStrong(WeakReference *ref) {
-  auto result = swift_weakLoadStrong(ref);
-  swift_weakDestroy(ref);
+  auto object = (HeapObject*) (ref->Value & ~WR_NATIVE);
+  if (object == nullptr) return nullptr;
+  auto result = swift_tryRetain(object);
+  ref->Value = (uintptr_t)nullptr;
+  swift_unownedRelease(object);
   return result;
 }
 
 void swift::swift_weakDestroy(WeakReference *ref) {
-  auto tmp = ref->Value;
-  ref->Value = nullptr;
+  auto tmp = (HeapObject*) (ref->Value & ~WR_NATIVE);
+  ref->Value = (uintptr_t)nullptr;
   SWIFT_RT_ENTRY_CALL(swift_unownedRelease)(tmp);
 }
 
 void swift::swift_weakCopyInit(WeakReference *dest, WeakReference *src) {
-  auto object = src->Value;
+  if (src->Value == (uintptr_t)nullptr) {
+    dest->Value = (uintptr_t)nullptr;
+    return;
+  }
+
+  // src might be visible to other threads
+  auto ptr = __atomic_fetch_or(&src->Value, WR_READING, __ATOMIC_RELAXED);
+  while (ptr & WR_READING) {
+    short c = 0;
+    while (__atomic_load_n(&src->Value, __ATOMIC_RELAXED) & WR_READING) {
+      if (++c == WR_SPINLIMIT) {
+        sched_yield();
+        c -= 1;
+      }
+    }
+    ptr = __atomic_fetch_or(&src->Value, WR_READING, __ATOMIC_RELAXED);
+  }
+
+  auto object = (HeapObject*)(ptr & ~WR_NATIVE);
   if (object == nullptr) {
-    dest->Value = nullptr;
+    __atomic_store_n(&src->Value, (uintptr_t)nullptr, __ATOMIC_RELAXED);
+    dest->Value = (uintptr_t)nullptr;
   } else if (object->refCount.isDeallocating()) {
-    src->Value = nullptr;
-    dest->Value = nullptr;
+    __atomic_store_n(&src->Value, (uintptr_t)nullptr, __ATOMIC_RELAXED);
     SWIFT_RT_ENTRY_CALL(swift_unownedRelease)(object);
+    dest->Value = (uintptr_t)nullptr;
   } else {
-    dest->Value = object;
     SWIFT_RT_ENTRY_CALL(swift_unownedRetain)(object);
+    __atomic_store_n(&src->Value, ptr, __ATOMIC_RELAXED);
+    dest->Value = (uintptr_t)object | WR_NATIVE;
   }
 }
 
 void swift::swift_weakTakeInit(WeakReference *dest, WeakReference *src) {
-  auto object = src->Value;
-  dest->Value = object;
-  if (object != nullptr && object->refCount.isDeallocating()) {
-    dest->Value = nullptr;
+  auto object = (HeapObject*) (src->Value & ~WR_NATIVE);
+  if (object == nullptr) {
+    dest->Value = (uintptr_t)nullptr;
+  } else if (object->refCount.isDeallocating()) {
+    dest->Value = (uintptr_t)nullptr;
     SWIFT_RT_ENTRY_CALL(swift_unownedRelease)(object);
+  } else {
+    dest->Value = (uintptr_t)object | WR_NATIVE;
   }
+  src->Value = (uintptr_t)nullptr;
 }
 
 void swift::swift_weakCopyAssign(WeakReference *dest, WeakReference *src) {
-  if (auto object = dest->Value) {
+  if (dest->Value) {
+    auto object = (HeapObject*) (dest->Value & ~WR_NATIVE);
     SWIFT_RT_ENTRY_CALL(swift_unownedRelease)(object);
   }
   swift_weakCopyInit(dest, src);
 }
 
 void swift::swift_weakTakeAssign(WeakReference *dest, WeakReference *src) {
-  if (auto object = dest->Value) {
+  if (dest->Value) {
+    auto object = (HeapObject*) (dest->Value & ~WR_NATIVE);
     SWIFT_RT_ENTRY_CALL(swift_unownedRelease)(object);
   }
   swift_weakTakeInit(dest, src);

stdlib/public/runtime/SwiftObject.mm

@@ -1022,7 +1022,7 @@ static void doWeakDestroy(WeakReference *addr, bool valueIsNative) {
 
 void swift::swift_unknownWeakInit(WeakReference *addr, void *value) {
   if (isObjCTaggedPointerOrNull(value)) {
-    addr->Value = (HeapObject*) value;
+    addr->Value = (uintptr_t) value;
     return;
   }
   doWeakInit(addr, value, usesNativeSwiftReferenceCounting_allocated(value));
@@ -1033,18 +1033,18 @@ static void doWeakDestroy(WeakReference *addr, bool valueIsNative) {
   // and re-initialize.
   if (isObjCTaggedPointerOrNull(newValue)) {
     swift_unknownWeakDestroy(addr);
-    addr->Value = (HeapObject*) newValue;
+    addr->Value = (uintptr_t) newValue;
     return;
   }
 
   bool newIsNative = usesNativeSwiftReferenceCounting_allocated(newValue);
 
   // If the existing value is not allocated, this is just an initialize.
-  void *oldValue = addr->Value;
+  void *oldValue = (void*) addr->Value;
   if (isObjCTaggedPointerOrNull(oldValue))
     return doWeakInit(addr, newValue, newIsNative);
 
-  bool oldIsNative = usesNativeSwiftReferenceCounting_allocated(oldValue);
+  bool oldIsNative = isNativeSwiftWeakReference(addr);
 
   // If they're both native, we can use the native function.
   if (oldIsNative && newIsNative)
@@ -1061,53 +1061,62 @@ static void doWeakDestroy(WeakReference *addr, bool valueIsNative) {
 }
 
 void *swift::swift_unknownWeakLoadStrong(WeakReference *addr) {
-  void *value = addr->Value;
-  if (isObjCTaggedPointerOrNull(value)) return value;
-
-  if (usesNativeSwiftReferenceCounting_allocated(value)) {
+  if (isNativeSwiftWeakReference(addr)) {
     return swift_weakLoadStrong(addr);
-  } else {
-    return (void*) objc_loadWeakRetained((id*) &addr->Value);
   }
-}
 
-void *swift::swift_unknownWeakTakeStrong(WeakReference *addr) {
-  void *value = addr->Value;
+  void *value = (void*) addr->Value;
   if (isObjCTaggedPointerOrNull(value)) return value;
 
-  if (usesNativeSwiftReferenceCounting_allocated(value)) {
+  return (void*) objc_loadWeakRetained((id*) &addr->Value);
+}
+
+void *swift::swift_unknownWeakTakeStrong(WeakReference *addr) {
+  if (isNativeSwiftWeakReference(addr)) {
     return swift_weakTakeStrong(addr);
-  } else {
-    void *result = (void*) objc_loadWeakRetained((id*) &addr->Value);
-    objc_destroyWeak((id*) &addr->Value);
-    return result;
   }
+
+  void *value = (void*) addr->Value;
+  if (isObjCTaggedPointerOrNull(value)) return value;
+
+  void *result = (void*) objc_loadWeakRetained((id*) &addr->Value);
+  objc_destroyWeak((id*) &addr->Value);
+  return result;
 }
 
 void swift::swift_unknownWeakDestroy(WeakReference *addr) {
+  if (isNativeSwiftWeakReference(addr)) {
+    return swift_weakDestroy(addr);
+  }
+
   id object = (id) addr->Value;
   if (isObjCTaggedPointerOrNull(object)) return;
-  doWeakDestroy(addr, usesNativeSwiftReferenceCounting_allocated(object));
+  objc_destroyWeak((id*) &addr->Value);
 }
+
 void swift::swift_unknownWeakCopyInit(WeakReference *dest, WeakReference *src) {
+  if (isNativeSwiftWeakReference(src)) {
+    return swift_weakCopyInit(dest, src);
+  }
+
   id object = (id) src->Value;
   if (isObjCTaggedPointerOrNull(object)) {
-    dest->Value = (HeapObject*) object;
-    return;
+    dest->Value = (uintptr_t) object;
+  } else {
+    objc_copyWeak((id*) &dest->Value, (id*) src);
   }
-  if (usesNativeSwiftReferenceCounting_allocated(object))
-    return swift_weakCopyInit(dest, src);
-  objc_copyWeak((id*) &dest->Value, (id*) src);
 }
 void swift::swift_unknownWeakTakeInit(WeakReference *dest, WeakReference *src) {
+  if (isNativeSwiftWeakReference(src)) {
+    return swift_weakTakeInit(dest, src);
+  }
+
   id object = (id) src->Value;
   if (isObjCTaggedPointerOrNull(object)) {
-    dest->Value = (HeapObject*) object;
-    return;
+    dest->Value = (uintptr_t) object;
+  } else {
+    objc_moveWeak((id*) &dest->Value, (id*) &src->Value);
   }
-  if (usesNativeSwiftReferenceCounting_allocated(object))
-    return swift_weakTakeInit(dest, src);
-  objc_moveWeak((id*) &dest->Value, (id*) &src->Value);
 }
 void swift::swift_unknownWeakCopyAssign(WeakReference *dest, WeakReference *src) {
   if (dest == src) return;