Fix nonblock Socket APIs accepting exception: false

Author: eregon

lib/truffle/socket/basic_socket.rb

@@ -155,15 +155,37 @@ def send(message, flags, dest_sockaddr = nil)
     bytes_sent
   end
 
-  def recv(bytes_to_read, flags = 0)
+  private def internal_recv(bytes_to_read, flags, buffer, exception)
+    raise ArgumentError, 'buffer argument not yet supported' if buffer
+
     Truffle::Socket::Foreign.memory_pointer(bytes_to_read) do |buf|
       n_bytes = Truffle::Socket::Foreign.recv(@descriptor, buf, bytes_to_read, flags)
-      Errno.handle('recv(2)') if n_bytes == -1
+
+      if n_bytes == -1
+        if exception
+          Truffle::Socket::Error.read_error('recv(2)', self)
+        else
+          return :wait_readable
+        end
+      end
+
       return buf.read_string(n_bytes)
     end
   end
 
-  def recvmsg(max_msg_len = nil, flags = 0, *_)
+  def recv(bytes_to_read, flags = 0, buf = nil)
+    internal_recv(bytes_to_read, flags, buf, true)
+  end
+
+  def recv_nonblock(bytes_to_read, flags = 0, buf = nil, exception: true)
+    fcntl(Fcntl::F_SETFL, Fcntl::O_NONBLOCK)
+
+    Truffle::Socket::Error.wrap_read_nonblock do
+      internal_recv(bytes_to_read, flags, buf, exception)
+    end
+  end
+
+  private def internal_recvmsg(max_msg_len, flags, max_control_len, scm_rights, exception)
     socket_type = getsockopt(:SOCKET, :TYPE).int
 
     if socket_type == Socket::SOCK_STREAM
@@ -188,7 +210,13 @@ def recvmsg(max_msg_len = nil, flags = 0, *_)
         msg_size = Truffle::Socket::Foreign
           .recvmsg(@descriptor, header.pointer, flags)
 
-        Truffle::Socket::Error.read_error('recvmsg(2)', self) if msg_size < 0
+        if msg_size < 0
+          if exception
+            Truffle::Socket::Error.read_error('recvmsg(2)', self)
+          else
+            return :wait_readable
+          end
+        end
 
         if grow_msg and header.message_truncated?
           need_more = true
@@ -216,13 +244,17 @@ def recvmsg(max_msg_len = nil, flags = 0, *_)
     nil
   end
 
-  def recvmsg_nonblock(max_msg_len = nil, flags = 0, *_)
+  def recvmsg(max_msg_len = nil, flags = 0, max_control_len = nil, scm_rights: false)
+    internal_recvmsg(max_msg_len, flags, max_control_len, scm_rights, true)
+  end
+
+  def recvmsg_nonblock(max_msg_len = nil, flags = 0, max_control_len = nil, exception: true, scm_rights: false)
     fcntl(Fcntl::F_SETFL, Fcntl::O_NONBLOCK)
 
-    recvmsg(max_msg_len, flags | Socket::MSG_DONTWAIT)
+    internal_recvmsg(max_msg_len, flags | Socket::MSG_DONTWAIT, max_control_len, scm_rights, exception)
   end
 
-  def sendmsg(message, flags = 0, dest_sockaddr = nil, *_)
+  private def internal_sendmsg(message, flags, dest_sockaddr, exception)
     msg_buffer = Truffle::Socket::Foreign.char_pointer(message.bytesize)
     io_vec     = Truffle::Socket::Foreign::Iovec.with_buffer(msg_buffer)
     header     = Truffle::Socket::Foreign::Msghdr.new
@@ -247,7 +279,13 @@ def sendmsg(message, flags = 0, dest_sockaddr = nil, *_)
       num_bytes = Truffle::Socket::Foreign
         .sendmsg(@descriptor, header.pointer, flags)
 
-      Truffle::Socket::Error.read_error('sendmsg(2)', self) if num_bytes < 0
+      if num_bytes < 0
+        if exception
+          Truffle::Socket::Error.read_error('sendmsg(2)', self)
+        else
+          return :wait_writable
+        end
+      end
 
       num_bytes
     ensure
@@ -258,10 +296,14 @@ def sendmsg(message, flags = 0, dest_sockaddr = nil, *_)
     end
   end
 
-  def sendmsg_nonblock(message, flags = 0, dest_sockaddr = nil, *_)
+  def sendmsg(message, flags = 0, dest_sockaddr = nil, *controls)
+    internal_sendmsg(message, flags, dest_sockaddr, true)
+  end
+
+  def sendmsg_nonblock(message, flags = 0, dest_sockaddr = nil, *controls, exception: true)
     fcntl(Fcntl::F_SETFL, Fcntl::O_NONBLOCK)
 
-    sendmsg(message, flags | Socket::MSG_DONTWAIT, dest_sockaddr)
+    internal_sendmsg(message, flags | Socket::MSG_DONTWAIT, dest_sockaddr, exception)
   end
 
   def close_read
@@ -292,14 +334,6 @@ def close_write
     nil
   end
 
-  def recv_nonblock(bytes_to_read, flags = 0)
-    fcntl(Fcntl::F_SETFL, Fcntl::O_NONBLOCK)
-
-    Truffle::Socket::Error.wrap_read_nonblock do
-      recv(bytes_to_read, flags)
-    end
-  end
-
   def shutdown(how = Socket::SHUT_RDWR)
     how = Truffle::Socket.shutdown_option(how)
     err = Truffle::Socket::Foreign.shutdown(@descriptor, how)

lib/truffle/socket/ip_socket.rb

@@ -39,10 +39,10 @@ def peeraddr(reverse_lookup=nil)
     Truffle::Socket.address_info(:getpeername, self, reverse_lookup)
   end
 
-  def recvfrom(maxlen, flags = 0)
-    flags = 0 if flags.nil?
+  private def internal_recvfrom(maxlen, flags, exception)
+    message, addr = internal_recvmsg(maxlen, flags, nil, false, exception)
 
-    message, addr = recvmsg(maxlen, flags)
+    return message if message == :wait_readable
 
     aname    = Truffle::Socket.address_family_name(addr.afamily)
     hostname = addr.ip_address
@@ -60,4 +60,9 @@ def recvfrom(maxlen, flags = 0)
 
     [message, [aname, addr.ip_port, hostname, addr.ip_address]]
   end
+
+  def recvfrom(maxlen, flags = 0)
+    flags = 0 if flags.nil?
+    internal_recvfrom(maxlen, flags, true)
+  end
 end

lib/truffle/socket/socket.rb

@@ -377,10 +377,14 @@ def recvfrom(bytes, flags = 0)
     [message, addr]
   end
 
-  def recvfrom_nonblock(bytes, flags = 0)
-    message, addr = recvmsg_nonblock(bytes, flags)
+  def recvfrom_nonblock(bytes, flags = 0, exception: true)
+    message, addr = recvmsg_nonblock(bytes, flags, exception: exception)
 
-    [message, addr]
+    if message == :wait_readable
+      message
+    else
+      [message, addr]
+    end
   end
 
   def listen(backlog)

lib/truffle/socket/udp_socket.rb

@@ -71,14 +71,12 @@ def send(message, flags, host = nil, port = nil)
     super(message, flags, addr)
   end
 
-  def recvfrom_nonblock(maxlen, flags = 0)
+  def recvfrom_nonblock(maxlen, flags = 0, exception: true)
     fcntl(Fcntl::F_SETFL, Fcntl::O_NONBLOCK)
 
     flags = 0 if flags.nil?
 
-    flags |= Socket::MSG_DONTWAIT
-
-    recvfrom(maxlen, flags)
+    internal_recvfrom(maxlen, flags | Socket::MSG_DONTWAIT, exception)
   end
 
   def inspect