Add timeout argument to Thread::SizedQueue#push

Author: itarato

CHANGELOG.md

@@ -29,6 +29,7 @@ Compatibility:
 * Add `Module#refinements` (#3039, @itarato).
 * Add `Refinement#refined_class` (#3039, @itarato).
 * Add `rb_hash_new_capa` function (#3039, @itarato).
+* Add `timeout` argument to `Thread::SizedQueue#push` (#3039, @itarato).
 
 Performance:
 

lib/truffle/thread.rb

@@ -50,5 +50,15 @@ def pop(non_block = false, timeout: nil)
     end
     alias_method :shift, :pop
     alias_method :deq, :pop
+
+    def push(value, non_block = false, timeout: nil)
+      Primitive.sized_queue_push(
+        self,
+        value,
+        Primitive.as_boolean(non_block),
+        Truffle::QueueOperations.validate_and_prepare_timeout_in_milliseconds(non_block, timeout))
+    end
+    alias_method :<<, :push
+    alias_method :enq, :push
   end
 end

spec/ruby/core/sizedqueue/append_spec.rb

@@ -1,6 +1,7 @@
 require_relative '../../spec_helper'
 require_relative '../../shared/queue/enque'
 require_relative '../../shared/sizedqueue/enque'
+require_relative '../../shared/types/rb_num2dbl_fails'
 
 describe "SizedQueue#<<" do
   it_behaves_like :queue_enq, :<<, -> { SizedQueue.new(10) }
@@ -9,3 +10,9 @@
 describe "SizedQueue#<<" do
   it_behaves_like :sizedqueue_enq, :<<, -> n { SizedQueue.new(n) }
 end
+
+describe "SizedQueue operations with timeout" do
+  ruby_version_is "3.2" do
+    it_behaves_like :rb_num2dbl_fails, nil, -> v { q = SizedQueue.new(1); q.send(:<<, 1, timeout: v) }
+  end
+end

spec/ruby/core/sizedqueue/enq_spec.rb

@@ -1,6 +1,7 @@
 require_relative '../../spec_helper'
 require_relative '../../shared/queue/enque'
 require_relative '../../shared/sizedqueue/enque'
+require_relative '../../shared/types/rb_num2dbl_fails'
 
 describe "SizedQueue#enq" do
   it_behaves_like :queue_enq, :enq, -> { SizedQueue.new(10) }
@@ -9,3 +10,9 @@
 describe "SizedQueue#enq" do
   it_behaves_like :sizedqueue_enq, :enq, -> n { SizedQueue.new(n) }
 end
+
+describe "SizedQueue operations with timeout" do
+  ruby_version_is "3.2" do
+    it_behaves_like :rb_num2dbl_fails, nil, -> v { q = SizedQueue.new(1); q.enq(1, timeout: v) }
+  end
+end

spec/ruby/core/sizedqueue/push_spec.rb

@@ -1,6 +1,7 @@
 require_relative '../../spec_helper'
 require_relative '../../shared/queue/enque'
 require_relative '../../shared/sizedqueue/enque'
+require_relative '../../shared/types/rb_num2dbl_fails'
 
 describe "SizedQueue#push" do
   it_behaves_like :queue_enq, :push, -> { SizedQueue.new(10) }
@@ -9,3 +10,9 @@
 describe "SizedQueue#push" do
   it_behaves_like :sizedqueue_enq, :push, -> n { SizedQueue.new(n) }
 end
+
+describe "SizedQueue operations with timeout" do
+  ruby_version_is "3.2" do
+    it_behaves_like :rb_num2dbl_fails, nil, -> v { q = SizedQueue.new(1); q.push(1, timeout: v) }
+  end
+end

spec/ruby/shared/sizedqueue/enque.rb

@@ -108,6 +108,25 @@
           "can't set a timeout if non_block is enabled",
         )
       end
+
+      it "raise ClosedQueueError when closed before enqueued" do
+        q = @object.call(1)
+        q.close
+        -> { q.send(@method, 2, timeout: 1) }.should raise_error(ClosedQueueError, "queue closed")
+      end
+
+      it "raise ClosedQueueError when getting closed during wait" do
+        q = @object.call(1)
+        q.push(1)
+
+        t = Thread.new {
+          -> { q.push(1, timeout: 0.1) }.should raise_error(ClosedQueueError, "queue closed")
+        }
+
+        sleep(0.05)
+        q.close
+        t.join
+      end
     end
   end
 end

spec/truffleruby.next-specs

@@ -34,3 +34,6 @@ spec/ruby/core/module/refinements_spec.rb
 spec/ruby/core/refinement/refined_class_spec.rb
 spec/ruby/core/module/used_refinements_spec.rb
 spec/ruby/optional/capi/hash_spec.rb
+spec/ruby/core/sizedqueue/append_spec.rb
+spec/ruby/core/sizedqueue/enq_spec.rb
+spec/ruby/core/sizedqueue/push_spec.rb

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

@@ -127,6 +127,37 @@ public boolean put(Object item) throws InterruptedException {
         }
     }
 
+    @TruffleBoundary
+    public Object put(Object item, long timeoutMilliseconds) throws InterruptedException {
+        lock.lock();
+        try {
+            if (closed) {
+                return CLOSED;
+            }
+
+            final long deadline = System.currentTimeMillis() + timeoutMilliseconds;
+
+            while (size == capacity) {
+                final long currentTimeout = deadline - System.currentTimeMillis();
+                final boolean signalled = ConcurrentOperations.awaitAndCheckInterrupt(canAdd, currentTimeout,
+                        TimeUnit.MILLISECONDS);
+
+                if (!signalled) {
+                    return false;
+                }
+
+                if (closed) {
+                    return CLOSED;
+                }
+            }
+
+            doAdd(item);
+            return true;
+        } finally {
+            lock.unlock();
+        }
+    }
+
     private void doAdd(Object item) {
         items[addEnd] = item;
         addEnd++;
@@ -144,17 +175,20 @@ public Object poll(long timeoutMilliseconds) throws InterruptedException {
         lock.lock();
 
         try {
-            if (size == 0) {
-                final boolean signalled = ConcurrentOperations.awaitAndCheckInterrupt(canTake, timeoutMilliseconds,
+            final long deadline = System.currentTimeMillis() + timeoutMilliseconds;
+
+            while (size == 0) {
+                final long currentTimeout = deadline - System.currentTimeMillis();
+                final boolean signalled = ConcurrentOperations.awaitAndCheckInterrupt(canTake, currentTimeout,
                         TimeUnit.MILLISECONDS);
 
-                if (!signalled) { // timed out
+                if (!signalled) { // Timed out.
                     return null;
                 }
+            }
 
-                if (closed) {
-                    return CLOSED;
-                }
+            if (closed) {
+                return CLOSED;
             }
 
             return doTake();

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

@@ -100,27 +100,29 @@ protected int max(RubySizedQueue self) {
 
     }
 
-    @CoreMethod(names = { "push", "<<", "enq" }, required = 1, optional = 1)
-    public abstract static class SizedQueuePushNode extends CoreMethodArrayArgumentsNode {
+    @Primitive(name = "sized_queue_push")
+    public abstract static class SizedQueuePushNode extends PrimitiveArrayArgumentsNode {
 
         @Specialization
-        protected RubySizedQueue doPush(RubySizedQueue self, final Object value, Object maybeNonBlocking,
+        protected Object doPush(
+                RubySizedQueue self, final Object value, Object maybeNonBlocking, Object timeoutMilliseconds,
                 @Cached BooleanCastWithDefaultNode booleanCastWithDefaultNode,
                 @Cached PushNode pushNode) {
             final boolean nonBlocking = booleanCastWithDefaultNode.execute(this, maybeNonBlocking, false);
-            return pushNode.execute(this, self, value, nonBlocking);
+            return pushNode.execute(this, self, value, nonBlocking, timeoutMilliseconds);
         }
     }
 
     @GenerateInline
     @GenerateCached(false)
     public abstract static class PushNode extends RubyBaseNode {
 
-        public abstract RubySizedQueue execute(Node node, RubySizedQueue self, final Object value, boolean nonBlocking);
+        public abstract Object execute(Node node, RubySizedQueue self, final Object value, boolean nonBlocking,
+                Object timeoutMilliseconds);
 
         @Specialization(guards = "!nonBlocking")
         protected static RubySizedQueue pushBlocking(
-                Node node, RubySizedQueue self, final Object value, boolean nonBlocking,
+                Node node, RubySizedQueue self, final Object value, boolean nonBlocking, Nil timeoutMilliseconds,
                 @Cached @Shared PropagateSharingNode propagateSharingNode) {
             final SizedQueue queue = self.queue;
 
@@ -141,9 +143,38 @@ private static void doPushBlocking(Node node, final Object value, final SizedQue
             });
         }
 
+        @Specialization(guards = "!nonBlocking")
+        protected static Object pushBlocking(
+                Node node, RubySizedQueue self, final Object value, boolean nonBlocking, long timeoutMilliseconds,
+                @Cached @Shared PropagateSharingNode propagateSharingNode) {
+            final SizedQueue queue = self.queue;
+            propagateSharingNode.execute(node, self, value);
+            final long deadline = System.currentTimeMillis() + timeoutMilliseconds;
+
+            var success = getContext(node).getThreadManager().runUntilResult(node, () -> {
+                final long currentTimeout = deadline - System.currentTimeMillis();
+                final Object result;
+
+                if (currentTimeout > 0) {
+                    result = queue.put(value, currentTimeout);
+                } else {
+                    result = queue.put(value, 0);
+                }
+
+                if (result == SizedQueue.CLOSED) {
+                    throw new RaiseException(getContext(node), coreExceptions(node).closedQueueError(node));
+                }
+                return result;
+            });
+            if ((boolean) success) {
+                return self;
+            }
+            return nil;
+        }
+
         @Specialization(guards = "nonBlocking")
         protected static RubySizedQueue pushNonBlock(
-                Node node, RubySizedQueue self, final Object value, boolean nonBlocking,
+                Node node, RubySizedQueue self, final Object value, boolean nonBlocking, Nil timeoutMilliseconds,
                 @Cached @Shared PropagateSharingNode propagateSharingNode,
                 @Cached InlinedBranchProfile errorProfile) {
             final SizedQueue queue = self.queue;

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

@@ -110,17 +110,20 @@ public Object poll(long timeoutMilliseconds) throws InterruptedException {
         lock.lock();
 
         try {
-            if (takeEnd == null) {
-                final boolean signalled = ConcurrentOperations.awaitAndCheckInterrupt(canTake, timeoutMilliseconds,
+            final long deadline = System.currentTimeMillis() + timeoutMilliseconds;
+
+            while (takeEnd == null) {
+                final long currentTimeout = deadline - System.currentTimeMillis();
+                final boolean signalled = ConcurrentOperations.awaitAndCheckInterrupt(canTake, currentTimeout,
                         TimeUnit.MILLISECONDS);
 
-                if (!signalled) { // timed out
+                if (!signalled) { // Timed out.
                     return null;
                 }
+            }
 
-                if (closed) {
-                    return CLOSED;
-                }
+            if (closed) {
+                return CLOSED;
             }
 
             return doTake();