Naive parallel Sieve of Eratosthenes in Java

Question

My naive version now is too slow. I think setting/accessing concurrent atomic bit is way slower compared to access/modify an array of boolean. Second, the parallel execution only happens on the beginning of the process. For bigger numbers, thread-pool is not really active.

public class EratosthenesSieve {

    private final ExecutorService executorService;
    private final int n;
    private final AtomicBitSet bitSet;
    private final int squareRoot;
    private final Node head;
    private final Node trailer;

    public EratosthenesSieve(int n) {
        this.n = n;
        this.squareRoot = (int) Math.sqrt(n);
        bitSet = new AtomicBitSet(n + 1);
        executorService = Executors.newCachedThreadPool();
        Node h = new Node(0, null, null);
        Node t = new Node(0, null, h);
        h.setNext(t);
        head = h;
        trailer = t;

    }

    public Node addLast(int i) {
        Node node;
        while ((node = trailer.prepend(i)) == null)
            ;
        return node;
    }

    class FlagIdeal
        implements Runnable {
        private final int candidate;

        FlagIdeal(int candidate) {
            this.candidate = candidate;
        }

        @Override
        public void run() {
            int j = bitSet.firstZeroAfter(candidate);
            if (j >= n) {
                executorService.shutdown();
                return;
            }
            int prime = j;
            for (Node<Integer> iterator = head.forward(); iterator != null; iterator = iterator.forward()) {
                // with a good scheduler, it should never happen
                if (prime > iterator.element) {
                    executorService.submit(this);
                    return;
                }
            }
            // after the iteration over the list of restrictions, something may change
            if (bitSet.get(prime)) {
                executorService.submit(this);
                return;
            }
            // work is over
            if (prime > squareRoot) {
                executorService.shutdown();
                return;
            }
            int lastNumber = prime * prime;
            // before adding work, add a new restriction
            Node activeRestriction = addLast(lastNumber);
            executorService.submit(new FlagIdeal(prime + 1));

            // now, mark all multiples of this prime number starting from the square
            bitSet.set(lastNumber);
            lastNumber = lastNumber + prime;
            while (lastNumber <= n) {
                bitSet.set(lastNumber);
                activeRestriction.element = lastNumber;
                lastNumber = lastNumber + prime;
            }
            // remove the restriction
            while (!activeRestriction.delete() && !activeRestriction.isDeleted())
                ;

        }

        private void surroundedSleep(int i) {
            try {
                Thread.sleep(i);
            } catch (InterruptedException e) {
                new RuntimeException(e);
            }
        }
    }

    public static void main(String[] args) throws InterruptedException {
        EratosthenesSieve eratosthenesSieve = new EratosthenesSieve(100000000);
        long initial = System.nanoTime();
        eratosthenesSieve.execute();
        System.out.println(System.nanoTime() - initial);

    }

    private void execute() throws InterruptedException {
        executorService.submit(new FlagIdeal(2));
        while (!executorService.isTerminated()) {
            Thread.sleep(100);
        }

        int count = 0;
        for (int i = 2; i < n; i++) {
            if (!bitSet.get(i)) {
                count++;
                // System.out.print(i);
                // System.out.print(' ');
            }
        }
        System.out.println();
        System.out.println(count);
    }
}

Here, kind of the bottleneck

public class AtomicBitSet {
    public static final int CHUNK_SIZE = 64;
    public static final int CHUNK_SIZE_MINUS_ONE = 63;
    public static final int CHUNK_BITS = 6;
    // Adapted <From StackOverflow url="http://stackoverflow.com/a/12425007/1879686" author="Peter Lawrey">
    private final AtomicLongArray array;

    public AtomicBitSet(int length) {
        int intLength = (length + CHUNK_SIZE_MINUS_ONE) / CHUNK_SIZE;
        array = new AtomicLongArray(intLength);
    }

    public void set(long n) {
        long bit = 1l << n;
        int idx = (int) (n >>> CHUNK_BITS);
        while (true) {
            long num = array.get(idx);
            long  num2 = num | bit;
            if (num == num2 || array.compareAndSet(idx, num, num2))
                return;
        }
    }

    public boolean get(long n) {
        long bit = 1l << n;
        int idx = (int) (n >>> CHUNK_BITS);
        long num = array.get(idx);
        return (num & bit) != 0;
    }
    // </From StackOverflow>

    public int firstZeroAfter(long n){
        int idx = (int) (n >>> CHUNK_BITS);
        long bit = 1l << (n % CHUNK_SIZE);
        long oneMask = bit - 1;

        long num = array.get(idx);
        int trailingOnes = Long.numberOfTrailingZeros(~num & ~oneMask);

        while(trailingOnes == CHUNK_SIZE && idx < array.length() -1){
            idx++;
            num = array.get(idx);
            trailingOnes = Long.numberOfTrailingZeros(~num);
        }
        return (idx << CHUNK_BITS) + trailingOnes;
    }
}

Suggestions?

Trivia: Parallel execution safety

I am using a ConcurrentDoubleLinkedList written by Doug Lea with assistance from members of JCP JSR-166 Expert Group to "lock" the bigger number that can be safely considered as a prime. When you read Node in the code, it is a node of this list.

Below is the Node slightly modified:

import java.util.concurrent.atomic.AtomicReference;

/**
 * Linked Nodes. As a minor efficiency hack, this class opportunistically inherits from AtomicReference, with the atomic
 * ref used as the "next" link.
 *
 * Nodes are in doubly-linked lists. There are three kinds of special nodes, distinguished by: * The list header has a
 * null prev link * The list trailer has a null next link * A deletion marker has a prev link pointing to itself. All
 * three kinds of special nodes have null element fields.
 *
 * Regular nodes have non-null element, next, and prev fields. To avoid visible inconsistencies when deletions overlap
 * element replacement, replacements are done by replacing the node, not just setting the element.
 *
 * Nodes can be traversed by read-only ConcurrentLinkedDeque class operations just by following raw next pointers, so
 * long as they ignore any special nodes seen along the way. (This is automated in method forward.) However, traversal
 * using prev pointers is not guaranteed to see all live nodes since a prev pointer of a deleted node can become
 * unrecoverably stale.
 */
public class Node
    extends AtomicReference<Node> {
    private volatile Node prev;

    int element;

    /** Creates a node with given contents */
    Node(int element, Node next, Node prev) {
        super(next);
        this.prev = prev;
        this.element = element;
    }

    /** Creates a marker node with given successor */
    Node(Node next) {
        super(next);
        this.prev = this;
        this.element = 0;
    }

    /**
     * Gets next link (which is actually the value held as atomic reference).
     */
    private Node getNext() {
        return get();
    }

    /**
     * Sets next link
     *
     * @param n
     *            the next node
     */
    void setNext(Node n) {
        set(n);
    }

    /**
     * compareAndSet next link
     */
    private boolean casNext(Node cmp, Node val) {
        return compareAndSet(cmp, val);
    }

    /**
     * Gets prev link
     */
    private Node getPrev() {
        return prev;
    }

    /**
     * Sets prev link
     *
     * @param b
     *            the previous node
     */
    void setPrev(Node b) {
        prev = b;
    }

    /**
     * Returns true if this is a header, trailer, or marker node
     */
    boolean isSpecial() {
        return element == 0;
    }

    /**
     * Returns true if this is a trailer node
     */
    boolean isTrailer() {
        return getNext() == null;
    }

    /**
     * Returns true if this is a header node
     */
    boolean isHeader() {
        return getPrev() == null;
    }

    /**
     * Returns true if this is a marker node
     */
    boolean isMarker() {
        return getPrev() == this;
    }

    /**
     * Returns true if this node is followed by a marker, meaning that it is deleted.
     *
     * @return true if this node is deleted
     */
    boolean isDeleted() {
        Node f = getNext();
        return f != null && f.isMarker();
    }

    /**
     * Returns next node, ignoring deletion marker
     */
    private Node nextNonmarker() {
        Node f = getNext();
        return (f == null || !f.isMarker()) ? f : f.getNext();
    }

    /**
     * Returns the next non-deleted node, swinging next pointer around any encountered deleted nodes, and also patching
     * up successor''s prev link to point back to this. Returns null if this node is trailer so has no successor.
     *
     * @return successor, or null if no such
     */
    Node successor() {
        Node f = nextNonmarker();
        for (;;) {
            if (f == null)
                return null;
            if (!f.isDeleted()) {
                if (f.getPrev() != this && !isDeleted())
                    f.setPrev(this); // relink f's prev
                return f;
            }
            Node s = f.nextNonmarker();
            if (f == getNext())
                casNext(f, s); // unlink f
            f = s;
        }
    }

    /**
     * Returns the apparent predecessor of target by searching forward for it starting at this node, patching up
     * pointers while traversing. Used by predecessor().
     *
     * @return target's predecessor, or null if not found
     */
    private Node findPredecessorOf(Node target) {
        Node n = this;
        for (;;) {
            Node f = n.successor();
            if (f == target)
                return n;
            if (f == null)
                return null;
            n = f;
        }
    }

    /**
     * Returns the previous non-deleted node, patching up pointers as needed. Returns null if this node is header so has
     * no successor. May also return null if this node is deleted, so doesn't have a distinct predecessor.
     *
     * @return predecessor or null if not found
     */
    Node predecessor() {
        Node n = this;
        for (;;) {
            Node b = n.getPrev();
            if (b == null)
                return n.findPredecessorOf(this);
            Node s = b.getNext();
            if (s == this)
                return b;
            if (s == null || !s.isMarker()) {
                Node p = b.findPredecessorOf(this);
                if (p != null)
                    return p;
            }
            n = b;
        }
    }

    /**
     * Returns the next node containing a nondeleted user element. Use for forward list traversal.
     *
     * @return successor, or null if no such
     */
    Node forward() {
        Node f = successor();
        return (f == null || f.isSpecial()) ? null : f;
    }

    /**
     * Returns previous node containing a nondeleted user element, if possible. Use for backward list traversal, but
     * beware that if this method is called from a deleted node, it might not be able to determine a usable predecessor.
     *
     * @return predecessor, or null if no such could be found
     */
    Node back() {
        Node f = predecessor();
        return (f == null || f.isSpecial()) ? null : f;
    }

    /**
     * Tries to insert a node holding element as successor, failing if this node is deleted.
     *
     * @param element
     *            the element
     * @return the new node, or null on failure.
     */
    Node append(int element) {
        for (;;) {
            Node f = getNext();
            if (f == null || f.isMarker())
                return null;
            Node x = new Node(element, f, this);
            if (casNext(f, x)) {
                f.setPrev(x); // optimistically link
                return x;
            }
        }
    }

    /**
     * Tries to insert a node holding element as predecessor, failing if no live predecessor can be found to link to.
     *
     * @param element
     *            the element
     * @return the new node, or null on failure.
     */
    Node prepend(int element) {
        for (;;) {
            Node b = predecessor();
            if (b == null)
                return null;
            Node x = new Node(element, this, b);
            if (b.casNext(this, x)) {
                setPrev(x); // optimistically link
                return x;
            }
        }
    }

    /**
     * Tries to mark this node as deleted, failing if already deleted or if this node is header or trailer
     *
     * @return true if successful
     */
    boolean delete() {
        Node b = getPrev();
        Node f = getNext();
        if (b != null && f != null && !f.isMarker() && casNext(f, new Node(f))) {
            if (b.casNext(this, f))
                f.setPrev(b);
            return true;
        }
        return false;
    }

    /**
     * Tries to insert a node holding element to replace this node. failing if already deleted.
     *
     * @param newElement
     *            the new element
     * @return the new node, or null on failure.
     */
    Node replace(int newElement) {
        for (;;) {
            Node b = getPrev();
            Node f = getNext();
            if (b == null || f == null || f.isMarker())
                return null;
            Node x = new Node(newElement, f, b);
            if (casNext(f, new Node(x))) {
                b.successor(); // to relink b
                x.successor(); // to relink f
                return x;
            }
        }
    }

    @Override
    public String toString() {
        return Integer.toString(element);
    }
}

Answer 1

Broken Code

There seems to be a great deal of confusion about the Node class. The EratosthenesSieve class you posted does not compile against the Node class you later posted, on account of the node class not being generic whereas your sieve class contains code that assumes it is. The sieve class also does not compile against the linked version of the Node class (which is generic), on account of Node.element being final in that version, but the sieve class attempting to modify that field.

The posted Node class has similar form as the linked one, but also considerable differences -- primarily, the former has been made specific to elements of type int instead of generic, but it has also been reformatted, been made public, and its Node.element has been made non-final. This review is based on the version you provided, assuming the aforementioned error in EratosthenesSieve is fixed (which can easily be done).

Naming

My apologies, but your naming is overall terrible, and that makes your code much more difficult to understand than it needs to be. Field and local variable names for the most part are meaningless or simply reflect data type; few of them clearly reflect the the significance or intended use of their values. Even the meaning of the name of inner class FlagIdeal is obscure, plus it appears to be a verb -- class names should be nouns. Some methods have obscure names, too.

In addition to being unclear, the names head and trailer in particular are also non-apposite. As a pair, they should be either head and tail or header and trailer.

Code Layout

I strongly recommend a more conventional layout with respect to placement of the inner class. In your code, it is placed between two methods of the main class, and that makes it harder than it needs to be to follow which method belongs to which class. If you're going to declare an inner class, then the best place for it is after all the containing class's fields and methods.

Correctness

The adaptations to the Node class make it possible to use it in a manner that is not thread safe, and it is in fact used in an unsafe manner because instances' element fields being accessed and modified by different threads without synchronization. This can be rescued by making the field volatile, but that substantially slows the program (as should be expected).

Even if it were implemented in a thread-safe way, your "restriction" mechanism would still have a race condition, because your code lies. The check against the restrictions interprets the restriction Node values as representing the last numbers flagged as composite by each running task, but when you initially create each one, you set its value to the first composite the task intends to mark, and then queue the next task before that composite is actually marked. It is possible for threads to be scheduled such that the next task starts before the one that enqueued it actually marks the designated number as composite.

Algorithmic Improvements

The most work is required to sieve the smallest primes, and the most work of all is required to sieve 2. You can efficiently handle 2 as a special case by initializing all the chunks in your bit set to the bit pattern 0x5555555555555555L, thus striking out all the even numbers very efficiently, and then starting at 3 for the rest of the algorithm. This modification speeds your original code by more than 30% for me.

That's in the same vein as your own suggested improvement. Yours is more comprehensive, and with the threshold for enabling the optimization set to 32, my variation on that approach cuts the elapsed time in half. The two can be combined, but using them together provides only a marginal additional improvement over yours alone, so that's probably not worth it. One important thing to glean from that, however, is that atomic operations and operations on volatile data are expensive.

You can get another significant boost by getting rid of your "restrictions" and the associated Node objects altogether, along with the operations on its necessarily volatile element. The objective of those is to ensure that each previous FlagIdeal task has progressed far enough that a given one can correctly identify the next prime from the sieve. We can avoid the need for that by simply not enqueuing the next FlagIdeal task until we are confident that it can immediately proceed.

We cannot know the earliest time when that is possible, but we do know that it is possible if the current task has already processed all values up to at least the square root of the sieve size (recognizing that the current task did not start until the previous task proceeded that far, etc.). That's an easy criterion to check, and it leaves plenty of room for many concurrent tasks, because the square root of the sieve size forms a much smaller proportion of the total sieve size than 1/(# cores). That optimization cuts the runtime by about 20% relative to my variation on your optimization alone, for a total savings of about 65%. The code is much simplified, too.

I have not tested this one, but I'm inclined to think that further speedup could be obtained by doing away with the AtomicBitSet altogether, and instead using an analog that relies on an ordinary array. To make that thread safe and still performant, you would implement coarse-grained locks protecting large sections of the array -- for example, blocks of sqrt(sieveMax) elements. That could trivialize the previous optimization, and although there would be increased contention between threads, I anticipate that the replacement of millions of atomic operations with normal operations plus a (comparatively) few lock operations would be a big win.

Additional Comments

• There is a bit of inconsistency about the meaning of EratosthenesSieve.n. In some places it is directly or indirectly handled as an inclusive upper bound on the numbers to be sieved, but in other places it is treated as an exclusive upper bound. In practice, though, this matters only if you happen to choose a prime for that value.

• Its unclear why you use temporary variables h and t in the EratosthenesSieve constructor. It would be simpler and clearer to initialize the head and trailer members directly.

• EratosthenesSieve.execute() waits for executaion to complete by sleeping for 100 ms at a time and checking each time it wakes whether the ExecutorService is terminated. That's a weird way of doing it when there is ExecutorService.awaitTermination(), which is clearer, and which also allows you to set an overall timeout.

• It's really inefficient to count the the number of elements in your AtomicBitSet by iterating over its entire range -- you end up performing 64 times as many (expensive) atomic reads as you need to do, not to mention a bunch of redundant arithmetic. You could instead read each chunk exactly once, and count the bits efficiently with, say, Long.bitCount(). Just make sure to mark 0 and 1 as composites first, or else subtract 2 from the result.

• Inner class FlagIdeal contains a private method surroundedSleep() that is never invoked

Update

Here is the fastest code I have been able to come up with. Note that it does not rely on atomic operations and it does not use any java.util.concurrent.locks.Lock implementation -- I tried several variations employing such objects, but in the end I found that structuring the code to use native synchronization on suitable coarse-grained objects produced the best result. The below code runs in about 10% of the time of the original code on my test system.

It should be noted, however, that the below code is only about twice as fast as the best sequential code I could come up with, so much of the performance gain is from general algorithmic improvement (especially getting rid of all the atomic operations, I suspect), not so much from improving the parallelism. The Sieve of Eratosthenes is simply not well suited to parallelization, though this code seems to demonstrate that it can obtain at least some advantage.

public class EratosthenesSieve {

    // This is a tuning parameter; performance degrades if it gets too large,
    // presumably because of additional synchronization overhead, but less
    // concurrency is possible the smaller it is.
    private final static int MAX_SEGMENTS = 24;

    // This is a tuning parameter affecting which primes' multiples are
    // accumulated into groups for joint flagging, and which are applied
    // individually.  Values less than 5 completely disable the
    // accumulator strategy, and values greater than 61 are unlikely to
    // be beneficial.
    private final static int ACCUMULATOR_STRATEGY_LIMIT = 32;

    private final static long WARMUP_TIME = 30000000000L;

    private final ExecutorService executorService;
    private final int sieveMax;
    private final int sqrtMax;
    private final SegmentedBitSet composites;
    private final long[][] compositeSegments;
    private final int segmentLength;
    private final int segmentDepth;

    public EratosthenesSieve(int max) {
        executorService = Executors.newCachedThreadPool();
        sieveMax = max;
        sqrtMax = (int) Math.sqrt(max);

        composites = new SegmentedBitSet(max + 1,
                Math.min((max / (sqrtMax + 1)) + 1, MAX_SEGMENTS),
                0xD75D75D75D75D75DL, 0x5D75D75D75D75D75L, 0x75D75D75D75D75D7L);
        composites.set(1);
        composites.clear(2);
        composites.clear(3);
        compositeSegments = composites.getSegments();
        segmentLength = compositeSegments[0].length;
        segmentDepth = segmentLength * SegmentedBitSet.CHUNK_SIZE;
    }

    public static void main(String[] args) throws InterruptedException {
        long initial = System.nanoTime();
        long total = 0L;

        System.err.println("Warming up:");
        do {        
            initial = System.nanoTime();
            new EratosthenesSieve(100000000).execute();
            total += (System.nanoTime() - initial);
            System.gc();
        } while (total < WARMUP_TIME);

        System.err.println("\nTesting:");
        initial = System.nanoTime();
        new EratosthenesSieve(100000000).execute();
        System.err.print("elapsed (ns): ");
        System.err.println(System.nanoTime() - initial);
    }

    private void execute() throws InterruptedException {
        executorService.submit(new CompositeFlagger(5));
        if (!executorService.awaitTermination(10, TimeUnit.SECONDS)) {
            System.err.println("Execution timed out - aborting");
            executorService.shutdownNow();
            return;
        }

        System.err.print("Found ");
        System.err.print(1 + sieveMax - composites.countBits());
        System.err.println(" primes");
    }

    class CompositeFlagger implements Runnable {

        private final int candidate;

        CompositeFlagger(int candidate) {
            this.candidate = candidate;
        }

        private int markSegmentStandard(long[] segment, int segmentIndex,
                int composite, int prime) {
            int segmentBound = (segmentIndex + 1) * segmentLength * SegmentedBitSet.CHUNK_SIZE;

            while (composite < segmentBound) {
                int chunkIndex = composite >>> SegmentedBitSet.CHUNK_BITS;
                segment[chunkIndex % segmentLength] |= (1L << composite);
                composite += prime;
            }

            return composite;
        }

        private int markSegmentAccumulator(long[] segment, int segmentIndex,
                int composite, int prime) {
            int segmentBound = (segmentIndex + 1) * segmentLength * SegmentedBitSet.CHUNK_SIZE;
            int chunkIndex = composite >>> SegmentedBitSet.CHUNK_BITS;
            int lastChunkIndex = chunkIndex;
            long accumulator = 1L << composite;

            do {
                composite += prime;
                chunkIndex = composite >>> SegmentedBitSet.CHUNK_BITS;

                if (chunkIndex == lastChunkIndex) {
                    accumulator |= (1L << composite);
                } else {
                    segment[lastChunkIndex % segmentLength] |= accumulator;
                    lastChunkIndex = chunkIndex;
                    accumulator = (1L << composite);
                }
            } while (composite < segmentBound);

            return composite;
        }

        @Override
        public void run() {
            int prime;
            int composite;

            synchronized (compositeSegments[0]) { 
                prime = composites.firstZeroAfter(candidate);
                composite = prime * prime;

                assert(prime <= sieveMax);
                if (composite > sieveMax) {
                    executorService.shutdown();
                    return;
                }

                // mark all multiples of the present prime in the first segment
                if (prime > ACCUMULATOR_STRATEGY_LIMIT) {
                    composite = markSegmentStandard(compositeSegments[0], 0, composite, prime);
                } else {
                    composite = markSegmentAccumulator(compositeSegments[0], 0, composite, prime);
                }
            }

            // queue the next task
            executorService.submit(new CompositeFlagger(prime + 1));

            // mark all multiples of the present prime in the remaining segments
            for (int segmentIndex = 1; segmentIndex < compositeSegments.length; segmentIndex++) {
                synchronized (compositeSegments[segmentIndex]) {
                    if (prime > ACCUMULATOR_STRATEGY_LIMIT) {
                        composite = markSegmentStandard(compositeSegments[segmentIndex],
                                segmentIndex, composite, prime);
                    } else {
                        composite = markSegmentAccumulator(compositeSegments[segmentIndex],
                                segmentIndex, composite, prime);
                    }
                }
            }
        }
    }
}

class SegmentedBitSet {
    // The number of bits in a long:
    public static final int CHUNK_SIZE = 64;

    // The number of bit-number bits that express an offset within a chunk
    public static final int CHUNK_BITS = 6;

    private final long[][] segments;
    private final int segmentSize;
    private final int lastSegmentSize;
    private final long length;

    public SegmentedBitSet(long length, int numSegments) {
        this(length, numSegments, new long[0]);
    }

    public SegmentedBitSet(long length, int numSegments, long... chunkPattern) {
        long totalChunks = (length + CHUNK_SIZE - 1) / CHUNK_SIZE;

        this.length = length;
        segmentSize = (int) ((totalChunks + numSegments - 1) / numSegments);
        segments = new long[numSegments][segmentSize];

        if ((totalChunks % segmentSize) != 0) {
            lastSegmentSize = (int) (totalChunks % segmentSize);
            segments[numSegments - 1] = new long[lastSegmentSize];
        } else {
            lastSegmentSize = segmentSize;
        }

        if (chunkPattern.length != 0) {
            for (long i = 0; i < totalChunks; i++) {
                segments[(int) (i / segmentSize)][(int) (i % segmentSize)]
                        = chunkPattern[(int) (i % chunkPattern.length)];
            }

            if (length % CHUNK_SIZE != 0) {
                segments[segments.length - 1][(int) ((totalChunks - 1) % segmentSize)]
                        &= ~(-1L << length);
            }
        }
    }

    public void set(long n) {
        long bit = 1l << n;
        int chunkIndex = (int) (n >>> CHUNK_BITS);

        segments[chunkIndex / segmentSize][chunkIndex % segmentSize] |= bit;
    }

    public void clear(long n) {
        long bit = 1l << n;
        int chunkIndex = (int) (n >>> CHUNK_BITS);

        segments[chunkIndex / segmentSize][chunkIndex % segmentSize] &= ~bit;
    }

    public boolean get(long n) {
        long bit = 1l << n;
        int chunkIndex = (int) (n >>> CHUNK_BITS);

        return ((segments[chunkIndex / segmentSize][chunkIndex % segmentSize] & bit) != 0);
    }

    public long[][] getSegments() {
        return segments;
    }

    public int firstZeroAfter(long n) {
        int chunkIndex = (int) (n >>> CHUNK_BITS);
        int segmentIndex = chunkIndex / segmentSize;
        int chunkOffset = chunkIndex % segmentSize;
        long chunk = segments[segmentIndex][chunkOffset];
        int trailingOnes = Long.numberOfTrailingZeros(~chunk & (-1L << n));

        while((trailingOnes == CHUNK_SIZE)
                && ((segmentIndex < segments.length - 1)
                        || (chunkOffset < segments[segmentIndex].length - 1))) {
            chunkOffset = (chunkOffset + 1) % segmentSize;
            segmentIndex += (chunkOffset == 0) ? 1 : 0;
            chunk = segments[segmentIndex][chunkOffset];
            trailingOnes = Long.numberOfTrailingZeros(~chunk);
        }

        return ((segmentIndex * segmentSize + chunkOffset) << CHUNK_BITS) + trailingOnes;
    }

    public int countBits() {
        int lastBits = (int) (((length - 1) % CHUNK_SIZE) + 1);
        int count = 0;

        for (long[] segment : segments) {
            for (long chunk : segment) {
                count += Long.bitCount(chunk);
            }
        }

        return count - Long.bitCount(
                segments[segments.length - 1][lastSegmentSize - 1] & (-1L << lastBits));
    }
}

Answer 2

Accumulator Strategy

One minor optimization can beat the single processing in a long run (\$n >= 100000000\$). Instead of update the directly AtomicSet for each non-prime number, you can update an intermediate variable with same size of each chunk of the AtomicSet and using this intermediate value to update the chunk once.

        // now, mark all multiples of this prime number starting from the square
        int idx = 0;
        long  num = 0;
        if (prime < OPTIMIZATION){
            long bit = 1l << lastNumber;
            int remainingBits = lastNumber % AtomicBitSet.CHUNK_SIZE;
            idx = lastNumber >>> AtomicBitSet.CHUNK_BITS;
            num = bit;


            //bitSet.set(lastNumber); let's avoid update the bitset too often
            for (lastNumber = lastNumber + prime; lastNumber <= n; lastNumber = lastNumber + prime) {
                remainingBits += prime;
                bit = 1l  << lastNumber;
                if (remainingBits >= AtomicBitSet.CHUNK_SIZE){
                    bitSet.set(idx, num);
                    idx++;
                    num = bit;
                    activeRestriction.element = lastNumber;
                    remainingBits = remainingBits % AtomicBitSet.CHUNK_SIZE;
                } else {
                    num = num | bit;
                }
            }
        } else {
            bitSet.set(lastNumber);

            for (lastNumber = lastNumber + prime; lastNumber <= n; lastNumber = lastNumber + prime) {
                bitSet.set(lastNumber);
                activeRestriction.element = lastNumber;
            }
        }
        if (prime <= OPTIMIZATION){
            bitSet.set(idx, num);
        }