From c478549ea7490de2432ed31ffee37a2bfc1d24f3 Mon Sep 17 00:00:00 2001 From: Eric Wong Date: Sat, 12 Mar 2011 12:56:15 -0800 Subject: [PATCH] allow reusing netlink socket for inet_diag No need to waste resources on creating/destroying a socket. --- examples/linux-listener-stats.rb | 3 ++- ext/raindrops/linux_inet_diag.c | 57 +++++++++++++++++++++++++++++++--------- test/test_inet_diag_socket.rb | 13 +++++++++ test/test_linux.rb | 26 ++++++++++++++++++ 4 files changed, 86 insertions(+), 13 deletions(-) create mode 100644 test/test_inet_diag_socket.rb 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 #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) -- 2.11.4.GIT