Use a better time wait poll

Author: amarziali

internal-api/internal-api-9/src/main/java/datadog/trace/util/queue/MpscBlockingConsumerArrayQueueVarHandle.java

@@ -98,7 +98,7 @@ public boolean offer(E e, long timeout, @Nonnull TimeUnit unit) throws Interrupt
   }
 
   /**
-   * Polls with a timeout using progressive backoff.
+   * Polls with a timeout.
    *
    * @param timeout max wait time
    * @param unit time unit
@@ -107,32 +107,25 @@ public boolean offer(E e, long timeout, @Nonnull TimeUnit unit) throws Interrupt
    */
   public E poll(long timeout, @Nonnull TimeUnit unit) throws InterruptedException {
     final long deadline = System.nanoTime() + unit.toNanos(timeout);
-    int idleCount = 0;
 
-    while (true) {
-      E value = poll();
-      if (value != null) {
-        return value;
-      }
+    E e = poll();
+    if (e != null) {
+      return e;
+    }
 
-      long remaining = deadline - System.nanoTime();
-      if (remaining <= 0) {
-        return null;
-      }
+    // register this thread as the waiting consumer
+    consumerThread = Thread.currentThread();
+    final long remaining = deadline - System.nanoTime();
 
-      // Progressive backoff
-      if (idleCount < 100) {
-        // spin
-      } else if (idleCount < 1_000) {
-        Thread.yield();
-      } else {
-        LockSupport.parkNanos(Math.min(remaining, 1_000_000L));
-      }
-      idleCount++;
+    if (remaining <= 0) {
+      consumerThread = null;
+      return null;
+    }
 
-      if (Thread.interrupted()) {
-        throw new InterruptedException();
-      }
+    LockSupport.parkNanos(this, remaining);
+    if (Thread.interrupted()) {
+      throw new InterruptedException();
     }
+    return poll();
   }
 }

internal-api/internal-api-9/src/test/groovy/datadog/trace/util/queue/AbstractQueueTest.groovy

@@ -1,13 +1,10 @@
 package datadog.trace.util.queue
 
-import static java.util.concurrent.TimeUnit.NANOSECONDS
 
 import datadog.trace.test.util.DDSpecification
-import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.AtomicBoolean
 import java.util.function.Consumer
 
-abstract class AbstractQueueTest<T extends AbstractQueue<Integer>> extends DDSpecification {
+abstract class AbstractQueueTest<T extends BaseQueue<Integer>> extends DDSpecification {
   abstract T createQueue(int capacity)
   protected T queue = createQueue(8)
 
@@ -103,7 +100,7 @@ abstract class AbstractQueueTest<T extends AbstractQueue<Integer>> extends DDSpe
 
   def "remainingCapacity should reflect current occupancy"() {
     given:
-    def q = new MpscArrayQueue<Integer>(4)
+    def q = createQueue(4)
     q.offer(1)
     q.offer(2)
 
@@ -116,48 +113,4 @@ abstract class AbstractQueueTest<T extends AbstractQueue<Integer>> extends DDSpe
     then:
     q.remainingCapacity() == 3
   }
-
-
-  def "poll with timeout returns null if no element becomes available"() {
-    when:
-    def start = System.nanoTime()
-    def value = queue.poll(200, TimeUnit.MILLISECONDS)
-    def elapsedMs = NANOSECONDS.toMillis(System.nanoTime() - start)
-
-    then:
-    value == null
-    elapsedMs >= 200  // waited approximately the timeout
-  }
-
-  def "poll with zero timeout behaves like immediate poll"() {
-    expect:
-    queue.poll(0, TimeUnit.MILLISECONDS) == null
-
-    when:
-    queue.offer(99)
-
-    then:
-    queue.poll(0, TimeUnit.MILLISECONDS) == 99
-  }
-
-  def "poll throws InterruptedException if interrupted"() {
-    given:
-    def thrown = new AtomicBoolean()
-    def thread = Thread.start {
-      try {
-        queue.poll(500, TimeUnit.MILLISECONDS)
-      } catch (InterruptedException ie) {
-        thrown.set(true)
-        Thread.currentThread().interrupt()
-      }
-    }
-
-    when:
-    Thread.sleep(50)
-    thread.interrupt()
-    thread.join()
-
-    then:
-    thrown.get()
-  }
 }

internal-api/internal-api-9/src/test/groovy/datadog/trace/util/queue/MpscBlockingConsumerArrayQueueVarHandleTest.groovy

@@ -1,7 +1,10 @@
 package datadog.trace.util.queue
 
+import static java.util.concurrent.TimeUnit.NANOSECONDS
 
 import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicReference
 import java.util.function.Consumer
 import java.util.function.Supplier
@@ -121,4 +124,47 @@ class MpscBlockingConsumerArrayQueueVarHandleTest extends AbstractQueueTest<Mpsc
     filled == 8
     queue.size() == 8
   }
+
+  def "poll with timeout returns null if no element becomes available"() {
+    when:
+    def start = System.nanoTime()
+    def value = queue.poll(200, TimeUnit.MILLISECONDS)
+    def elapsedMs = NANOSECONDS.toMillis(System.nanoTime() - start)
+
+    then:
+    value == null
+    elapsedMs >= 200  // waited approximately the timeout
+  }
+
+  def "poll with zero timeout behaves like immediate poll"() {
+    expect:
+    queue.poll(0, TimeUnit.MILLISECONDS) == null
+
+    when:
+    queue.offer(99)
+
+    then:
+    queue.poll(0, TimeUnit.MILLISECONDS) == 99
+  }
+
+  def "poll throws InterruptedException if interrupted"() {
+    given:
+    def thrown = new AtomicBoolean()
+    def thread = Thread.start {
+      try {
+        queue.poll(500, TimeUnit.MILLISECONDS)
+      } catch (InterruptedException ie) {
+        thrown.set(true)
+        Thread.currentThread().interrupt()
+      }
+    }
+
+    when:
+    Thread.sleep(50)
+    thread.interrupt()
+    thread.join()
+
+    then:
+    thrown.get()
+  }
 }