BinaryLatch.java

/*
 * Written by Chris Vest and released to the public domain, as explained at
 * http://creativecommons.org/publicdomain/zero/1.0/
 */
package binarylatch;

import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.LockSupport;

/**
 * Similar to a CountDownLatch(1), but faster because it can optimise for this special case.
 */
public interface BinaryLatch
{
    /**
     * Create a new BinaryLatch instance.
     */
    public static BinaryLatch create()
    {
        return new BinaryLatchImpl();
    }

    /**
     * Release the latch, thereby unblocking all current and future calls to {@link #await()}.
     */
    void release();

    /**
     * Wait for the latch to be released, blocking the current thread if necessary.
     *
     * This method returns immediately if the latch has already been released.
     */
    void await();
}

final class BinaryLatchImpl extends AtomicReference<BinaryLatchImpl.Node> implements BinaryLatch
{
    static class Node
    {
        volatile Node next;
    }

    private static final class Waiter extends Node
    {
        final Thread waitingThread = Thread.currentThread();
        volatile boolean successor;
    }

    private static final Node end = new Node();
    private static final Node released = new Node();

    @Override
    public void release()
    {
        // Once the release sentinel is on the stack, it can never (observably) leave.
        // Waiters might accidentally remove the released sentinel from the stack for brief periods of time, but then
        // they are required to fix the situation and put it back.
        // Atomically swapping the release sentinel onto the stack will give us back all the waiters, if any.
        Node waiters = getAndSet( released );
        if ( waiters == null )
        {
            // There are no waiters to unpark, so don't bother.
            return;
        }
        unparkSuccessor( waiters );
    }

    private void unparkSuccessor( Node waiters )
    {
        if ( waiters.getClass() == Waiter.class )
        {
            Waiter waiter = (Waiter) waiters;
            waiter.successor = true;
            LockSupport.unpark( waiter.waitingThread );
        }
    }

    @Override
    public void await()
    {
        // Put in a local variable to avoid volatile reads we don't need.
        Node state = get();
        if ( state != released )
        {
            // The latch hasn't obviously already been released, so we want to add a waiter to the stack. Trouble is,
            // we might race with release here, so we need to re-check for release after we've modified the stack.
            Waiter waiter = new Waiter();
            state = getAndSet( waiter );
            if ( state == released )
            {
                // If we get 'released' back from the swap, then we raced with release, and it is our job to put the
                // released sentinel back. Doing so can, however, return more waiters that have added themselves in
                // the mean time. If we find such waiters, then we must make sure to unpark them. Note that we will
                // never get a null back from this swap, because we at least added our own waiter earlier.
                Node others = getAndSet( released );
                // Set our next pointer to 'released' as a signal to other threads who might be going through the
                // stack in the isReleased check.
                waiter.next = released;
                unparkAll( others );
            }
            else
            {
                // It looks like the latch hasn't yet been released, so we are going to park. Before that, we must
                // assign a non-null value to our next pointer, so other threads will know that we have been properly
                // enqueued. We use the 'end' sentinel as a marker when there's otherwise no other next node.
                waiter.next = state == null? end : state;
                do
                {
                    // Park may wake up spuriously, so we have to loop on it until we observe from the state of the
                    // stack, that the latch has been released.
                    LockSupport.park( this );
                }
                while ( !isReleased( waiter ) );
            }
        }
    }

    private boolean isReleased( Waiter waiter )
    {
        // If we are the must recently enqueued waiter on the stack before the release, then that makes us the
        // successor. As the successor, it is our job to wake up all the other threads. We can *only* become the
        // successor if the latch has been released, so there's no need to check anything else in this case.
        if ( waiter.successor )
        {
            unparkAll( waiter.next );
            return true;
        }

        // Otherwise we have to go through the entire stack and look for the 'released' sentinel, since we might be
        // racing with the 'state == released' branch in await.
        Node state = get();
        do
        {
            if ( state == released )
            {
                // We've been released!
                return true;
            }

            Node next;
            do
            {
                // We loop on reading the next pointer because we might observe an enqueued node before its next
                // pointer has been properly assigned. This is a benign race because we know that the next pointer of a
                // properly enqueued node is never null.
                next = state.next;
            }
            while ( next == null );
            state = next;
        }
        while ( state != end );
        // Reaching the end of the stack without seeing 'released' means we're not released.
        return false;
    }

    private void unparkAll( Node waiters )
    {
        // If we find a node that is not a waiter, then it is either 'end' or 'released'. Looking at the type pointer
        // is the cheapest way to make this check.
        while ( waiters.getClass() == Waiter.class )
        {
            Waiter waiter = (Waiter) waiters;
            LockSupport.unpark( waiter.waitingThread );
            Node next;
            do
            {
                // Just like in isReleased, loop if the next pointer is null.
                next = waiters.next;
            }
            while ( next == null );
            waiters = next;
        }
    }
}