Lock-free implementation of getAndUpdate() using atomic CAS (Compare-And-Swap) operation

Question

We have the following class written in Kotlin Native with the new Memory Manager (which doesn't require to freeze objects):

import kotlin.native.concurrent.AtomicReference

class AtomicRef<V>(initialValue: V) {

    private val atom = AtomicReference(initialValue)

    /**
     * Get the current value
     */
    fun get(): V = atom.get()

    /**
     * Atomically compare the current value with `expected`, 
     * and set it to `new` if the current and the expected values are the same.
     *
     * @return `true` if successful, `false` return indicates that
     *   the actual value was not equal to the expected value.
     */
    fun compareAndSet(expected: V, new: V): Boolean = 
        atom.compareAndSet(expected, new)

    /**
     * Atomically compare the current value with `expected`, 
     * and set it to `new` if the current and the expected values are the same.
     *
     * @return the witness value, 
     *   which will be the same as the `expected` value if successful.
     */
    fun compareAndExchange(expected: V, new: V): V =
        atom.compareAndSwap(expected, new)

    /**
     * Atomically updates the current value with the results 
     * of applying the given function, returning the previous value.
     *
     * @param updateFunction a side-effect-free function
     * @return the previous value
     */
    fun getAndUpdate(updateFunction: (V) -> V): V {
        while(true) {
            val prev = get()
            val next = updateFunction(prev)
            if (compareAndSet(prev, next)) return prev
        }
    }
}

compareAndSet and compareAndExchange implementations delegate to the corresponding Kotlin stdlib AtomicReference methods (compareAndSet and compareAndSwap, respectively), which provide native bug- and livelock-free implementations.

getAndUpdate implementation is well known, and can be found, for example, in the Java j.u.c.AtomicReference class, or in the atomicfu library by JetBrains.

I'd like to rewrite the implementation of the getAndUpdate() function to use compareAndExchange and avoid calling get() on every loop iteration.

I've come up with the following implementation:

    fun getAndUpdate(updateFunction: (V) -> V): V {
        var prev = get()
        while(true) {
            val expected = prev
            val next = updateFunction(expected)
            prev = compareAndExchange(expected, next)
            if (prev == expected) return prev
        }
    }

Is it correct, and will it work the same way as the one using compareAndSet?

UPDATE

I've found a ticket with a discussion in the OpenJDK tracker, where one of the OpenJDK members suggests a similar Java implementation for AtomicInteger.getAndUpdate:

public final int getAndUpdate(IntUnaryOperator updateFunction) {
    for (int prev = get() ;;) {
        int next = updateFunction.applyAsInt(prev);
        if (prev == (prev = compareAndExchange(prev, next)))
            return prev;
    } 
}

This seems optimal when compareAndExchange gets compiled to a native instruction. If it does not, and the hardware has spurious failures, compareAndExchange gets compiled to a loop, but there's no wasted extra volatile read on failure because the value read is returned to our loop.

If this is not a good use case for compareAndExchange, then what is? It's currently unused.

The Java code above can be expressed in Kotlin as:

fun getAndUpdate(updateFunction: (V) -> V): V {
    var prev = get()
    while (true) {
        val next = updateFunction(prev)
        val witnessed = compareAndExchange(prev, next)
        if (prev == witnessed) return prev
        prev = witnessed
    }
}

which is quite close to the code I planned to use for getAndUpdate.

Answer 1

You could rewrite the getAndUpdate() function using the compareAndExchange() method as follows:

fun getAndUpdate(updateFunction: (V) -> V): V {
    var prev: V
    var next: V
    do {
       prev = compareAndExchange(null as V, null as V)
       next = updateFunction(prev)
    } while (!compareAndSet(prev, next))
    return prev
}

Here, the getAndUpdate() function now utilizes the compareAndExchange() method. It will attempt to perform the update in a loop until the compareAndSet() method returns true, which means the update was successful. The compareAndExchange() method is called with null values to obtain the current value without modifying it, which effectively replaces the get() call