Add timeout argument to Thread::SizedQueue#pop.

Author: itarato

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:
 

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

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

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

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

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
 

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

src/main/java/org/truffleruby/core/queue/SizedQueue.java

@@ -11,6 +11,7 @@
 
 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;
 
@@ -137,6 +138,40 @@ private void doAdd(Object item) {
         }
     }
 
+    @TruffleBoundary
+    public Object poll(long timeoutMilliseconds) throws InterruptedException {
+        lock.lock();
+
+        try {
+            if (size == 0) {
+                final boolean signalled = canTake.await(timeoutMilliseconds, TimeUnit.MILLISECONDS);
+
+                /* 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();
+                }
+
+                if (!signalled) {
+                    // Timed out.
+                    return null;
+                }
+
+                if (closed) {
+                    return CLOSED;
+                }
+            }
+
+            return doTake();
+        } finally {
+            lock.unlock();
+        }
+    }
+
     @TruffleBoundary
     public Object poll() {
         lock.lock();

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;
 

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 timeout_seconds_or_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