Test & Set

Working with threads always feels hazardous to me.

Non-determinism limits how much we can exercise with simple tests. Stress tests can shake out some bugs, but they’re slow and potentially flaky. I don’t like flaky tests.

So I have a few concurrency patterns that I trust, and I use ’em a lot. AtomicReference is one of them. It’s got a lot of concurrency power and a easy-enough API. I can make sense of programs that use AtomicReference.

The drawback of AtomicReference is that I usually need a while loop to update it. The atomic compareAndSet() function returns false if my thread loses a race and must retry. After writing a whole bunch of these tricky loops yesterday, I factored out a helper testAndSet() function in order to split difficult business logic from difficult concurrency code.

In one situation, my code sets a state variable to Canceled, but only if the current state is cancelable. It must also tell the task it has been canceled.

private val state = AtomicReference<State>(State.Idle)

override fun cancel() {
  while (true) {
    val previous = state.get()
    if (previous is State.Running || previous == State.Idle) {
      if (!state.compareAndSet(previous, State.Canceled)) {
        continue // Lost a race.
      }
      (previous as? State.Running)
        ?.callback
        ?.onFailure(this, IOException("canceled"))
    }
    return
  }
}

That loop really complicates this code! And it’s an ugly loop, with a return break out at the bottom.

Here’s the cleaned up version that uses testAndSet():

override fun cancel() {
  val previous = state.testAndSet(
    condition = { it is State.Running || it == State.Idle },
    newValue = State.Canceled,
  )
  (previous as? State.Running)
    ?.callback
    ?.onFailure(this, IOException("canceled"))
}

Getting the signature for this was difficult! Once I realized that the function should return the previous value (and not a Boolean), everything snapped into place.

The testAndSet() code is a tad bit fancy ’cause I’m using tailrec instead of a conventional loop.

/**
 * Updates this reference to [newValue] if the previous value matches
 * [condition].
 *
 * It is possible that a [condition] is invoked multiple times in a
 * single call to this function. This will occur if this state object
 * is updated between fetching previous value and replacing it.
 *
 * Returns the previous value. This will always be the replaced value
 * if the value was updated.
 */
internal tailrec fun <T> AtomicReference<T>.testAndSet(
  newValue: T,
  condition: (T) -> Boolean,
): T {
  val previous = get()
  if (!condition(previous)) return previous
  if (compareAndSet(previous, newValue)) return previous
  return testAndSet(newValue, condition) // Lost a race, retry.
}

I’m particularly proud of how this small function successfully separates business logic and concurrency logic.

Programming. It’s fun.