Lock-free in Swift: ABA Problem

In the previous article, we talked about using atomic primitives. Today, we will cover one of the main challenges in lock-free algorithms: the ABA problem. This issue is hard to detect, and its solutions are often not ideal and can seem tricky.

We have already learned that the loop with compareAndSwap is a popular construction in lock-free algorithms, but this comes with a common issue known as the ABA problem. Since most of these algorithms are based on a loop with the compareAndSwap operation, which compares the expected value with the value at a specific address, the ABA problem can occur when multiple threads work with the shared address simultaneously. Let’s consider a simple example, one thread performs a compareAndSwap on a memory address M that contains a memory pointer A that refers to data C, while another thread changes the data in the following manner:

M = A -> C
M = B -> D
M = A -> E

The issue here is that the same pointer A can refer to the completely different data E which is not equal initial C. But for compareAndSwap in the first thread, it doesn’t make a difference, it checks that M contains expected A, and, as a result, returns true, potentially breaking the consistency of the data structure.

For a better understanding of the ABA problem, let’s check an example of a lock-free structure that intentionally contains a flaw. This is a lock-free stack based on a while loop and compareAndSwap. It’s called ABAStack because it has a vulnerability in its design which makes the ABA problem possible.

The stack has an internal type Node that represents an element in the stack. Each Node has a reference to the next element and a generic value. The stack also holds a reference to the top element, which points to nil when it is initialized. All the references have ManagedAtomic type, allowing atomic operations on them. “Managed” in this context means that Automatic Reference Counting (ARC) will handle memory management, so we don’t need to deallocate memory manually. Additionally, we need to make Node conform to AtomicReference so an instance of the type can be held as ManagedAtomic.

final class ABAStack<T> {
  
  final class Node<T>: AtomicReference {
    var value: T
    var next: ManagedAtomic<Node?>
    
    init(next: Node?, value: T) {
      self.next = ManagedAtomic(next)
      self.value = value
    }
  }
  
  private var top: ManagedAtomic<Node<T>?> = ManagedAtomic(nil)
}

Here, we implement the push method. This is done using a while-loop with compareAndSwap that keeps repeating until the value is successfully set. This approach can introduce overhead compared to traditional algorithms, and in some cases, lock-free algorithms can even be slower. Another common pattern is to use a local reference for the variable within the scope of the loop. This ensures that we work with the same value through the entire iteration. In the code snippet below, we create a new node with a passed parameter. Then, we try to set it as the top element using compareExchange until it succeeds.

func push(value: T) {
  let newNode = Node(next: nil, value: value)
  while(true) {
    let topNode = top.load(ordering: .relaxed)
    newNode.next = ManagedAtomic(topNode)
    
    if top.compareExchange(expected: topNode, desired: newNode, ordering: .releasing).exchanged {
      return
    }
  }
}

In the push method, we use releasing ordering in compareExchange to ensure that after the operation finishes, subsequent operations will use the most recent value. In other words, we are pushing the store buffer to store the value immediately in memory. You can read more about the store buffer in the first lock-free article. To get the topNode, we can use relaxed ordering because if it doesn’t have the most recent value, it will fail at the compareExchange step and then go through another iteration in the while(true) loop.

The pop method contains a similar while loop. Initially, we get a reference to the top element of the stack. Then, we try to exchange the top element with its next element until it succeeds. In the case where the top node is nil, we return nil as well.

func pop() -> T? {
  while(true) {
    guard let topNode = top.load(ordering: .relaxed) else { return nil }
    let nextNode = topNode.next.load(ordering: .relaxed)

    if top.compareExchange(expected: topNode, desired: nextNode, ordering: .acquiring).exchanged {
      return topNode.value
    }
  }
}

We use acquiring ordering in compareExchange because we want the top value not to be rearranged with the previous values. In other, we are emptying the invalidation queue to retrieve it directly from the memory. Similar to the push method, we use relaxed ordering to obtain the topNode. However, in the case of pop, if top is nil (indicating an empty stack), the method will return nil. This scenario is not covered by the while(true) loop in terms of reorderings, but since it’s the initial state and there are no previous values for top, we are free to use relaxed ordering.

It’s worth to mention that determining the appropriate (or even optimal) memory ordering requires careful analysis and can lead to errors that are difficult to reproduce.

Before attempting to reproduce the ABA problem, let’s discuss an important detail about memory allocation. Memory allocation can reuse an address that was previously allocated in the application. This means that if you had a reference to some data and then freed it, that address can be reused in the following memory allocation. This is important because it can be a cause of the ABA problem.

To illustrate it let’s see how can it happen with ABAStack:

let stack = ABAStack<Int>()

stack.push(value: 1)
stack.push(value: 2)
stack.push(value: 3)

let group = DispatchGroup()

// Thread 1
group.enter()
DispatchQueue.global().async {
   _ = stack.pop()
   _ = stack.pop()
   stack.push(value: 3)
   group.leave()
}

// Thread 2
group.enter()
DispatchQueue.global().async {
   _ = stack.pop()
   group.leave()
}

group.wait()

Let’s consider the following situation that can occur with the code:

1. Thread 2 begins execution and encounters a context switch on the line with top.compareExchange(expected: topNode, desired: nextNode, ordering: .releasing) within the pop method.

2. Meanwhile, Thread 1 executes stack.pop(), which returns 3, and then executes stack.pop() again, returning 2. As previously discussed, we’re using automatic memory deallocation through ARC, so let’s assume that the nodes that were withdrawn (3 and 2) have already been deallocated. Next, a push operation occurs. Inside the push method, a new node is allocated, and in this scenario, it returns the same address previously used for the node with value 3. The newly allocated node with the old address becomes the top of the stack.

3. Another context switch occurs, and we resume with Thread 2 inside the pop operation. The compareExchange operation verifies the address of the top and confirms it’s the same as before and swaps the top of the stack with its next, which is the node with value 2. With manual memory management, the next element could have even been deallocated. This situation exemplifies the ABA problem.

It’s important to understand that stack.push(value: 3) could push any other value, and the ABA problem could still occur because we are comparing the addresses of the nodes. Additionally, the ABA problem can happen in more complex scenarios and combinations, which can sometimes be very difficult to reproduce and understand.

Today, we’ve learned about the ABA problem and implemented our first lock-free structure, even with the flaw. There are multiple ways of solving the ABA problem but there is no perfect solution. In the upcoming articles, we will focus on possible solutions to the ABA problem and refine the stack structure accordingly.