[GR-19220] Add timeout argument to Thread::SizedQueue#pop (#3046)

eregon · eregon · commit 9a349526d2fb · 2023-06-21T19:29:49.000Z
PullRequest: truffleruby/3877
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -23,6 +23,7 @@ Compatibility:
 * Add `String#byteindex` and `String#byterindex` (#3039, @itarato).
 * Add implementations of `rb_proc_call_with_block`, `rb_proc_call_kw`, `rb_proc_call_with_block_kw` and `rb_funcall_with_block_kw` (#3068, @andrykonchin).
 * Add optional `timeout` argument to `Thread::Queue#pop` (#3039, @itarato).
+* Add optional `timeout` argument to `Thread::SizedsQueue#pop` (#3039, @itarato).
 
 Performance:
 
diff --git a/lib/truffle/thread.rb b/lib/truffle/thread.rb
@@ -35,7 +35,18 @@ def pop(non_block = false, timeout: nil)
       Primitive.queue_pop(
         self,
         Primitive.as_boolean(non_block),
-        Truffle::QueueOperations.validate_and_prepare_timeout(non_block, timeout))
+        Truffle::QueueOperations.validate_and_prepare_timeout_in_milliseconds(non_block, timeout))
+    end
+    alias_method :shift, :pop
+    alias_method :deq, :pop
+  end
+
+  class SizedQueue
+    def pop(non_block = false, timeout: nil)
+      Primitive.sized_queue_pop(
+        self,
+        Primitive.as_boolean(non_block),
+        Truffle::QueueOperations.validate_and_prepare_timeout_in_milliseconds(non_block, timeout))
     end
     alias_method :shift, :pop
     alias_method :deq, :pop
diff --git a/spec/ruby/core/sizedqueue/deq_spec.rb b/spec/ruby/core/sizedqueue/deq_spec.rb
@@ -1,6 +1,13 @@
 require_relative '../../spec_helper'
 require_relative '../../shared/queue/deque'
+require_relative '../../shared/types/rb_num2dbl_fails'
 
 describe "SizedQueue#deq" do
   it_behaves_like :queue_deq, :deq, -> { SizedQueue.new(10) }
 end
+
+describe "SizedQueue operations with timeout" do
+  ruby_version_is "3.2" do
+    it_behaves_like :rb_num2dbl_fails, nil, -> v { q = SizedQueue.new(10); q.push(1); q.deq(timeout: v) }
+  end
+end
diff --git a/spec/ruby/core/sizedqueue/pop_spec.rb b/spec/ruby/core/sizedqueue/pop_spec.rb
@@ -1,6 +1,13 @@
 require_relative '../../spec_helper'
 require_relative '../../shared/queue/deque'
+require_relative '../../shared/types/rb_num2dbl_fails'
 
 describe "SizedQueue#pop" do
   it_behaves_like :queue_deq, :pop, -> { SizedQueue.new(10) }
 end
+
+describe "SizedQueue operations with timeout" do
+  ruby_version_is "3.2" do
+    it_behaves_like :rb_num2dbl_fails, nil, -> v { q = SizedQueue.new(10); q.push(1); q.pop(timeout: v) }
+  end
+end
diff --git a/spec/ruby/core/sizedqueue/shift_spec.rb b/spec/ruby/core/sizedqueue/shift_spec.rb
@@ -1,6 +1,13 @@
 require_relative '../../spec_helper'
 require_relative '../../shared/queue/deque'
+require_relative '../../shared/types/rb_num2dbl_fails'
 
 describe "SizedQueue#shift" do
   it_behaves_like :queue_deq, :shift, -> { SizedQueue.new(10) }
 end
+
+describe "SizedQueue operations with timeout" do
+  ruby_version_is "3.2" do
+    it_behaves_like :rb_num2dbl_fails, nil, -> v { q = SizedQueue.new(10); q.push(1); q.shift(timeout: v) }
+  end
+end
diff --git a/spec/ruby/shared/queue/deque.rb b/spec/ruby/shared/queue/deque.rb
@@ -123,6 +123,12 @@
           "can't set a timeout if non_block is enabled",
         )
       end
+
+      it "returns nil for a closed empty queue" do
+        q = @object.call
+        q.close
+        q.send(@method, timeout: 0).should == nil
+      end
     end
   end
 
diff --git a/spec/truffleruby.next-specs b/spec/truffleruby.next-specs
@@ -25,3 +25,7 @@ spec/ruby/core/string/byterindex_spec.rb
 spec/ruby/core/queue/deq_spec.rb
 spec/ruby/core/queue/pop_spec.rb
 spec/ruby/core/queue/shift_spec.rb
+
+spec/ruby/core/sizedqueue/deq_spec.rb
+spec/ruby/core/sizedqueue/pop_spec.rb
+spec/ruby/core/sizedqueue/shift_spec.rb
diff --git a/src/main/.checkstyle_checks.xml b/src/main/.checkstyle_checks.xml
@@ -149,6 +149,10 @@
       <property name="format" value="\.computeIfAbsent"/>
       <property name="message" value="Use ConcurrentOperations.getOrCompute() instead of ConcurrentHashMap.computeIfAbsent() which does not scale"/>
     </module>
+    <module name="RegexpSinglelineJava">
+      <property name="format" value="\b(?!\w*[Ll]atch)\w+\.await\(.*\)"/>
+      <property name="message" value="Use ConcurrentOperations.awaitAndCheckInterrupt() instead"/>
+    </module>
     <module name="RegexpSinglelineJava">
       <property name="format" value="@NodeChildren"/>
       <property name="message" value="Do not use @NodeChildren, use multiple @Child annotations directly"/>
diff --git a/src/main/java/org/truffleruby/collections/ConcurrentOperations.java b/src/main/java/org/truffleruby/collections/ConcurrentOperations.java
@@ -13,6 +13,10 @@
 
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.AbstractQueuedSynchronizer;
+import java.util.concurrent.locks.Condition;
+import java.util.concurrent.locks.ReentrantLock;
 import java.util.function.Function;
 
 public abstract class ConcurrentOperations {
@@ -43,4 +47,40 @@ public static <K, V> boolean replace(Map<K, V> map, K key, V oldValue, V newValu
         return oldValue == null ? map.putIfAbsent(key, newValue) == null : map.replace(key, oldValue, newValue);
     }
 
+    /** {@link Condition#await()} might return and not throw InterruptedException if there was both a signal and an
+     * interrupt before the lock could be re-acquired by await(). We always want the InterruptedException if both a
+     * signal and an interrupt happen.
+     * <hr>
+     * {@link Condition#await()} when it sees both a signal and a {@link Thread#interrupt()} has to decide whether it
+     * prefers the signal or the interrupt. For a {@link ReentrantLock#newCondition()} condition, the implementation in
+     * {@link AbstractQueuedSynchronizer.ConditionObject#await()} sometimes throws InterruptedException and sometimes
+     * returns, depending on which happens first. Because await() needs to reacquire the lock, even if the
+     * InterruptedException happens before it might still become a situation of "both signal and interrupt". If it
+     * prefers the signal and just returns, it does {@code Thread.currentThread().interrupt()} to let us know there was
+     * an interrupt too. In any case, if there was any interrupt we want to throw InterruptedException, regardless of
+     * what the implementation prefers. */
+    public static void awaitAndCheckInterrupt(Condition condition) throws InterruptedException {
+        // Checkstyle: stop
+        condition.await();
+        // Checkstyle: resume
+
+        if (Thread.interrupted()) {
+            throw new InterruptedException();
+        }
+    }
+
+    /** See {@link ConcurrentOperations#awaitAndCheckInterrupt(java.util.concurrent.locks.Condition)} */
+    public static boolean awaitAndCheckInterrupt(Condition condition, long time, TimeUnit unit)
+            throws InterruptedException {
+        // Checkstyle: stop
+        boolean value = condition.await(time, unit);
+        // Checkstyle: resume
+
+        if (Thread.interrupted()) {
+            throw new InterruptedException();
+        }
+
+        return value;
+    }
+
 }
diff --git a/src/main/java/org/truffleruby/core/mutex/ConditionVariableNodes.java b/src/main/java/org/truffleruby/core/mutex/ConditionVariableNodes.java
@@ -21,6 +21,7 @@
 import org.truffleruby.annotations.CoreModule;
 import org.truffleruby.annotations.Primitive;
 import org.truffleruby.builtins.PrimitiveArrayArgumentsNode;
+import org.truffleruby.collections.ConcurrentOperations;
 import org.truffleruby.collections.Memo;
 import org.truffleruby.core.klass.RubyClass;
 import org.truffleruby.core.thread.RubyThread;
@@ -98,7 +99,7 @@ private void waitInternal(RubyConditionVariable conditionVariable, ReentrantLock
 
             // condLock must be locked before unlocking mutexLock, to avoid losing potential signals.
             // We must not change the Ruby Thread status and not consume a Java thread interrupt while locking condLock.
-            // If there is an interrupt, it should be consumed by condition.await() and the Ruby Thread sleep status
+            // If there is an interrupt, it should be consumed by Condition#await() and the Ruby Thread sleep status
             // must imply being ready to be interrupted by Thread#{run,wakeup}.
             condLock.lock();
             int holdCount = 0; // can be > 1 for MonitorMixin
@@ -164,9 +165,10 @@ private void awaitSignal(RubyConditionVariable self, RubyThread thread, long dur
                          * ReentrantLock exposed to Ruby, e.g., via a Ruby Mutex, as that could be held forever by
                          * another Ruby thread (which did {code mutex.lock} after the #wait), and we would never be able
                          * to interrupt both threads at the same time for a synchronous ThreadLocalAction). */
-                        condition.await(endNanoTime - currentTime, TimeUnit.NANOSECONDS);
+                        ConcurrentOperations.awaitAndCheckInterrupt(condition, endNanoTime - currentTime,
+                                TimeUnit.NANOSECONDS);
                     } else {
-                        condition.await();
+                        ConcurrentOperations.awaitAndCheckInterrupt(condition);
                     }
                     if (consumeSignal(self)) {
                         return BlockingAction.SUCCESS;
diff --git a/src/main/java/org/truffleruby/core/queue/SizedQueue.java b/src/main/java/org/truffleruby/core/queue/SizedQueue.java
@@ -11,10 +11,12 @@
 
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.ReentrantLock;
 
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
+import org.truffleruby.collections.ConcurrentOperations;
 
 public class SizedQueue {
 
@@ -111,7 +113,7 @@ public boolean put(Object item) throws InterruptedException {
                     return false;
                 }
 
-                canAdd.await();
+                ConcurrentOperations.awaitAndCheckInterrupt(canAdd);
             }
 
             if (closed) {
@@ -137,6 +139,30 @@ private void doAdd(Object item) {
         }
     }
 
+    @TruffleBoundary
+    public Object poll(long timeoutMilliseconds) throws InterruptedException {
+        lock.lock();
+
+        try {
+            if (size == 0) {
+                final boolean signalled = ConcurrentOperations.awaitAndCheckInterrupt(canTake, timeoutMilliseconds,
+                        TimeUnit.MILLISECONDS);
+
+                if (!signalled) { // timed out
+                    return null;
+                }
+
+                if (closed) {
+                    return CLOSED;
+                }
+            }
+
+            return doTake();
+        } finally {
+            lock.unlock();
+        }
+    }
+
     @TruffleBoundary
     public Object poll() {
         lock.lock();
@@ -162,7 +188,7 @@ public Object take() throws InterruptedException {
                     return CLOSED;
                 }
 
-                canTake.await();
+                ConcurrentOperations.awaitAndCheckInterrupt(canTake);
             }
 
             return doTake();
diff --git a/src/main/java/org/truffleruby/core/queue/SizedQueueNodes.java b/src/main/java/org/truffleruby/core/queue/SizedQueueNodes.java
@@ -16,12 +16,15 @@
 import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.profiles.InlinedBranchProfile;
 import org.truffleruby.annotations.CoreMethod;
+import org.truffleruby.annotations.Primitive;
 import org.truffleruby.builtins.CoreMethodArrayArgumentsNode;
 import org.truffleruby.builtins.CoreMethodNode;
 import org.truffleruby.annotations.CoreModule;
+import org.truffleruby.builtins.PrimitiveArrayArgumentsNode;
 import org.truffleruby.core.cast.BooleanCastWithDefaultNode;
 import org.truffleruby.core.klass.RubyClass;
 import org.truffleruby.core.thread.ThreadManager.BlockingAction;
+import org.truffleruby.language.Nil;
 import org.truffleruby.language.RubyBaseNodeWithExecute;
 import org.truffleruby.language.RubyNode;
 import org.truffleruby.annotations.Visibility;
@@ -154,18 +157,11 @@ protected RubySizedQueue pushNonBlock(RubySizedQueue self, final Object value, b
 
     }
 
-    @CoreMethod(names = { "pop", "shift", "deq" }, optional = 1)
-    @NodeChild(value = "queue", type = RubyNode.class)
-    @NodeChild(value = "nonBlocking", type = RubyBaseNodeWithExecute.class)
-    public abstract static class PopNode extends CoreMethodNode {
-
-        @CreateCast("nonBlocking")
-        protected RubyBaseNodeWithExecute coerceToBoolean(RubyBaseNodeWithExecute nonBlocking) {
-            return BooleanCastWithDefaultNode.create(false, nonBlocking);
-        }
+    @Primitive(name = "sized_queue_pop")
+    public abstract static class PopNode extends PrimitiveArrayArgumentsNode {
 
         @Specialization(guards = "!nonBlocking")
-        protected Object popBlocking(RubySizedQueue self, boolean nonBlocking) {
+        protected Object popBlocking(RubySizedQueue self, boolean nonBlocking, Nil timeoutMilliseconds) {
             final SizedQueue queue = self.queue;
 
             final Object value = doPop(queue);
@@ -182,8 +178,31 @@ private Object doPop(SizedQueue queue) {
             return getContext().getThreadManager().runUntilResult(this, queue::take);
         }
 
+        @Specialization(guards = "!nonBlocking")
+        protected Object popBlocking(RubySizedQueue self, boolean nonBlocking, long timeoutMilliseconds) {
+            final SizedQueue queue = self.queue;
+            final long deadline = System.currentTimeMillis() + timeoutMilliseconds;
+
+            return getContext().getThreadManager().runUntilResult(this, () -> {
+                final long currentTimeout = deadline - System.currentTimeMillis();
+                final Object value;
+
+                if (currentTimeout > 0) {
+                    value = queue.poll(currentTimeout);
+                } else {
+                    value = queue.poll();
+                }
+
+                if (value == SizedQueue.CLOSED || value == null) {
+                    return nil;
+                } else {
+                    return value;
+                }
+            });
+        }
+
         @Specialization(guards = "nonBlocking")
-        protected Object popNonBlock(RubySizedQueue self, boolean nonBlocking,
+        protected Object popNonBlock(RubySizedQueue self, boolean nonBlocking, Nil timeoutMilliseconds,
                 @Cached InlinedBranchProfile errorProfile) {
             final SizedQueue queue = self.queue;
 
diff --git a/src/main/java/org/truffleruby/core/queue/UnsizedQueue.java b/src/main/java/org/truffleruby/core/queue/UnsizedQueue.java
@@ -17,6 +17,7 @@
 import java.util.concurrent.locks.ReentrantLock;
 
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
+import org.truffleruby.collections.ConcurrentOperations;
 
 public class UnsizedQueue {
 
@@ -95,17 +96,7 @@ public Object take() throws InterruptedException {
                     return CLOSED;
                 }
 
-                canTake.await();
-
-                /* We need to check the interrupted flag here, while holding the lock and after the await(). Otherwise,
-                 * if another thread does {@code th.raise(exc); q.push(1)} like in core/thread/handle_interrupt_spec.rb
-                 * then we might miss the ThreadLocalAction and instead return, which would incorrectly delay the
-                 * ThreadLocalAction if under {@code Thread.handle_interrupt(Object => :on_blocking)}. The other thread
-                 * cannot wait on the Future, as that might block arbitrary long when side-effects are disabled on this
-                 * thread. */
-                if (Thread.interrupted()) {
-                    throw new InterruptedException();
-                }
+                ConcurrentOperations.awaitAndCheckInterrupt(canTake);
             }
 
             return doTake();
@@ -120,13 +111,10 @@ public Object poll(long timeoutMilliseconds) throws InterruptedException {
 
         try {
             if (takeEnd == null) {
-                final boolean signalled = canTake.await(timeoutMilliseconds, TimeUnit.MILLISECONDS);
-
-                if (Thread.interrupted()) { // See comment of take()
-                    throw new InterruptedException();
-                }
+                final boolean signalled = ConcurrentOperations.awaitAndCheckInterrupt(canTake, timeoutMilliseconds,
+                        TimeUnit.MILLISECONDS);
 
-                if (!signalled) {
+                if (!signalled) { // timed out
                     return null;
                 }
 
diff --git a/src/main/java/org/truffleruby/core/thread/ThreadManager.java b/src/main/java/org/truffleruby/core/thread/ThreadManager.java
@@ -568,7 +568,7 @@ public Object apply(State state) {
      * given action should throw an {@link InterruptedException} when {@link Thread#interrupt()} is called. Otherwise,
      * the {@link SafepointManager} will not be able to interrupt this action. See
      * {@link ThreadNodes.ThreadRunBlockingSystemCallNode} for blocking native calls. If the action throws an
-     * {@link InterruptedException}, it will be retried until it returns a non-null value.
+     * {@link InterruptedException}, it will be retried.
      *
      * @param action must not touch any Ruby state
      * @return the return value from {@code action} */
diff --git a/src/main/ruby/truffleruby/core/truffle/queue_operations.rb b/src/main/ruby/truffleruby/core/truffle/queue_operations.rb
@@ -10,17 +10,11 @@
 
 module Truffle
   module QueueOperations
-    def self.validate_and_prepare_timeout(non_block, timeout)
-      if non_block
-        raise ArgumentError, "can't set a timeout if non_block is enabled" unless Primitive.nil?(timeout)
-      else
-        unless Primitive.nil?(timeout)
-          # Convert seconds to milliseconds.
-          return (Truffle::Type.rb_num2dbl(timeout) * 1000).to_i
-        end
-      end
+    def self.validate_and_prepare_timeout_in_milliseconds(non_block, timeout_seconds_or_nil)
+      return nil if Primitive.nil?(timeout_seconds_or_nil)
+      raise ArgumentError, "can't set a timeout if non_block is enabled" if non_block
 
-      timeout
+      (Truffle::Type.rb_num2dbl(timeout_seconds_or_nil) * 1000).to_i
     end
   end
 end
