fix(blocks runtime): Implement _Block_tryRetain

Prevents race conditions when we try to convert a weak reference to a block to a strong one while another thread is concurrently deallocating the block.

arc.mm

@@ -23,7 +23,6 @@
 #import "visibility.h"
 #import "objc/hooks.h"
 #import "objc/objc-arc.h"
-#import "objc/blocks_runtime.h"
 #include "objc/message.h"
 
 /**
@@ -953,13 +952,24 @@ static BOOL setObjectHasWeakRefs(id obj)
 	else if (UNLIKELY(objc_test_class_flag(cls, objc_class_flag_is_block)))
 	{
 		obj = static_cast<id>(block_load_weak(obj));
-	}
+		if (obj == nil)
+		{
+			return nil;
+		}
+		// This is a defeasible retain operation that protects against another thread concurrently
+		// starting to deallocate the block.
+		if (_Block_tryRetain(obj))
+		{
+			return obj;
+		}
+		return nil;
 
+	}
 	else if (!objc_test_class_flag(cls, objc_class_flag_fast_arc))
 	{
 		obj = _objc_weak_load(obj);
 	}
-	// block_load_weak() or _objc_weak_load() can return nil
+	// _objc_weak_load() can return nil
 	if (obj == nil) { return nil; }
 	return retain(obj, YES);
 }

blocks_runtime.m

@@ -24,6 +24,7 @@
  *  OTHER DEALINGS IN THE SOFTWARE.
  */
 #import "objc/blocks_runtime.h"
+#include "objc/blocks_private.h"
 #import "objc/runtime.h"
 #import "objc/objc-arc.h"
 #include "blocks_runtime.h"
@@ -299,3 +300,25 @@ OBJC_PUBLIC bool _Block_isDeallocating(const void* arg)
 	int refCount = __sync_fetch_and_add(refCountPtr, 0);
 	return refCount == 0;
 }
+
+OBJC_PUBLIC bool _Block_tryRetain(const void* arg)
+{
+	/* This is used by the weak reference management in ARC. The implementation
+	 * follows the reasoning of `retain_fast()` in arc.mm: We want to abandon the
+	 * retain operation if another thread has started deallocating the object between
+	 * loading the weak pointer and executing the retain operation.
+	 */
+	struct Block_layout *block = (struct Block_layout*)arg;
+	int *refCountPtr = &block->reserved;
+	int refCountVal = __sync_fetch_and_add(refCountPtr, 0);
+	int newVal = refCountVal;
+	do {
+		refCountVal = newVal;
+		if (refCountVal <= 0)
+		{
+			return false;
+		}
+		newVal = __sync_val_compare_and_swap(refCountPtr, refCountVal, newVal + 1);
+	} while (newVal != refCountVal);
+	return true;
+}

objc/blocks_private.h

@@ -45,7 +45,22 @@ struct Block_descriptor
 	const char *encoding;
 };
 
+/** 
+  * Checks whether the block is currently being deallocated.
+  *
+  * Used by ARC weak reference management. Only call this after the weak
+  * reference lock is acquired. 
+  */
 OBJC_PUBLIC BLOCKS_EXPORT bool _Block_isDeallocating(const void *aBlock);
+/**
+ *  Atomically increments the reference count of the block.
+ *  Returns true if the block was retained, and false if it is already 
+ *  being deallocated.
+ *
+ * Used by ARC weak reference management. Only call this after the weak
+ * reference lock is acquired.
+ */
+OBJC_PUBLIC BLOCKS_EXPORT bool _Block_tryRetain(const void *aBlock);
 
 // Helper structure
 struct Block_layout