allow reusing netlink socket for inet_diag

No need to waste resources on creating/destroying
a socket.

4 files changed, 86 insertions, 13 deletions

diff --git a/examples/linux-listener-stats.rb b/examples/linux-listener-stats.rb
index 5826296..70d60c6 100755
--- a/examples/linux-listener-stats.rb
+++ b/examples/linux-listener-stats.rb
@@ -79,6 +79,7 @@ end
 combined = {}
 tcp = nil if tcp.empty?
 unix = nil if unix.empty?
+sock = Raindrops::InetDiagSocket.new if tcp
 
 len = 35 if len > 35
 fmt = "%20s % #{len}s % 10u % 10u\n"
@@ -89,7 +90,7 @@ begin
     combined.clear
     now = nil
   end
-  tcp and combined.merge! Raindrops::Linux.tcp_listener_stats(tcp)
+  tcp and combined.merge! Raindrops::Linux.tcp_listener_stats(tcp, sock)
   unix and combined.merge! Raindrops::Linux.unix_listener_stats(unix)
   combined.each do |addr,stats|
     active, queued = stats.active, stats.queued
diff --git a/ext/raindrops/linux_inet_diag.c b/ext/raindrops/linux_inet_diag.c
index 4d3b399..13ba4ff 100644
--- a/ext/raindrops/linux_inet_diag.c
+++ b/ext/raindrops/linux_inet_diag.c
@@ -4,6 +4,7 @@
 #else
 #  include <st.h>
 #endif
+#include "my_fileno.h"
 #ifdef __linux__
 
 /* Ruby 1.8.6+ macros (for compatibility with Ruby 1.9) */
@@ -48,7 +49,8 @@ rb_thread_blocking_region(
 
 static size_t page_size;
 static unsigned g_seq;
-static VALUE cListenStats;
+static VALUE cListenStats, cIDSock;
+static ID id_new;
 
 struct listen_stats {
         uint32_t active;
@@ -64,8 +66,25 @@ struct nogvl_args {
         st_table *table;
         struct iovec iov[3]; /* last iov holds inet_diag bytecode */
         struct listen_stats stats;
+        int fd;
 };
 
+/*
+ * call-seq:
+ *        Raindrops::InetDiagSocket.new        -> Socket
+ *
+ * Creates a new Socket object for the netlink inet_diag facility
+ */
+static VALUE ids_s_new(VALUE klass)
+{
+        VALUE argv[3];
+        argv[0] = INT2NUM(AF_NETLINK);
+        argv[1] = INT2NUM(SOCK_RAW);
+        argv[2] = INT2NUM(NETLINK_INET_DIAG);
+
+        return rb_call_super(3, argv);
+}
+
 /* creates a Ruby ListenStats Struct based on our internal listen_stats */
 static VALUE rb_listen_stats(struct listen_stats *stats)
 {
@@ -231,7 +250,6 @@ static inline void r_acc(struct nogvl_args *args, struct inet_diag_msg *r)
          */
 }
 
-static const char err_socket[] = "socket";
 static const char err_sendmsg[] = "sendmsg";
 static const char err_recvmsg[] = "recvmsg";
 static const char err_nlmsg[] = "nlmsg";
@@ -300,15 +318,11 @@ static VALUE diag(void *ptr)
         struct msghdr msg;
         const char *err = NULL;
         unsigned seq = ++g_seq;
-        int fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_INET_DIAG);
-
-        if (fd < 0)
-                return (VALUE)err_socket;
 
         prep_diag_args(args, &nladdr, &rta, &req, &msg);
         req.nlh.nlmsg_seq = seq;
 
-        if (sendmsg(fd, &msg, 0) < 0) {
+        if (sendmsg(args->fd, &msg, 0) < 0) {
                 err = err_sendmsg;
                 goto out;
         }
@@ -321,7 +335,7 @@ static VALUE diag(void *ptr)
                 struct nlmsghdr *h = (struct nlmsghdr *)args->iov[0].iov_base;
 
                 prep_msghdr(&msg, args, &nladdr, 1);
-                readed = recvmsg(fd, &msg, 0);
+                readed = recvmsg(args->fd, &msg, 0);
                 if (readed < 0) {
                         if (errno == EINTR)
                                 continue;
@@ -346,7 +360,6 @@ static VALUE diag(void *ptr)
 out:
         {
                 int save_errno = errno;
-                close(fd);
                 if (err && args->table) {
                         st_foreach(args->table, st_free_data, 0);
                         st_free_table(args->table);
@@ -491,7 +504,7 @@ static VALUE tcp_stats(struct nogvl_args *args, VALUE addr)
 
 /*
  * call-seq:
- *      Raindrops::Linux.tcp_listener_stats([addrs]) => hash
+ *      Raindrops::Linux.tcp_listener_stats([addrs[, sock]]) => hash
  *
  * If specified, +addr+ may be a string or array of strings representing
  * listen addresses to filter for. Returns a hash with given addresses as
@@ -500,6 +513,7 @@ static VALUE tcp_stats(struct nogvl_args *args, VALUE addr)
  *      addrs = %w(0.0.0.0:80 127.0.0.1:8080)
  *
  * If +addr+ is nil or not specified, all (IPv4) addresses are returned.
+ * If +sock+ is specified, it should be a Raindrops::InetDiagSock object.
  */
 static VALUE tcp_listener_stats(int argc, VALUE *argv, VALUE self)
 {
@@ -507,9 +521,9 @@ static VALUE tcp_listener_stats(int argc, VALUE *argv, VALUE self)
         long i;
         VALUE rv = rb_hash_new();
         struct nogvl_args args;
-        VALUE addrs;
+        VALUE addrs, sock;
 
-        rb_scan_args(argc, argv, "01", &addrs);
+        rb_scan_args(argc, argv, "02", &addrs, &sock);
 
         /*
          * allocating page_size instead of OP_LEN since we'll reuse the
@@ -519,6 +533,9 @@ static VALUE tcp_listener_stats(int argc, VALUE *argv, VALUE self)
         args.iov[2].iov_len = OPLEN;
         args.iov[2].iov_base = alloca(page_size);
         args.table = NULL;
+        if (NIL_P(sock))
+                sock = rb_funcall(cIDSock, id_new, 0);
+        args.fd = my_fileno(sock);
 
         switch (TYPE(addrs)) {
         case T_STRING:
@@ -551,6 +568,9 @@ static VALUE tcp_listener_stats(int argc, VALUE *argv, VALUE self)
 
         st_foreach(args.table, NIL_P(addrs) ? st_to_hash : st_AND_hash, rv);
         st_free_table(args.table);
+
+        /* let GC deal with corner cases */
+        if (argc < 2) rb_io_close(sock);
         return rv;
 }
 
@@ -559,6 +579,19 @@ void Init_raindrops_linux_inet_diag(void)
         VALUE cRaindrops = rb_const_get(rb_cObject, rb_intern("Raindrops"));
         VALUE mLinux = rb_define_module_under(cRaindrops, "Linux");
 
+        rb_require("socket");
+        cIDSock = rb_const_get(rb_cObject, rb_intern("Socket"));
+        id_new = rb_intern("new");
+
+        /*
+         * Document-class: Raindrops::InetDiag::Socket
+         *
+         * This is a subclass of +Socket+ specifically for talking
+         * to the inet_diag facility of Netlink.
+         */
+        cIDSock = rb_define_class_under(cRaindrops, "InetDiagSocket", cIDSock);
+        rb_define_singleton_method(cIDSock, "new", ids_s_new, 0);
+
         cListenStats = rb_const_get(cRaindrops, rb_intern("ListenStats"));
 
         rb_define_module_function(mLinux, "tcp_listener_stats",
diff --git a/test/test_inet_diag_socket.rb b/test/test_inet_diag_socket.rb
new file mode 100644
index 0000000..c40d50e
--- /dev/null
+++ b/test/test_inet_diag_socket.rb
@@ -0,0 +1,13 @@
+# -*- encoding: binary -*-
+require 'test/unit'
+require 'raindrops'
+$stderr.sync = $stdout.sync = true
+
+class TestInetDiagSocket < Test::Unit::TestCase
+  def test_new
+    sock = Raindrops::InetDiagSocket.new
+    assert_kind_of Socket, sock
+    assert_kind_of Fixnum, sock.fileno
+    assert_nil sock.close
+  end
+end if RUBY_PLATFORM =~ /linux/
diff --git a/test/test_linux.rb b/test/test_linux.rb
index 5a7269d..d3c8da7 100644
--- a/test/test_linux.rb
+++ b/test/test_linux.rb
@@ -62,6 +62,32 @@ class TestLinux < Test::Unit::TestCase
     assert_equal 1, stats[addr].active
   end
 
+  def test_tcp_reuse_sock
+    nlsock = Raindrops::InetDiagSocket.new
+    s = TCPServer.new(TEST_ADDR, 0)
+    port = s.addr[1]
+    addr = "#{TEST_ADDR}:#{port}"
+    addrs = [ addr ]
+    stats = tcp_listener_stats(addrs, nlsock)
+    assert_equal 1, stats.size
+    assert_equal 0, stats[addr].queued
+    assert_equal 0, stats[addr].active
+
+    c = TCPSocket.new(TEST_ADDR, port)
+    stats = tcp_listener_stats(addrs, nlsock)
+    assert_equal 1, stats.size
+    assert_equal 1, stats[addr].queued
+    assert_equal 0, stats[addr].active
+
+    sc = s.accept
+    stats = tcp_listener_stats(addrs, nlsock)
+    assert_equal 1, stats.size
+    assert_equal 0, stats[addr].queued
+    assert_equal 1, stats[addr].active
+    ensure
+      nlsock.close
+  end
+
   def test_tcp_multi
     s1 = TCPServer.new(TEST_ADDR, 0)
     s2 = TCPServer.new(TEST_ADDR, 0)