Improve Signal state management to avoid an ARC race.

Author: andersio

Sources/Atomic.swift

@@ -14,65 +14,118 @@ import MachO
 /// A simple, generic lock-free finite state machine.
 ///
 /// - warning: `deinitialize` must be called to dispose of the consumed memory.
-internal struct UnsafeAtomicState<State: RawRepresentable> where State.RawValue == Int32 {
-	internal typealias Transition = (expected: State, next: State)
-#if os(macOS) || os(iOS) || os(tvOS) || os(watchOS)
-	private let value: UnsafeMutablePointer<Int32>
+internal struct UnsafeAtomicInt32 {
+	#if os(macOS) || os(iOS) || os(tvOS) || os(watchOS)
+	private let _value: UnsafeMutablePointer<Int32>
+
+	internal var value: Int32 {
+		return _value.pointee
+	}
 
 	/// Create a finite state machine with the specified initial state.
 	///
 	/// - parameters:
 	///   - initial: The desired initial state.
-	internal init(_ initial: State) {
-		value = UnsafeMutablePointer<Int32>.allocate(capacity: 1)
-		value.initialize(to: initial.rawValue)
+	internal init(_ initial: Int32) {
+		_value = UnsafeMutablePointer<Int32>.allocate(capacity: 1)
+		_value.initialize(to: initial)
 	}
 
 	/// Deinitialize the finite state machine.
 	internal func deinitialize() {
-		value.deinitialize()
-		value.deallocate(capacity: 1)
+		_value.deinitialize()
+		_value.deallocate(capacity: 1)
 	}
 
-	/// Compare the current state with the specified state.
+	internal func swap(from expected: Int32, to next: Int32) -> Bool {
+		return OSAtomicCompareAndSwap32Barrier(expected, next, _value)
+	}
+
+	/// Increment the atomic integer.
 	///
-	/// - parameters:
-	///   - expected: The expected state.
+	/// - returns: The new value.
+	@discardableResult
+	internal func increment() -> Int32 {
+		return OSAtomicIncrement32Barrier(_value)
+	}
+
+	/// Decrement the atomic integer.
 	///
-	/// - returns: `true` if the current state matches the expected state.
-	///            `false` otherwise.
-	internal func `is`(_ expected: State) -> Bool {
-		return OSAtomicCompareAndSwap32Barrier(expected.rawValue,
-		                                       expected.rawValue,
-		                                       value)
+	/// - returns: The new value.
+	@discardableResult
+	internal func decrement() -> Int32 {
+		return OSAtomicDecrement32Barrier(_value)
 	}
+	#else
+	private let _value: Atomic<Int32>
 
-	/// Try to transition from the expected current state to the specified next
-	/// state.
+	private var value: Int32 {
+		return _value.pointee
+	}
+
+	/// Create a finite state machine with the specified initial state.
 	///
 	/// - parameters:
-	///   - expected: The expected state.
-	///   - next: The state to transition to.
+	///   - initial: The desired initial state.
+	internal init(_ initial: Int32) {
+		_value = Atomic(initial)
+	}
+
+	/// Deinitialize the finite state machine.
+	internal func deinitialize() {}
+
+	internal func swap(from expected: Int32, to next: Int32) -> Bool {
+		return _value.modify { value in
+			guard value == expected else { return false }
+			value = next
+			return true
+		}
+	}
+
+	/// Increment the atomic integer.
 	///
-	/// - returns: `true` if the transition succeeds. `false` otherwise.
-	internal func tryTransition(from expected: State, to next: State) -> Bool {
-		return OSAtomicCompareAndSwap32Barrier(expected.rawValue,
-		                                       next.rawValue,
-		                                       value)
+	/// - returns: The new value.
+	@discardableResult
+	internal func increment() -> Int32 {
+		return _value.modify { value in
+			value += 1
+			return value
+		}
 	}
-#else
-	private let value: Atomic<Int32>
+
+	/// Decrement the atomic integer.
+	///
+	/// - returns: The new value.
+	@discardableResult
+	internal func decrement() -> Int32 {
+		return _value.modify { value in
+			value -= 1
+			return value
+		}
+	}
+	#endif
+}
+
+/// A simple, generic lock-free finite state machine.
+///
+/// - warning: `deinitialize` must be called to dispose of the consumed memory.
+internal struct UnsafeAtomicState<State: RawRepresentable> where State.RawValue == Int32 {
+	internal typealias Transition = (expected: State, next: State)
+
+	private let value: UnsafeAtomicInt32
 
 	/// Create a finite state machine with the specified initial state.
 	///
 	/// - parameters:
 	///   - initial: The desired initial state.
 	internal init(_ initial: State) {
-		value = Atomic(initial.rawValue)
+		value = UnsafeAtomicInt32(initial.rawValue)
 	}
 
 	/// Deinitialize the finite state machine.
-	internal func deinitialize() {}
+	internal func deinitialize() {
+		value.deinitialize()
+	}
 
 	/// Compare the current state with the specified state.
 	///
@@ -82,26 +135,20 @@ internal struct UnsafeAtomicState<State: RawRepresentable> where State.RawValue
 	/// - returns: `true` if the current state matches the expected state.
 	///            `false` otherwise.
 	internal func `is`(_ expected: State) -> Bool {
-		return value.modify { $0 == expected.rawValue }
+		return value.swap(from: expected.rawValue, to: expected.rawValue)
 	}
 
 	/// Try to transition from the expected current state to the specified next
 	/// state.
 	///
 	/// - parameters:
 	///   - expected: The expected state.
+	///   - next: The state to transition to.
 	///
 	/// - returns: `true` if the transition succeeds. `false` otherwise.
 	internal func tryTransition(from expected: State, to next: State) -> Bool {
-		return value.modify { value in
-			if value == expected.rawValue {
-				value = next.rawValue
-				return true
-			}
-			return false
-		}
+		return value.swap(from: expected.rawValue, to: next.rawValue)
 	}
-#endif
 }
 
 /// `Lock` exposes `os_unfair_lock` on supported platforms, with pthread mutex as the