Support timeout for Addrinfo

Addrinfo.getaddrinfo and .foreach now accepts :timeout in seconds as
a keyword argument. If getaddrinfo_a(3) is available, the timeout will be
applied for name resolution. Otherwise, it will be ignored.

Socket.tcp accepts :resolv_timeout to use this feature.

This commit is retry of 6382f5cc91ac9e36776bc854632d9a1237250da7.
Test was failed on Solaris machines which don't have "http" in
/etc/services. In this commit, use "ssh" instead.

6 files changed, 157 insertions, 13 deletions

diff --git a/ext/socket/extconf.rb b/ext/socket/extconf.rb
index 0cc8a88d5c..42cfb064c9 100644
--- a/ext/socket/extconf.rb
+++ b/ext/socket/extconf.rb
@@ -444,6 +444,7 @@ else
   test_func = "socket(0,0,0)"
   have_library("nsl", 't_open("", 0, (struct t_info *)NULL)', headers) # SunOS
   have_library("socket", "socket(0,0,0)", headers) # SunOS
+  have_library("anl", 'getaddrinfo_a', headers)
 end
 
 if have_func(test_func, headers)
@@ -505,6 +506,7 @@ EOF
   unless have_func("gethostname((char *)0, 0)", headers)
     have_func("uname((struct utsname *)NULL)", headers)
   end
+  have_func("getaddrinfo_a", headers)
 
   ipv6 = false
   default_ipv6 = /haiku/ !~ RUBY_PLATFORM
diff --git a/ext/socket/lib/socket.rb b/ext/socket/lib/socket.rb
index 4ed2df23e6..d756a32a5a 100644
--- a/ext/socket/lib/socket.rb
+++ b/ext/socket/lib/socket.rb
@@ -223,8 +223,8 @@ class Addrinfo
   #   #   #<Addrinfo: [::1]:80 TCP (:80)>
   #   #   #<Addrinfo: [::1]:80 UDP (:80)>
   #
-  def self.foreach(nodename, service, family=nil, socktype=nil, protocol=nil, flags=nil, &block)
-    Addrinfo.getaddrinfo(nodename, service, family, socktype, protocol, flags).each(&block)
+  def self.foreach(nodename, service, family=nil, socktype=nil, protocol=nil, flags=nil, timeout: nil, &block)
+    Addrinfo.getaddrinfo(nodename, service, family, socktype, protocol, flags, timeout: timeout).each(&block)
   end
 end
 
@@ -606,6 +606,7 @@ class Socket < BasicSocket
   # _opts_ may have following options:
   #
   # [:connect_timeout] specify the timeout in seconds.
+  # [:resolv_timeout] specify the name resolution timeout in seconds.
   #
   # If a block is given, the block is called with the socket.
   # The value of the block is returned.
@@ -619,7 +620,7 @@ class Socket < BasicSocket
   #     puts sock.read
   #   }
   #
-  def self.tcp(host, port, local_host = nil, local_port = nil, connect_timeout: nil) # :yield: socket
+  def self.tcp(host, port, local_host = nil, local_port = nil, connect_timeout: nil, resolv_timeout: nil) # :yield: socket
     last_error = nil
     ret = nil
 
@@ -628,7 +629,7 @@ class Socket < BasicSocket
       local_addr_list = Addrinfo.getaddrinfo(local_host, local_port, nil, :STREAM, nil)
     end
 
-    Addrinfo.foreach(host, port, nil, :STREAM) {|ai|
+    Addrinfo.foreach(host, port, nil, :STREAM, timeout: resolv_timeout) {|ai|
       if local_addr_list
         local_addr = local_addr_list.find {|local_ai| local_ai.afamily == ai.afamily }
         next unless local_addr
diff --git a/ext/socket/raddrinfo.c b/ext/socket/raddrinfo.c
index dad11ad1ac..0499b845b5 100644
--- a/ext/socket/raddrinfo.c
+++ b/ext/socket/raddrinfo.c
@@ -199,6 +199,27 @@ nogvl_getaddrinfo(void *arg)
 }
 #endif
 
+#ifdef HAVE_GETADDRINFO_A
+struct gai_suspend_arg
+{
+    struct gaicb *req;
+    struct timespec *timeout;
+};
+
+static void *
+nogvl_gai_suspend(void *arg)
+{
+    int ret;
+    struct gai_suspend_arg *ptr = arg;
+    struct gaicb const *wait_reqs[1];
+
+    wait_reqs[0] = ptr->req;
+    ret = gai_suspend(wait_reqs, 1, ptr->timeout);
+
+    return (void *)(VALUE)ret;
+}
+#endif
+
 static int
 numeric_getaddrinfo(const char *node, const char *service,
         const struct addrinfo *hints,
@@ -318,6 +339,59 @@ rb_getaddrinfo(const char *node, const char *service,
     return ret;
 }
 
+#ifdef HAVE_GETADDRINFO_A
+int
+rb_getaddrinfo_a(const char *node, const char *service,
+               const struct addrinfo *hints,
+               struct rb_addrinfo **res, struct timespec *timeout)
+{
+    struct addrinfo *ai;
+    int ret;
+    int allocated_by_malloc = 0;
+
+    ret = numeric_getaddrinfo(node, service, hints, &ai);
+    if (ret == 0)
+        allocated_by_malloc = 1;
+    else {
+	struct gai_suspend_arg arg;
+	struct gaicb *reqs[1];
+	struct gaicb req;
+
+	req.ar_name = node;
+	req.ar_service = service;
+	req.ar_request = hints;
+
+	reqs[0] = &req;
+	ret = getaddrinfo_a(GAI_NOWAIT, reqs, 1, NULL);
+	if (ret) return ret;
+
+	arg.req = &req;
+	arg.timeout = timeout;
+
+	ret = (int)(VALUE)rb_thread_call_without_gvl(nogvl_gai_suspend, &arg, RUBY_UBF_IO, 0);
+
+	if (ret) {
+	    /* on Ubuntu 18.04 (or other systems), gai_suspend(3) returns EAI_SYSTEM/ENOENT on timeout */
+	    if (ret == EAI_SYSTEM && errno == ENOENT) {
+		return EAI_AGAIN;
+	    } else {
+		return ret;
+	    }
+	}
+
+	ret = gai_error(reqs[0]);
+	ai = reqs[0]->ar_result;
+    }
+
+    if (ret == 0) {
+        *res = (struct rb_addrinfo *)xmalloc(sizeof(struct rb_addrinfo));
+        (*res)->allocated_by_malloc = allocated_by_malloc;
+        (*res)->ai = ai;
+    }
+    return ret;
+}
+#endif
+
 void
 rb_freeaddrinfo(struct rb_addrinfo *ai)
 {
@@ -530,6 +604,42 @@ rsock_getaddrinfo(VALUE host, VALUE port, struct addrinfo *hints, int socktype_h
     return res;
 }
 
+#ifdef HAVE_GETADDRINFO_A
+static struct rb_addrinfo*
+rsock_getaddrinfo_a(VALUE host, VALUE port, struct addrinfo *hints, int socktype_hack, VALUE timeout)
+{
+    struct rb_addrinfo* res = NULL;
+    char *hostp, *portp;
+    int error;
+    char hbuf[NI_MAXHOST], pbuf[NI_MAXSERV];
+    int additional_flags = 0;
+
+    hostp = host_str(host, hbuf, sizeof(hbuf), &additional_flags);
+    portp = port_str(port, pbuf, sizeof(pbuf), &additional_flags);
+
+    if (socktype_hack && hints->ai_socktype == 0 && str_is_number(portp)) {
+       hints->ai_socktype = SOCK_DGRAM;
+    }
+    hints->ai_flags |= additional_flags;
+
+    if (NIL_P(timeout)) {
+	error = rb_getaddrinfo(hostp, portp, hints, &res);
+    } else {
+	struct timespec _timeout = rb_time_timespec_interval(timeout);
+	error = rb_getaddrinfo_a(hostp, portp, hints, &res, &_timeout);
+    }
+
+    if (error) {
+        if (hostp && hostp[strlen(hostp)-1] == '\n') {
+            rb_raise(rb_eSocket, "newline at the end of hostname");
+        }
+        rsock_raise_socket_error("getaddrinfo_a", error);
+    }
+
+    return res;
+}
+#endif
+
 int
 rsock_fd_family(int fd)
 {
@@ -817,7 +927,7 @@ rsock_addrinfo_new(struct sockaddr *addr, socklen_t len,
 static struct rb_addrinfo *
 call_getaddrinfo(VALUE node, VALUE service,
                  VALUE family, VALUE socktype, VALUE protocol, VALUE flags,
-                 int socktype_hack)
+                 int socktype_hack, VALUE timeout)
 {
     struct addrinfo hints;
     struct rb_addrinfo *res;
@@ -834,7 +944,16 @@ call_getaddrinfo(VALUE node, VALUE service,
     if (!NIL_P(flags)) {
 	hints.ai_flags = NUM2INT(flags);
     }
+
+#ifdef HAVE_GETADDRINFO_A
+    if (NIL_P(timeout)) {
+	res = rsock_getaddrinfo(node, service, &hints, socktype_hack);
+    } else {
+	res = rsock_getaddrinfo_a(node, service, &hints, socktype_hack, timeout);
+    }
+#else
     res = rsock_getaddrinfo(node, service, &hints, socktype_hack);
+#endif
 
     if (res == NULL)
 	rb_raise(rb_eSocket, "host not found");
@@ -848,7 +967,7 @@ init_addrinfo_getaddrinfo(rb_addrinfo_t *rai, VALUE node, VALUE service,
                           VALUE family, VALUE socktype, VALUE protocol, VALUE flags,
                           VALUE inspectnode, VALUE inspectservice)
 {
-    struct rb_addrinfo *res = call_getaddrinfo(node, service, family, socktype, protocol, flags, 1);
+    struct rb_addrinfo *res = call_getaddrinfo(node, service, family, socktype, protocol, flags, 1, Qnil);
     VALUE canonname;
     VALUE inspectname = rb_str_equal(node, inspectnode) ? Qnil : make_inspectname(inspectnode, inspectservice, res->ai);
 
@@ -918,7 +1037,7 @@ addrinfo_firstonly_new(VALUE node, VALUE service, VALUE family, VALUE socktype,
     VALUE canonname;
     VALUE inspectname;
 
-    struct rb_addrinfo *res = call_getaddrinfo(node, service, family, socktype, protocol, flags, 0);
+    struct rb_addrinfo *res = call_getaddrinfo(node, service, family, socktype, protocol, flags, 0, Qnil);
 
     inspectname = make_inspectname(node, service, res->ai);
 
@@ -938,13 +1057,13 @@ addrinfo_firstonly_new(VALUE node, VALUE service, VALUE family, VALUE socktype,
 }
 
 static VALUE
-addrinfo_list_new(VALUE node, VALUE service, VALUE family, VALUE socktype, VALUE protocol, VALUE flags)
+addrinfo_list_new(VALUE node, VALUE service, VALUE family, VALUE socktype, VALUE protocol, VALUE flags, VALUE timeout)
 {
     VALUE ret;
     struct addrinfo *r;
     VALUE inspectname;
 
-    struct rb_addrinfo *res = call_getaddrinfo(node, service, family, socktype, protocol, flags, 0);
+    struct rb_addrinfo *res = call_getaddrinfo(node, service, family, socktype, protocol, flags, 0, timeout);
 
     inspectname = make_inspectname(node, service, res->ai);
 
@@ -1696,7 +1815,7 @@ addrinfo_mload(VALUE self, VALUE ary)
 #endif
         res = call_getaddrinfo(rb_ary_entry(pair, 0), rb_ary_entry(pair, 1),
                                INT2NUM(pfamily), INT2NUM(socktype), INT2NUM(protocol),
-                               INT2NUM(flags), 1);
+                               INT2NUM(flags), 1, Qnil);
 
         len = res->ai->ai_addrlen;
         memcpy(&ss, res->ai->ai_addr, res->ai->ai_addrlen);
@@ -2331,6 +2450,8 @@ addrinfo_unix_path(VALUE self)
 }
 #endif
 
+static ID id_timeout;
+
 /*
  * call-seq:
  *   Addrinfo.getaddrinfo(nodename, service, family, socktype, protocol, flags) => [addrinfo, ...]
@@ -2377,10 +2498,16 @@ addrinfo_unix_path(VALUE self)
 static VALUE
 addrinfo_s_getaddrinfo(int argc, VALUE *argv, VALUE self)
 {
-    VALUE node, service, family, socktype, protocol, flags;
+    VALUE node, service, family, socktype, protocol, flags, opts, timeout;
+
+    rb_scan_args(argc, argv, "24:", &node, &service, &family, &socktype,
+		 &protocol, &flags, &opts);
+    rb_get_kwargs(opts, &id_timeout, 0, 1, &timeout);
+    if (timeout == Qundef) {
+	timeout = Qnil;
+    }
 
-    rb_scan_args(argc, argv, "24", &node, &service, &family, &socktype, &protocol, &flags);
-    return addrinfo_list_new(node, service, family, socktype, protocol, flags);
+    return addrinfo_list_new(node, service, family, socktype, protocol, flags, timeout);
 }
 
 /*
@@ -2565,6 +2692,8 @@ rsock_init_addrinfo(void)
      * The Addrinfo class maps <tt>struct addrinfo</tt> to ruby.  This
      * structure identifies an Internet host and a service.
      */
+    id_timeout = rb_intern("timeout");
+
     rb_cAddrinfo = rb_define_class("Addrinfo", rb_cData);
     rb_define_alloc_func(rb_cAddrinfo, addrinfo_s_allocate);
     rb_define_method(rb_cAddrinfo, "initialize", addrinfo_initialize, -1);
diff --git a/include/ruby/intern.h b/include/ruby/intern.h
index ed1a4cfdd5..4fa00d55e6 100644
--- a/include/ruby/intern.h
+++ b/include/ruby/intern.h
@@ -936,6 +936,7 @@ VALUE rb_time_num_new(VALUE, VALUE);
 struct timeval rb_time_interval(VALUE num);
 struct timeval rb_time_timeval(VALUE time);
 struct timespec rb_time_timespec(VALUE time);
+struct timespec rb_time_timespec_interval(VALUE num);
 VALUE rb_time_utc_offset(VALUE time);
 /* variable.c */
 VALUE rb_mod_name(VALUE);
diff --git a/test/socket/test_addrinfo.rb b/test/socket/test_addrinfo.rb
index a06f3eb451..68341c2148 100644
--- a/test/socket/test_addrinfo.rb
+++ b/test/socket/test_addrinfo.rb
@@ -688,5 +688,10 @@ class TestSocketAddrinfo < Test::Unit::TestCase
       assert_equal(ai1.canonname, ai2.canonname)
     end
 
+    def test_addrinfo_timeout
+      ai = Addrinfo.getaddrinfo("localhost", "ssh", Socket::PF_INET, Socket::SOCK_STREAM, timeout: 1).fetch(0)
+      assert_equal(6, ai.protocol)
+      assert_equal(22, ai.ip_port)
+    end
   end
 end
diff --git a/time.c b/time.c
index 1e760f6ca5..18bed6d68a 100644
--- a/time.c
+++ b/time.c
@@ -2698,6 +2698,12 @@ rb_time_timespec(VALUE time)
     return time_timespec(time, FALSE);
 }
 
+struct timespec
+rb_time_timespec_interval(VALUE num)
+{
+    return time_timespec(num, TRUE);
+}
+
 enum {
     TMOPT_IN,
     TMOPT_MAX_