remove ThreadStamp

Author: Scottmitch

servicetalk-concurrent-api/src/main/java/io/servicetalk/concurrent/api/PublisherAsBlockingIterable.java

@@ -34,7 +34,7 @@
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
-import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
+import java.util.concurrent.atomic.AtomicLongFieldUpdater;
 import java.util.concurrent.locks.LockSupport;
 import javax.annotation.Nullable;
 
@@ -57,7 +57,7 @@
  * @param <T> Type of items emitted by the {@link Publisher} from which this {@link BlockingIterable} is created.
  */
 final class PublisherAsBlockingIterable<T> implements BlockingIterable<T> {
-    private static final int MAX_OUTSTANDING_DEMAND = min(128, SpscBlockingQueue.ThreadStamp.MAX_DEMAND);
+    private static final int MAX_OUTSTANDING_DEMAND = 128;
     final Publisher<T> original;
     private final int queueCapacityHint;
 
@@ -287,10 +287,17 @@ private static final class SpscBlockingQueue<T> implements BlockingQueue<T> {
                 getLong("io.servicetalk.concurrent.internal.blockingIterableYieldNs", 1024);
 
         @SuppressWarnings("rawtypes")
-        private static final AtomicReferenceFieldUpdater<SpscBlockingQueue, ThreadStamp> threadStampUpdater =
-                AtomicReferenceFieldUpdater.newUpdater(SpscBlockingQueue.class, ThreadStamp.class, "threadStamp");
+        private static final AtomicLongFieldUpdater<SpscBlockingQueue> pcIndexUpdater =
+                AtomicLongFieldUpdater.newUpdater(SpscBlockingQueue.class, "pcIndex");
         private final Queue<T> spscQueue;
-        private volatile ThreadStamp threadStamp = new ThreadStamp(null);
+        @Nullable
+        private Thread consumerThread;
+        /**
+         * high 32 bits == producer index (see {@link #pIndex(long)})
+         * low 32 bits == consumer index (see {@link #cIndex(long)}}
+         * @see #computePCIndex(int, int)
+         */
+        private volatile long pcIndex;
 
         SpscBlockingQueue(Queue<T> spscQueue) {
             this.spscQueue = requireNonNull(spscQueue);
@@ -505,21 +512,15 @@ public String toString() {
         }
 
         private void producerSignalAdded() {
-            ThreadStamp nextStamp = null;
             for (;;) {
-                final ThreadStamp currStamp = threadStamp;
-                if (nextStamp == null) {
-                    nextStamp = new ThreadStamp(null, currStamp.incProducerIndex(), currStamp.consumerIndex);
-                } else {
-                    // nextStamp.thread = null; -> not necessary, this is always the case.
-                    nextStamp.producerIndex = currStamp.incProducerIndex();
-                    nextStamp.consumerIndex = currStamp.consumerIndex;
-                }
-                if (threadStampUpdater.compareAndSet(this, currStamp, nextStamp)) {
-                    if (currStamp.thread != null) {
-                        LockSupport.unpark(currStamp.thread);
-                        // Producer should be equal/behind consumer in order for an unpark to be necessary.
-                        assert currStamp.producerIndex - currStamp.consumerIndex <= 0;
+                final long currIndex = pcIndex;
+                final int producer = pIndex(currIndex);
+                final int consumer = cIndex(currIndex);
+                if (pcIndexUpdater.compareAndSet(this, currIndex, computePCIndex(producer + 1, consumer))) {
+                    if (producer - consumer <= 0 && consumerThread != null) {
+                        final Thread wakeThread = consumerThread;
+                        consumerThread = null;
+                        LockSupport.unpark(wakeThread);
                     }
                     break;
                 }
@@ -528,38 +529,21 @@ private void producerSignalAdded() {
 
         private T take0(BiLongFunction<TimeUnit, T> taker, long timeout, TimeUnit unit) throws InterruptedException {
             final Thread currentThread = Thread.currentThread();
-            ThreadStamp nextStamp = null;
             for (;;) {
-                ThreadStamp currStamp = threadStamp;
-                if (currStamp.producerIndex == currStamp.consumerIndex) {
-                    if (nextStamp == null) {
-                        nextStamp = new ThreadStamp(currentThread, currStamp.producerIndex,
-                                currStamp.incConsumerIndex());
-                    } else {
-                        // nextStamp.thread = currentThread; -> not necessary, this is always the case.
-                        nextStamp.producerIndex = currStamp.producerIndex;
-                        nextStamp.consumerIndex = currStamp.incConsumerIndex();
-                    }
-
-                    if (threadStampUpdater.compareAndSet(this, currStamp, nextStamp)) {
+                long currIndex = pcIndex;
+                final int producer = pIndex(currIndex);
+                final int consumer = cIndex(currIndex);
+                if (producer == consumer) {
+                    // Set consumerThread before pcIndex, to establish happens-before with producer thread.
+                    consumerThread = currentThread;
+                    if (pcIndexUpdater.compareAndSet(this, currIndex, computePCIndex(producer, consumer + 1))) {
                         return taker.apply(timeout, unit);
                     }
                 } else {
                     final T item = spscQueue.poll();
                     if (item != null) {
-                        for (;;) {
-                            if (nextStamp == null) {
-                                nextStamp = new ThreadStamp(null, currStamp.producerIndex,
-                                        currStamp.incConsumerIndex());
-                            } else {
-                                nextStamp.thread = null;
-                                nextStamp.producerIndex = currStamp.producerIndex;
-                                nextStamp.consumerIndex = currStamp.incConsumerIndex();
-                            }
-                            if (threadStampUpdater.compareAndSet(this, currStamp, nextStamp)) {
-                                break;
-                            }
-                            currStamp = threadStamp;
+                        while (!pcIndexUpdater.compareAndSet(this, currIndex, computePCIndex(producer, consumer + 1))) {
+                            currIndex = pcIndex;
                         }
                         return item;
                     }
@@ -569,21 +553,12 @@ private T take0(BiLongFunction<TimeUnit, T> taker, long timeout, TimeUnit unit)
             }
         }
 
-        private void consumerSignalRemoved(int i) {
-            ThreadStamp nextStamp = null;
+        private void consumerSignalRemoved(final int i) {
             for (;;) {
-                final ThreadStamp currStamp = threadStamp;
-                assert (currStamp.producerIndex - currStamp.consumerIndex) + i < ThreadStamp.MAX_DEMAND;
-                if (nextStamp == null) {
-                    nextStamp = new ThreadStamp(null, currStamp.producerIndex,
-                            currStamp.incConsumerIndex(i));
-                } else {
-                    // nextStamp.thread = null; -> not necessary, this is always the case.
-                    nextStamp.producerIndex = currStamp.producerIndex;
-                    nextStamp.consumerIndex = currStamp.incConsumerIndex(i);
-                }
-
-                if (threadStampUpdater.compareAndSet(this, currStamp, nextStamp)) {
+                final long currIndex = pcIndex;
+                final int producer = pIndex(currIndex);
+                final int consumer = cIndex(currIndex);
+                if (pcIndexUpdater.compareAndSet(this, currIndex, computePCIndex(producer, consumer + i))) {
                     break;
                 }
             }
@@ -641,44 +616,16 @@ private static void checkInterrupted() throws InterruptedException {
             }
         }
 
-        /**
-         * The producer thread may produce multiple items before the consumer thread consume the events. If the consumer
-         * thread changes we need to make sure the new consumer thread observes production events so this object
-         * contains the thread to wakeup and a count of how many items are in the queue and not yet consumed.
-         */
-        private static final class ThreadStamp {
-            private static final Short MAX_DEMAND = Short.MAX_VALUE;
-            @Nullable
-            Thread thread;
-            short producerIndex;
-            short consumerIndex;
-
-            ThreadStamp(@Nullable Thread thread) {
-                this.thread = thread;
-            }
-
-            ThreadStamp(@Nullable Thread thread, short producerIndex, short consumerIndex) {
-                this.thread = thread;
-                this.producerIndex = producerIndex;
-                this.consumerIndex = consumerIndex;
-            }
-
-            short incProducerIndex() {
-                return (short) (producerIndex + 1);
-            }
-
-            short incConsumerIndex() {
-                return (short) (consumerIndex + 1);
-            }
+        private static long computePCIndex(int producer, int consumer) {
+            return ((long) producer << 32) | consumer;
+        }
 
-            short incConsumerIndex(int amt) {
-                return (short) (consumerIndex + amt);
-            }
+        private static int cIndex(long pcIndex) {
+            return (int) pcIndex;
+        }
 
-            @Override
-            public String toString() {
-                return "thread: " + thread + " producerIndex: " + producerIndex + " consumerIndex: " + consumerIndex;
-            }
+        private static int pIndex(long pcIndex) {
+            return (int) (pcIndex >>> 32);
         }
 
         private interface BiLongFunction<T, R> {