Move queue pop to ruby and add timeout keyword argument.

Author: itarato

CHANGELOG.md

@@ -20,6 +20,7 @@ Compatibility:
 * Add `String#bytesplice` (#3039, @itarato).
 * 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).
 
 Performance:
 

lib/truffle/thread.rb

@@ -29,4 +29,15 @@ class ThreadError < StandardError
 class Thread
   Queue = ::Queue
   SizedQueue = ::SizedQueue
+
+  class Queue
+    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))
+    end
+    alias_method :shift, :pop
+    alias_method :deq, :pop
+  end
 end

spec/ruby/core/queue/deq_spec.rb

@@ -1,6 +1,13 @@
 require_relative '../../spec_helper'
 require_relative '../../shared/queue/deque'
+require_relative '../../shared/types/rb_num2dbl_fails'
 
 describe "Queue#deq" do
   it_behaves_like :queue_deq, :deq, -> { Queue.new }
 end
+
+describe "Queue operations with timeout" do
+  ruby_version_is "3.2" do
+    it_behaves_like :rb_num2dbl_fails, nil, -> v { q = Queue.new; q.push(1); q.deq(timeout: v) }
+  end
+end

spec/ruby/core/queue/pop_spec.rb

@@ -1,6 +1,13 @@
 require_relative '../../spec_helper'
 require_relative '../../shared/queue/deque'
+require_relative '../../shared/types/rb_num2dbl_fails'
 
 describe "Queue#pop" do
   it_behaves_like :queue_deq, :pop, -> { Queue.new }
 end
+
+describe "Queue operations with timeout" do
+  ruby_version_is "3.2" do
+    it_behaves_like :rb_num2dbl_fails, nil, -> v { q = Queue.new; q.push(1); q.pop(timeout: v) }
+  end
+end

spec/ruby/core/queue/shift_spec.rb

@@ -1,6 +1,13 @@
 require_relative '../../spec_helper'
 require_relative '../../shared/queue/deque'
+require_relative '../../shared/types/rb_num2dbl_fails'
 
 describe "Queue#shift" do
   it_behaves_like :queue_deq, :shift, -> { Queue.new }
 end
+
+describe "Queue operations with timeout" do
+  ruby_version_is "3.2" do
+    it_behaves_like :rb_num2dbl_fails, nil, -> v { q = Queue.new; q.push(1); q.shift(timeout: v) }
+  end
+end

spec/ruby/shared/queue/deque.rb

@@ -37,6 +37,15 @@
     q.send(@method).should == 1
   end
 
+  it "converts false-ish for non_blocking to boolean" do
+    q = @object.call
+    q << 1
+    q << 2
+
+    q.send(@method, false).should == 1
+    q.send(@method, nil).should == 2
+  end
+
   it "returns nil for a closed empty queue" do
     q = @object.call
     q.close
@@ -143,5 +152,13 @@
       q.close
       -> { q.send(@method, true) }.should raise_error(ThreadError)
     end
+
+    it "converts true-ish non_blocking argument to true" do
+      q = @object.call
+
+      -> { q.send(@method, true) }.should raise_error(ThreadError)
+      -> { q.send(@method, 1) }.should raise_error(ThreadError)
+      -> { q.send(@method, "") }.should raise_error(ThreadError)
+    end
   end
 end

spec/ruby/shared/types/rb_num2dbl_fails.rb

@@ -0,0 +1,17 @@
+#
+# Shared tests for rb_num2dbl related conversion failures.
+#
+# Usage example:
+#   it_behaves_like :rb_num2dbl_fails, nil, -> v { o = A.new; o.foo(v) }
+#
+
+describe :rb_num2dbl_fails, shared: true do
+  it "fails if string is provided" do
+    -> { @object.call("123") }.should raise_error(TypeError, "no implicit conversion to float from string")
+  end
+
+  it "fails if boolean is provided" do
+    -> { @object.call(true) }.should raise_error(TypeError, "no implicit conversion to float from true")
+    -> { @object.call(false) }.should raise_error(TypeError, "no implicit conversion to float from false")
+  end
+end

spec/truffleruby.next-specs

@@ -21,3 +21,7 @@ spec/ruby/core/string/bytesplice_spec.rb
 
 spec/ruby/core/string/byteindex_spec.rb
 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

src/main/java/org/truffleruby/core/CoreLibrary.java

@@ -1053,6 +1053,7 @@ public boolean isTruffleBootMainMethod(SharedMethodInfo info) {
             "/core/posix.rb",
             "/core/main.rb",
             "/core/post.rb",
+            "/core/truffle/queue_operations.rb",
             POST_BOOT_FILE
     };
 

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

@@ -13,19 +13,17 @@
 import com.oracle.truffle.api.library.CachedLibrary;
 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.NonStandard;
+import org.truffleruby.builtins.PrimitiveArrayArgumentsNode;
 import org.truffleruby.core.array.ArrayGuards;
 import org.truffleruby.core.array.RubyArray;
 import org.truffleruby.core.array.library.ArrayStoreLibrary;
-import org.truffleruby.core.cast.BooleanCastWithDefaultNode;
 import org.truffleruby.core.cast.ToANode;
 import org.truffleruby.core.klass.RubyClass;
+import org.truffleruby.language.Nil;
 import org.truffleruby.language.NotProvided;
-import org.truffleruby.language.RubyBaseNodeWithExecute;
-import org.truffleruby.language.RubyNode;
 import org.truffleruby.annotations.Visibility;
 import org.truffleruby.language.control.RaiseException;
 import org.truffleruby.language.objects.AllocationTracing;
@@ -34,8 +32,6 @@
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.Cached.Exclusive;
-import com.oracle.truffle.api.dsl.CreateCast;
-import com.oracle.truffle.api.dsl.NodeChild;
 import com.oracle.truffle.api.dsl.Specialization;
 
 @CoreModule(value = "Queue", isClass = true)
@@ -73,18 +69,11 @@ protected RubyQueue push(RubyQueue self, final Object value) {
 
     }
 
-    @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 = "queue_pop")
+    public abstract static class PopNode extends PrimitiveArrayArgumentsNode {
 
         @Specialization(guards = "!nonBlocking")
-        protected Object popBlocking(RubyQueue self, boolean nonBlocking,
+        protected Object popBlocking(RubyQueue self, boolean nonBlocking, Nil timeoutMilliseconds,
                 @Exclusive @Cached InlinedBranchProfile closedProfile) {
             final UnsizedQueue queue = self.queue;
 
@@ -103,61 +92,48 @@ private Object doPop(UnsizedQueue queue) {
             return getContext().getThreadManager().runUntilResult(this, queue::take);
         }
 
-        @Specialization(guards = "nonBlocking")
-        protected Object popNonBlock(RubyQueue self, boolean nonBlocking,
-                @Exclusive @Cached InlinedBranchProfile errorProfile) {
+        @Specialization(guards = "!nonBlocking")
+        protected Object popBlocking(RubyQueue self, boolean nonBlocking, long timeoutMilliseconds,
+                @Exclusive @Cached InlinedBranchProfile closedProfile) {
             final UnsizedQueue queue = self.queue;
+            final long deadline = System.currentTimeMillis() + timeoutMilliseconds;
 
-            final Object value = queue.poll();
-
-            if (value == null) {
-                errorProfile.enter(this);
-                throw new RaiseException(getContext(), coreExceptions().threadError("queue empty", this));
-            } else {
-                return value;
-            }
-        }
+            var result = getContext().getThreadManager().runUntilResult(this, () -> {
+                final long currentTimeout = deadline - System.currentTimeMillis();
+                final Object value;
 
-    }
+                if (currentTimeout > 0) {
+                    value = queue.poll(currentTimeout);
+                } else {
+                    value = queue.poll();
+                }
 
-    @NonStandard
-    @CoreMethod(names = "receive_timeout", required = 1, visibility = Visibility.PRIVATE, lowerFixnum = 1)
-    public abstract static class ReceiveTimeoutNode extends CoreMethodArrayArgumentsNode {
+                if (value == UnsizedQueue.CLOSED) {
+                    closedProfile.enter(this);
+                    return nil;
+                } else {
+                    return value;
+                }
+            });
 
-        @Specialization
-        protected Object receiveTimeout(RubyQueue self, int duration) {
-            return receiveTimeout(self, (double) duration);
+            if (result == null) {
+                return nil;
+            }
+            return result;
         }
 
-        @TruffleBoundary
-        @Specialization
-        protected Object receiveTimeout(RubyQueue self, double duration) {
+        @Specialization(guards = "nonBlocking")
+        protected Object popNonBlock(RubyQueue self, boolean nonBlocking, Nil timeoutMilliseconds,
+                @Exclusive @Cached InlinedBranchProfile errorProfile) {
             final UnsizedQueue queue = self.queue;
 
-            final long durationInMillis = (long) (duration * 1000.0);
-            final long start = System.currentTimeMillis();
-
-            return getContext().getThreadManager().runUntilResult(this, () -> {
-                long now = System.currentTimeMillis();
-                long waited = now - start;
-                if (waited >= durationInMillis) {
-                    // Try again to make sure we at least tried once
-                    final Object result = queue.poll();
-                    return translateResult(result);
-                }
-
-                final Object result = queue.poll(durationInMillis);
-                return translateResult(result);
-            });
-        }
+            final Object value = queue.poll();
 
-        private Object translateResult(Object result) {
-            if (result == null) {
-                return false;
-            } else if (result == UnsizedQueue.CLOSED) {
-                return nil;
+            if (value == null) {
+                errorProfile.enter(this);
+                throw new RaiseException(getContext(), coreExceptions().threadError("queue empty", this));
             } else {
-                return result;
+                return value;
             }
         }