From f9612a2cf49698311ad20dd94ceda06679b1a6b1 Mon Sep 17 00:00:00 2001
From: Peter Arato <it.arato@gmail.com>
Date: Tue, 23 May 2023 08:34:30 -0400
Subject: [PATCH] Adding IO#timeout.

---
 CHANGELOG.md                            |  1 +
 src/main/ruby/truffleruby/core/io.rb    | 14 ++++++
 src/main/ruby/truffleruby/core/posix.rb | 64 ++++++++++++++++++++++---
 3 files changed, 72 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 865e733142c2..103d51dd82f7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -17,6 +17,7 @@ Compatibility:
 * Fix `Range#size` and return `nil` for beginningless Range when end isn't Numeric (#3039, @rwstauner).
 * Alias `String#-@` to `String#dedup` (#3039, @itarato).
 * Fix `Pathname#relative_path_from` to convert string arguments to Pathname objects (@rwstauner).
+* Adding `IO#timeout` and `IO#timeout=` (#3039, @itarato).
 
 Performance:
 
diff --git a/src/main/ruby/truffleruby/core/io.rb b/src/main/ruby/truffleruby/core/io.rb
index 1fa4efb5c29c..1a5cba73678c 100644
--- a/src/main/ruby/truffleruby/core/io.rb
+++ b/src/main/ruby/truffleruby/core/io.rb
@@ -38,6 +38,8 @@ class IO
 
   include Enumerable
 
+  class TimeoutError < IOError; end
+
   module WaitReadable; end
   module WaitWritable; end
 
@@ -1648,6 +1650,18 @@ def printf(fmt, *args)
     write sprintf(fmt, *args)
   end
 
+  def timeout
+    @timeout || nil
+  end
+
+  def timeout=(new_timeout)
+    if Primitive.nil?(timeout) ^ Primitive.nil?(new_timeout)
+      self.nonblock = !Primitive.nil?(new_timeout)
+    end
+
+    @timeout = new_timeout
+  end
+
   def read(length = nil, buffer = nil)
     ensure_open_and_readable
     buffer = StringValue(buffer) if buffer
diff --git a/src/main/ruby/truffleruby/core/posix.rb b/src/main/ruby/truffleruby/core/posix.rb
index 5de7668ec74d..043fa1b1bdb6 100644
--- a/src/main/ruby/truffleruby/core/posix.rb
+++ b/src/main/ruby/truffleruby/core/posix.rb
@@ -401,10 +401,10 @@ def self.read_string_nonblock(io, count, exception)
   # by IO#sysread
 
   def self.read_string_native(io, length)
-    fd = io.fileno
     buffer = Primitive.io_thread_buffer_allocate(length)
     begin
-      bytes_read = Truffle::POSIX.read(fd, buffer, length)
+      bytes_read = execute_posix_read(io, buffer, length)
+
       if bytes_read < 0
         bytes_read, errno = bytes_read, Errno.errno
       elsif bytes_read == 0 # EOF
@@ -425,11 +425,62 @@ def self.read_string_native(io, length)
     end
   end
 
-  def self.read_to_buffer_native(io, length)
+  def self.execute_posix_read(io, buffer, length)
+    fd = io.fileno
+    return Truffle::POSIX.read(fd, buffer, length) unless io.timeout
+
+    deadline = Time.now.to_f + io.timeout
+
+    loop do
+      current_timeout = deadline - Time.now.to_f
+      raise IO::TimeoutError if current_timeout < 0
+
+      poll_result = Truffle::IOOperations.poll(io, Truffle::IOOperations::POLLIN, current_timeout)
+      if poll_result == 0
+        raise IO::TimeoutError
+      elsif poll_result == -1
+        Errno.handle_errno(Errno.errno)
+      end
+
+      if (bytes_read = Truffle::POSIX.read(fd, buffer, length)) == -1
+        continue if Errno.errno == Errno::EAGAIN
+        break
+      end
+
+      return bytes_read
+    end
+  end
+
+  def self.execute_posix_write(io, buffer, length)
     fd = io.fileno
+    return Truffle::POSIX.write(fd, buffer, length) unless io.timeout
+
+    deadline = Time.now.to_f + io.timeout
+
+    loop do
+      current_timeout = deadline - Time.now.to_f
+      raise IO::TimeoutError if current_timeout < 0
+
+      poll_result = Truffle::IOOperations.poll(io, Truffle::IOOperations::POLLOUT, current_timeout)
+      if poll_result == 0
+        raise IO::TimeoutError
+      elsif poll_result == -1
+        Errno.handle_errno(Errno.errno)
+      end
+
+      if (bytes_written = Truffle::POSIX.write(fd, buffer, length)) == -1
+        continue if Errno.errno == Errno::EAGAIN
+        break
+      end
+
+      return bytes_written
+    end
+  end
+
+  def self.read_to_buffer_native(io, length)
     buffer = Primitive.io_thread_buffer_allocate(length)
     begin
-      bytes_read = Truffle::POSIX.read(fd, buffer, length)
+      bytes_read = execute_posix_read(io, buffer, length)
       if bytes_read < 0
         bytes_read, errno = bytes_read, Errno.errno
       elsif bytes_read == 0 # EOF
@@ -495,7 +546,7 @@ def self.write_string_native(io, string, continue_on_eagain)
 
       written = 0
       while written < length
-        ret = Truffle::POSIX.write(fd, buffer + written, length - written)
+        ret = execute_posix_write(io, buffer + written, length - written)
         if ret < 0
           errno = Errno.errno
           if errno == EAGAIN_ERRNO
@@ -540,12 +591,11 @@ def self.write_string_polyglot(io, string, continue_on_eagain)
   # #write_string_nonblock_polylgot) is called by IO#write_nonblock
 
   def self.write_string_nonblock_native(io, string)
-    fd = io.fileno
     length = string.bytesize
     buffer = Primitive.io_thread_buffer_allocate(length)
     begin
       buffer.write_bytes string
-      written = Truffle::POSIX.write(fd, buffer, length)
+      written = execute_posix_write(io, buffer, length)
 
       if written < 0
         errno = Errno.errno