raindrops 0.4.1 → 0.5.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (49) hide show
  1. data/.document +2 -1
  2. data/.gitignore +4 -0
  3. data/.wrongdoc.yml +4 -0
  4. data/GIT-VERSION-GEN +1 -1
  5. data/GNUmakefile +2 -196
  6. data/Gemfile +7 -0
  7. data/LICENSE +1 -1
  8. data/README +17 -47
  9. data/Rakefile +0 -104
  10. data/examples/linux-listener-stats.rb +123 -0
  11. data/examples/{config.ru → middleware.ru} +1 -1
  12. data/examples/watcher.ru +4 -0
  13. data/examples/watcher_demo.ru +13 -0
  14. data/examples/zbatery.conf.rb +13 -0
  15. data/ext/raindrops/extconf.rb +5 -0
  16. data/ext/raindrops/linux_inet_diag.c +449 -151
  17. data/ext/raindrops/linux_tcp_info.c +170 -0
  18. data/ext/raindrops/my_fileno.h +36 -0
  19. data/ext/raindrops/raindrops.c +232 -20
  20. data/lib/raindrops.rb +20 -7
  21. data/lib/raindrops/aggregate.rb +8 -0
  22. data/lib/raindrops/aggregate/last_data_recv.rb +86 -0
  23. data/lib/raindrops/aggregate/pmq.rb +239 -0
  24. data/lib/raindrops/last_data_recv.rb +100 -0
  25. data/lib/raindrops/linux.rb +26 -16
  26. data/lib/raindrops/middleware.rb +112 -41
  27. data/lib/raindrops/middleware/proxy.rb +34 -0
  28. data/lib/raindrops/struct.rb +15 -0
  29. data/lib/raindrops/watcher.rb +362 -0
  30. data/pkg.mk +171 -0
  31. data/raindrops.gemspec +10 -20
  32. data/test/ipv6_enabled.rb +10 -0
  33. data/test/rack_unicorn.rb +12 -0
  34. data/test/test_aggregate_pmq.rb +65 -0
  35. data/test/test_inet_diag_socket.rb +13 -0
  36. data/test/test_last_data_recv_unicorn.rb +69 -0
  37. data/test/test_linux.rb +55 -57
  38. data/test/test_linux_all_tcp_listen_stats.rb +66 -0
  39. data/test/test_linux_all_tcp_listen_stats_leak.rb +43 -0
  40. data/test/test_linux_ipv6.rb +158 -0
  41. data/test/test_linux_tcp_info.rb +61 -0
  42. data/test/test_middleware.rb +15 -2
  43. data/test/test_middleware_unicorn.rb +37 -0
  44. data/test/test_middleware_unicorn_ipv6.rb +37 -0
  45. data/test/test_raindrops.rb +65 -1
  46. data/test/test_raindrops_gc.rb +23 -1
  47. data/test/test_watcher.rb +85 -0
  48. metadata +69 -22
  49. data/examples/linux-tcp-listener-stats.rb +0 -44
@@ -0,0 +1,123 @@
1
+ #!/usr/bin/ruby
2
+ # -*- encoding: binary -*-
3
+ $stdout.sync = $stderr.sync = true
4
+ # this is used to show or watch the number of active and queued
5
+ # connections on any listener socket from the command line
6
+
7
+ require 'raindrops'
8
+ require 'optparse'
9
+ require 'ipaddr'
10
+ require 'time'
11
+ begin
12
+ require 'sleepy_penguin'
13
+ rescue LoadError
14
+ end
15
+ usage = "Usage: #$0 [-d DELAY] [-t QUEUED_THRESHOLD] ADDR..."
16
+ ARGV.size > 0 or abort usage
17
+ delay = false
18
+ all = false
19
+ queued_thresh = -1
20
+ # "normal" exits when driven on the command-line
21
+ trap(:INT) { exit 130 }
22
+ trap(:PIPE) { exit 0 }
23
+
24
+ opts = OptionParser.new('', 24, ' ') do |opts|
25
+ opts.banner = usage
26
+ opts.on('-d', '--delay=DELAY', Float) { |n| delay = n }
27
+ opts.on('-t', '--queued-threshold=INT', Integer) { |n| queued_thresh = n }
28
+ opts.on('-a', '--all') { all = true }
29
+ opts.parse! ARGV
30
+ end
31
+
32
+ begin
33
+ require 'aggregate'
34
+ rescue LoadError
35
+ $stderr.puts "Aggregate missing, USR1 and USR2 handlers unavailable"
36
+ end if delay
37
+
38
+ if delay && defined?(SleepyPenguin::TimerFD)
39
+ @tfd = SleepyPenguin::TimerFD.new
40
+ @tfd.settime nil, delay, delay
41
+ def delay_for(seconds)
42
+ @tfd.expirations
43
+ end
44
+ else
45
+ alias delay_for sleep
46
+ end
47
+
48
+ agg_active = agg_queued = nil
49
+ if delay && defined?(Aggregate)
50
+ agg_active = Aggregate.new
51
+ agg_queued = Aggregate.new
52
+
53
+ def dump_aggregate(label, agg)
54
+ $stderr.write "--- #{label} ---\n"
55
+ %w(count min max outliers_low outliers_high mean std_dev).each do |f|
56
+ $stderr.write "#{f}=#{agg.__send__ f}\n"
57
+ end
58
+ $stderr.write "#{agg}\n\n"
59
+ end
60
+
61
+ trap(:USR1) do
62
+ dump_aggregate "active", agg_active
63
+ dump_aggregate "queued", agg_queued
64
+ end
65
+ trap(:USR2) do
66
+ agg_active = Aggregate.new
67
+ agg_queued = Aggregate.new
68
+ end
69
+ $stderr.puts "USR1(dump_aggregate) and USR2(reset) handlers ready for PID=#$$"
70
+ end
71
+
72
+ ARGV.each do |addr|
73
+ addr =~ %r{\A(127\..+):(\d+)\z} or next
74
+ host, port = $1, $2
75
+ hex_port = '%X' % port.to_i
76
+ ip_addr = IPAddr.new(host)
77
+ hex_host = ip_addr.hton.each_byte.inject('') { |s,o| s << '%02X' % o }
78
+ socks = File.readlines('/proc/net/tcp')
79
+ hex_addr = "#{hex_host}:#{hex_port}"
80
+ if socks.grep(/^\s+\d+:\s+#{hex_addr}\s+/).empty? &&
81
+ ! socks.grep(/^\s+\d+:\s+00000000:#{hex_port}\s+/).empty?
82
+ warn "W: #{host}:#{port} (#{hex_addr}) not found in /proc/net/tcp"
83
+ warn "W: Did you mean 0.0.0.0:#{port}?"
84
+ end
85
+ end
86
+
87
+ len = "address".size
88
+ now = nil
89
+ tcp, unix = [], []
90
+ ARGV.each do |addr|
91
+ bs = addr.respond_to?(:bytesize) ? addr.bytesize : addr.size
92
+ len = bs if bs > len
93
+ (addr =~ %r{\A/} ? unix : tcp) << addr
94
+ end
95
+ combined = {}
96
+ tcp_args = unix_args = nil
97
+ unless tcp.empty? && unix.empty?
98
+ tcp_args = tcp
99
+ unix_args = unix
100
+ end
101
+ sock = Raindrops::InetDiagSocket.new if tcp
102
+
103
+ len = 35 if len > 35
104
+ fmt = "%20s % #{len}s % 10u % 10u\n"
105
+ $stderr.printf fmt.tr('u','s'), *%w(timestamp address active queued)
106
+
107
+ begin
108
+ if now
109
+ combined.clear
110
+ now = nil
111
+ end
112
+ combined.merge! Raindrops::Linux.tcp_listener_stats(tcp_args, sock)
113
+ combined.merge! Raindrops::Linux.unix_listener_stats(unix_args)
114
+ combined.each do |addr,stats|
115
+ active, queued = stats.active, stats.queued
116
+ if agg_active
117
+ agg_active << active
118
+ agg_queued << queued
119
+ end
120
+ next if queued < queued_thresh
121
+ printf fmt, now ||= Time.now.utc.iso8601, addr, active, queued
122
+ end
123
+ end while delay && delay_for(delay)
@@ -1,4 +1,4 @@
1
- # sample stand-alone rackup application
1
+ # sample stand-alone rackup application for Raindrops::Middleware
2
2
  require 'rack/lobster'
3
3
  require 'raindrops'
4
4
  use Raindrops::Middleware
@@ -0,0 +1,4 @@
1
+ # Sample standalone Rack application, recommended use is with Zbatery
2
+ # See zbatery.conf.rb
3
+ require "raindrops"
4
+ run Raindrops::Watcher.new
@@ -0,0 +1,13 @@
1
+ # This is the exact config that powers http://raindrops-demo.bogomips.org/
2
+ # This is used with zbatery.conf.rb
3
+ #
4
+ # zbatery -c zbatery.conf.ru watcher_demo.ru -E none
5
+ require "raindrops"
6
+ use Raindrops::Middleware
7
+ listeners = %w(
8
+ 0.0.0.0:9418
9
+ 0.0.0.0:80
10
+ /tmp/.raindrops
11
+ /tmp/.r
12
+ )
13
+ run Raindrops::Watcher.new :listeners => listeners
@@ -0,0 +1,13 @@
1
+ # Used for running Raindrops::Watcher, which requires a multi-threaded
2
+ # Rack server capable of streaming a response. Threads must be used,
3
+ # so Zbatery is recommended: http://zbatery.bogomip.org/
4
+ Rainbows! do
5
+ use :ThreadSpawn
6
+ end
7
+ log_dir = "/var/log/zbatery"
8
+ if File.writable?(log_dir) && File.directory?(log_dir)
9
+ stderr_path "#{log_dir}/raindrops-demo.stderr.log"
10
+ stdout_path "#{log_dir}/raindrops-demo.stdout.log"
11
+ user "www-data", "www-data"
12
+ listen "/tmp/.raindrops"
13
+ end
@@ -3,6 +3,11 @@ require 'mkmf'
3
3
  have_func('mmap', 'sys/mman.h') or abort 'mmap() not found'
4
4
  have_func('munmap', 'sys/mman.h') or abort 'munmap() not found'
5
5
 
6
+ $CPPFLAGS += " -D_GNU_SOURCE "
7
+ have_func('mremap', 'sys/mman.h')
8
+
9
+ $CPPFLAGS += " -D_BSD_SOURCE -D_XOPEN_SOURCE=600 "
10
+ have_func("getpagesize", "unistd.h")
6
11
  have_func("rb_struct_alloc_noinit")
7
12
  have_func('rb_thread_blocking_region')
8
13
 
@@ -1,34 +1,20 @@
1
1
  #include <ruby.h>
2
+ #ifdef HAVE_RUBY_ST_H
3
+ # include <ruby/st.h>
4
+ #else
5
+ # include <st.h>
6
+ #endif
7
+ #include "my_fileno.h"
2
8
  #ifdef __linux__
3
9
 
4
10
  /* Ruby 1.8.6+ macros (for compatibility with Ruby 1.9) */
5
- #ifndef RSTRING_PTR
6
- # define RSTRING_PTR(s) (RSTRING(s)->ptr)
7
- #endif
8
11
  #ifndef RSTRING_LEN
9
12
  # define RSTRING_LEN(s) (RSTRING(s)->len)
10
13
  #endif
11
- #ifdef RSTRUCT
12
- # ifndef RSTRUCT_PTR
13
- # define RSTRUCT_PTR(s) (RSTRUCT(s)->ptr)
14
- # endif
15
- # ifndef RSTRUCT_LEN
16
- # define RSTRUCT_LEN(s) (RSTRUCT(s)->len)
17
- # endif
18
- #endif
19
-
20
- #ifndef HAVE_RB_STRUCT_ALLOC_NOINIT
21
- static ID id_new;
22
- static VALUE rb_struct_alloc_noinit(VALUE class)
23
- {
24
- return rb_funcall(class, id_new, 0, 0);
25
- }
26
- #endif /* !defined(HAVE_RB_STRUCT_ALLOC_NOINIT) */
27
14
 
28
15
  /* partial emulation of the 1.9 rb_thread_blocking_region under 1.8 */
29
16
  #ifndef HAVE_RB_THREAD_BLOCKING_REGION
30
17
  # include <rubysig.h>
31
- # define RUBY_UBF_IO ((rb_unblock_function_t *)-1)
32
18
  typedef void rb_unblock_function_t(void *);
33
19
  typedef VALUE rb_blocking_function_t(void *);
34
20
  static VALUE
@@ -50,6 +36,7 @@ rb_thread_blocking_region(
50
36
  #include <errno.h>
51
37
  #include <sys/socket.h>
52
38
  #include <sys/types.h>
39
+ #include <netdb.h>
53
40
  #include <unistd.h>
54
41
  #include <string.h>
55
42
  #include <asm/types.h>
@@ -62,59 +49,220 @@ rb_thread_blocking_region(
62
49
 
63
50
  static size_t page_size;
64
51
  static unsigned g_seq;
65
- static VALUE cListenStats;
66
-
67
- struct my_addr {
68
- in_addr_t addr;
69
- uint16_t port;
70
- };
52
+ static VALUE cListenStats, cIDSock;
53
+ static ID id_new;
71
54
 
72
55
  struct listen_stats {
73
- long active;
74
- long queued;
56
+ uint32_t active;
57
+ uint32_t listener_p:1;
58
+ uint32_t queued:31;
75
59
  };
76
60
 
77
61
  #define OPLEN (sizeof(struct inet_diag_bc_op) + \
78
- sizeof(struct inet_diag_hostcond) + \
79
- sizeof(in_addr_t))
62
+ sizeof(struct inet_diag_hostcond) + \
63
+ sizeof(struct sockaddr_storage))
80
64
 
81
65
  struct nogvl_args {
66
+ st_table *table;
82
67
  struct iovec iov[3]; /* last iov holds inet_diag bytecode */
83
- struct my_addr query_addr;
84
68
  struct listen_stats stats;
69
+ int fd;
85
70
  };
86
71
 
72
+ /*
73
+ * call-seq:
74
+ * Raindrops::InetDiagSocket.new -> Socket
75
+ *
76
+ * Creates a new Socket object for the netlink inet_diag facility
77
+ */
78
+ static VALUE ids_s_new(VALUE klass)
79
+ {
80
+ VALUE argv[3];
81
+ argv[0] = INT2NUM(AF_NETLINK);
82
+ argv[1] = INT2NUM(SOCK_RAW);
83
+ argv[2] = INT2NUM(NETLINK_INET_DIAG);
84
+
85
+ return rb_call_super(3, argv);
86
+ }
87
+
87
88
  /* creates a Ruby ListenStats Struct based on our internal listen_stats */
88
89
  static VALUE rb_listen_stats(struct listen_stats *stats)
89
90
  {
90
- VALUE rv = rb_struct_alloc_noinit(cListenStats);
91
- VALUE active = LONG2NUM(stats->active);
92
- VALUE queued = LONG2NUM(stats->queued);
93
-
94
- #ifdef RSTRUCT_PTR
95
- VALUE *ptr = RSTRUCT_PTR(rv);
96
- ptr[0] = active;
97
- ptr[1] = queued;
98
- #else /* Rubinius */
99
- rb_funcall(rv, rb_intern("active="), 1, active);
100
- rb_funcall(rv, rb_intern("queued="), 1, queued);
101
- #endif /* ! Rubinius */
102
- return rv;
91
+ VALUE active = UINT2NUM(stats->active);
92
+ VALUE queued = UINT2NUM(stats->queued);
93
+
94
+ return rb_struct_new(cListenStats, active, queued);
103
95
  }
104
96
 
105
- /*
106
- * converts a base 10 string representing a port number into
107
- * an unsigned 16 bit integer. Raises ArgumentError on failure
108
- */
109
- static uint16_t my_inet_port(const char *port)
97
+ static int st_free_data(st_data_t key, st_data_t value, st_data_t ignored)
98
+ {
99
+ xfree((void *)key);
100
+ xfree((void *)value);
101
+
102
+ return ST_DELETE;
103
+ }
104
+
105
+ static int st_to_hash(st_data_t key, st_data_t value, VALUE hash)
110
106
  {
111
- char *err;
112
- unsigned long tmp = strtoul(port, &err, 10);
107
+ struct listen_stats *stats = (struct listen_stats *)value;
113
108
 
114
- if (*err != 0 || tmp > 0xffff)
115
- rb_raise(rb_eArgError, "port not parsable: `%s'\n", port);
109
+ if (stats->listener_p) {
110
+ VALUE k = rb_str_new2((const char *)key);
111
+ VALUE v = rb_listen_stats(stats);
116
112
 
117
- return (uint16_t)tmp;
113
+ OBJ_FREEZE(k);
114
+ rb_hash_aset(hash, k, v);
115
+ }
116
+ return st_free_data(key, value, 0);
117
+ }
118
+
119
+ static int st_AND_hash(st_data_t key, st_data_t value, VALUE hash)
120
+ {
121
+ struct listen_stats *stats = (struct listen_stats *)value;
122
+
123
+ if (stats->listener_p) {
124
+ VALUE k = rb_str_new2((const char *)key);
125
+
126
+ if (rb_hash_lookup(hash, k) == Qtrue) {
127
+ VALUE v = rb_listen_stats(stats);
128
+ OBJ_FREEZE(k);
129
+ rb_hash_aset(hash, k, v);
130
+ }
131
+ }
132
+ return st_free_data(key, value, 0);
133
+ }
134
+
135
+ static const char *addr_any(sa_family_t family)
136
+ {
137
+ static const char ipv4[] = "0.0.0.0";
138
+ static const char ipv6[] = "[::]";
139
+
140
+ if (family == AF_INET)
141
+ return ipv4;
142
+ assert(family == AF_INET6 && "unknown family");
143
+ return ipv6;
144
+ }
145
+
146
+ static void bug_warn(void)
147
+ {
148
+ fprintf(stderr, "Please report how you produced this at "\
149
+ "raindrops@librelist.com\n");
150
+ fflush(stderr);
151
+ }
152
+
153
+ static struct listen_stats *stats_for(st_table *table, struct inet_diag_msg *r)
154
+ {
155
+ char *key, *port, *old_key;
156
+ size_t alloca_len;
157
+ struct listen_stats *stats;
158
+ size_t keylen;
159
+ size_t portlen = sizeof("65535");
160
+ struct sockaddr_storage ss = { 0 };
161
+ socklen_t len = sizeof(struct sockaddr_storage);
162
+ int rc;
163
+ int flags = NI_NUMERICHOST | NI_NUMERICSERV;
164
+
165
+ switch ((ss.ss_family = r->idiag_family)) {
166
+ case AF_INET: {
167
+ struct sockaddr_in *in = (struct sockaddr_in *)&ss;
168
+
169
+ in->sin_port = r->id.idiag_sport;
170
+ in->sin_addr.s_addr = r->id.idiag_src[0];
171
+ keylen = INET_ADDRSTRLEN;
172
+ alloca_len = keylen + 1 + portlen;
173
+ key = alloca(alloca_len);
174
+ key[keylen] = 0; /* will be ':' later */
175
+ port = key + keylen + 1;
176
+ rc = getnameinfo((struct sockaddr *)&ss, len,
177
+ key, keylen, port, portlen, flags);
178
+ break;
179
+ }
180
+ case AF_INET6: {
181
+ struct sockaddr_in6 *in6 = (struct sockaddr_in6 *)&ss;
182
+ in6->sin6_port = r->id.idiag_sport;
183
+ memcpy(&in6->sin6_addr.in6_u.u6_addr32,
184
+ &r->id.idiag_src, sizeof(__be32[4]));
185
+ keylen = INET6_ADDRSTRLEN;
186
+ /* [ ] */
187
+ alloca_len = 1 + keylen + 1 + 1 + portlen;
188
+ key = alloca(alloca_len);
189
+ *key = '[';
190
+ key[1 + keylen + 1] = 0; /* will be ':' later */
191
+ port = 1 + key + keylen + 1 + 1;
192
+ rc = getnameinfo((struct sockaddr *)&ss, len,
193
+ key + 1, keylen, port, portlen, flags);
194
+ break;
195
+ }
196
+ default:
197
+ assert(0 && "unsupported address family, could that be IPv7?!");
198
+ }
199
+ if (rc != 0) {
200
+ fprintf(stderr, "BUG: getnameinfo: %s\n", gai_strerror(rc));
201
+ bug_warn();
202
+ *key = 0;
203
+ }
204
+
205
+ keylen = strlen(key);
206
+ portlen = strlen(port);
207
+
208
+ switch (ss.ss_family) {
209
+ case AF_INET:
210
+ key[keylen] = ':';
211
+ memmove(key + keylen + 1, port, portlen + 1);
212
+ break;
213
+ case AF_INET6:
214
+ key[keylen] = ']';
215
+ key[keylen + 1] = ':';
216
+ memmove(key + keylen + 2, port, portlen + 1);
217
+ keylen++;
218
+ break;
219
+ default:
220
+ assert(0 && "unsupported address family, could that be IPv7?!");
221
+ }
222
+
223
+ if (st_lookup(table, (st_data_t)key, (st_data_t *)&stats))
224
+ return stats;
225
+
226
+ old_key = key;
227
+
228
+ if (r->idiag_state == TCP_ESTABLISHED) {
229
+ int n = snprintf(key, alloca_len, "%s:%u",
230
+ addr_any(ss.ss_family),
231
+ ntohs(r->id.idiag_sport));
232
+ if (n <= 0) {
233
+ fprintf(stderr, "BUG: snprintf: %d\n", n);
234
+ bug_warn();
235
+ }
236
+ if (st_lookup(table, (st_data_t)key, (st_data_t *)&stats))
237
+ return stats;
238
+ if (n <= 0) {
239
+ key = xmalloc(1);
240
+ *key = '\0';
241
+ } else {
242
+ old_key = key;
243
+ key = xmalloc(n + 1);
244
+ memcpy(key, old_key, n + 1);
245
+ }
246
+ } else {
247
+ key = xmalloc(keylen + 1 + portlen + 1);
248
+ memcpy(key, old_key, keylen + 1 + portlen + 1);
249
+ }
250
+ stats = xcalloc(1, sizeof(struct listen_stats));
251
+ st_insert(table, (st_data_t)key, (st_data_t)stats);
252
+ return stats;
253
+ }
254
+
255
+ static void table_incr_active(st_table *table, struct inet_diag_msg *r)
256
+ {
257
+ struct listen_stats *stats = stats_for(table, r);
258
+ ++stats->active;
259
+ }
260
+
261
+ static void table_set_queued(st_table *table, struct inet_diag_msg *r)
262
+ {
263
+ struct listen_stats *stats = stats_for(table, r);
264
+ stats->listener_p = 1;
265
+ stats->queued = r->idiag_rqueue;
118
266
  }
119
267
 
120
268
  /* inner loop of inet_diag, called for every socket returned by netlink */
@@ -127,85 +275,108 @@ static inline void r_acc(struct nogvl_args *args, struct inet_diag_msg *r)
127
275
  */
128
276
  if (r->idiag_inode == 0)
129
277
  return;
130
- if (r->idiag_state == TCP_ESTABLISHED)
131
- args->stats.active++;
132
- else /* if (r->idiag_state == TCP_LISTEN) */
133
- args->stats.queued = r->idiag_rqueue;
278
+ if (r->idiag_state == TCP_ESTABLISHED) {
279
+ if (args->table)
280
+ table_incr_active(args->table, r);
281
+ else
282
+ args->stats.active++;
283
+ } else { /* if (r->idiag_state == TCP_LISTEN) */
284
+ if (args->table)
285
+ table_set_queued(args->table, r);
286
+ else
287
+ args->stats.queued = r->idiag_rqueue;
288
+ }
134
289
  /*
135
290
  * we wont get anything else because of the idiag_states filter
136
291
  */
137
292
  }
138
293
 
139
- static const char err_socket[] = "socket";
140
294
  static const char err_sendmsg[] = "sendmsg";
141
295
  static const char err_recvmsg[] = "recvmsg";
142
296
  static const char err_nlmsg[] = "nlmsg";
143
297
 
298
+ struct diag_req {
299
+ struct nlmsghdr nlh;
300
+ struct inet_diag_req r;
301
+ };
302
+
303
+ static void prep_msghdr(
304
+ struct msghdr *msg,
305
+ struct nogvl_args *args,
306
+ struct sockaddr_nl *nladdr,
307
+ size_t iovlen)
308
+ {
309
+ memset(msg, 0, sizeof(struct msghdr));
310
+ msg->msg_name = (void *)nladdr;
311
+ msg->msg_namelen = sizeof(struct sockaddr_nl);
312
+ msg->msg_iov = args->iov;
313
+ msg->msg_iovlen = iovlen;
314
+ }
315
+
316
+ static void prep_diag_args(
317
+ struct nogvl_args *args,
318
+ struct sockaddr_nl *nladdr,
319
+ struct rtattr *rta,
320
+ struct diag_req *req,
321
+ struct msghdr *msg)
322
+ {
323
+ memset(req, 0, sizeof(struct diag_req));
324
+ memset(nladdr, 0, sizeof(struct sockaddr_nl));
325
+
326
+ nladdr->nl_family = AF_NETLINK;
327
+
328
+ req->nlh.nlmsg_len = sizeof(struct diag_req) +
329
+ RTA_LENGTH(args->iov[2].iov_len);
330
+ req->nlh.nlmsg_type = TCPDIAG_GETSOCK;
331
+ req->nlh.nlmsg_flags = NLM_F_ROOT | NLM_F_MATCH | NLM_F_REQUEST;
332
+ req->nlh.nlmsg_pid = getpid();
333
+ req->r.idiag_states = (1<<TCP_ESTABLISHED) | (1<<TCP_LISTEN);
334
+ rta->rta_type = INET_DIAG_REQ_BYTECODE;
335
+ rta->rta_len = RTA_LENGTH(args->iov[2].iov_len);
336
+
337
+ args->iov[0].iov_base = req;
338
+ args->iov[0].iov_len = sizeof(struct diag_req);
339
+ args->iov[1].iov_base = rta;
340
+ args->iov[1].iov_len = sizeof(struct rtattr);
341
+
342
+ prep_msghdr(msg, args, nladdr, 3);
343
+ }
344
+
345
+ static void prep_recvmsg_buf(struct nogvl_args *args)
346
+ {
347
+ /* reuse buffer that was allocated for bytecode */
348
+ args->iov[0].iov_len = page_size;
349
+ args->iov[0].iov_base = args->iov[2].iov_base;
350
+ }
351
+
144
352
  /* does the inet_diag stuff with netlink(), this is called w/o GVL */
145
353
  static VALUE diag(void *ptr)
146
354
  {
147
355
  struct nogvl_args *args = ptr;
148
356
  struct sockaddr_nl nladdr;
149
357
  struct rtattr rta;
150
- struct {
151
- struct nlmsghdr nlh;
152
- struct inet_diag_req r;
153
- } req;
358
+ struct diag_req req;
154
359
  struct msghdr msg;
155
360
  const char *err = NULL;
156
- unsigned seq = ++g_seq; /* not atomic, rely on GVL for now */
157
- int fd = socket(AF_NETLINK, SOCK_RAW, NETLINK_INET_DIAG);
158
-
159
- if (fd < 0)
160
- return (VALUE)err_socket;
361
+ unsigned seq = ++g_seq;
161
362
 
162
- memset(&args->stats, 0, sizeof(struct listen_stats));
163
-
164
- memset(&nladdr, 0, sizeof(nladdr));
165
- nladdr.nl_family = AF_NETLINK;
166
-
167
- memset(&req, 0, sizeof(req));
168
- req.nlh.nlmsg_len = sizeof(req) + RTA_LENGTH(args->iov[2].iov_len);
169
- req.nlh.nlmsg_type = TCPDIAG_GETSOCK;
170
- req.nlh.nlmsg_flags = NLM_F_ROOT | NLM_F_MATCH | NLM_F_REQUEST;
171
- req.nlh.nlmsg_pid = getpid();
363
+ prep_diag_args(args, &nladdr, &rta, &req, &msg);
172
364
  req.nlh.nlmsg_seq = seq;
173
- req.r.idiag_family = AF_INET;
174
- req.r.idiag_states = (1<<TCP_ESTABLISHED) | (1<<TCP_LISTEN);
175
- rta.rta_type = INET_DIAG_REQ_BYTECODE;
176
- rta.rta_len = RTA_LENGTH(args->iov[2].iov_len);
177
-
178
- args->iov[0].iov_base = &req;
179
- args->iov[0].iov_len = sizeof(req);
180
- args->iov[1].iov_base = &rta;
181
- args->iov[1].iov_len = sizeof(rta);
182
-
183
- memset(&msg, 0, sizeof(msg));
184
- msg.msg_name = (void *)&nladdr;
185
- msg.msg_namelen = sizeof(nladdr);
186
- msg.msg_iov = args->iov;
187
- msg.msg_iovlen = 3;
188
-
189
- if (sendmsg(fd, &msg, 0) < 0) {
365
+
366
+ if (sendmsg(args->fd, &msg, 0) < 0) {
190
367
  err = err_sendmsg;
191
368
  goto out;
192
369
  }
193
370
 
194
- /* reuse buffer that was allocated for bytecode */
195
- args->iov[0].iov_len = page_size;
196
- args->iov[0].iov_base = args->iov[2].iov_base;
371
+ prep_recvmsg_buf(args);
197
372
 
198
373
  while (1) {
199
374
  ssize_t readed;
375
+ size_t r;
200
376
  struct nlmsghdr *h = (struct nlmsghdr *)args->iov[0].iov_base;
201
377
 
202
- memset(&msg, 0, sizeof(msg));
203
- msg.msg_name = (void *)&nladdr;
204
- msg.msg_namelen = sizeof(nladdr);
205
- msg.msg_iov = args->iov;
206
- msg.msg_iovlen = 1;
207
-
208
- readed = recvmsg(fd, &msg, 0);
378
+ prep_msghdr(&msg, args, &nladdr, 1);
379
+ readed = recvmsg(args->fd, &msg, 0);
209
380
  if (readed < 0) {
210
381
  if (errno == EINTR)
211
382
  continue;
@@ -214,8 +385,8 @@ static VALUE diag(void *ptr)
214
385
  }
215
386
  if (readed == 0)
216
387
  goto out;
217
-
218
- for ( ; NLMSG_OK(h, readed); h = NLMSG_NEXT(h, readed)) {
388
+ r = (size_t)readed;
389
+ for ( ; NLMSG_OK(h, r); h = NLMSG_NEXT(h, r)) {
219
390
  if (h->nlmsg_seq != seq)
220
391
  continue;
221
392
  if (h->nlmsg_type == NLMSG_DONE)
@@ -230,33 +401,90 @@ static VALUE diag(void *ptr)
230
401
  out:
231
402
  {
232
403
  int save_errno = errno;
233
- close(fd);
404
+ if (err && args->table) {
405
+ st_foreach(args->table, st_free_data, 0);
406
+ st_free_table(args->table);
407
+ }
234
408
  errno = save_errno;
235
409
  }
236
410
  return (VALUE)err;
237
411
  }
238
412
 
239
- /* populates inet my_addr struct by parsing +addr+ */
240
- static void parse_addr(struct my_addr *inet, VALUE addr)
413
+ /* populates sockaddr_storage struct by parsing +addr+ */
414
+ static void parse_addr(struct sockaddr_storage *inet, VALUE addr)
241
415
  {
242
- char *host_port, *colon;
243
-
244
- if (TYPE(addr) != T_STRING)
245
- rb_raise(rb_eArgError, "addrs must be an Array of Strings");
416
+ char *host_ptr;
417
+ char *check;
418
+ char *colon = NULL;
419
+ char *rbracket = NULL;
420
+ void *dst;
421
+ long host_len;
422
+ int af, rc;
423
+ uint16_t *portdst;
424
+ unsigned long port;
425
+
426
+ Check_Type(addr, T_STRING);
427
+ host_ptr = StringValueCStr(addr);
428
+ host_len = RSTRING_LEN(addr);
429
+ if (*host_ptr == '[') { /* ipv6 address format (rfc2732) */
430
+ struct sockaddr_in6 *in6 = (struct sockaddr_in6 *)inet;
431
+ rbracket = memchr(host_ptr + 1, ']', host_len - 1);
432
+
433
+ if (rbracket == NULL)
434
+ rb_raise(rb_eArgError, "']' not found in IPv6 addr=%s",
435
+ host_ptr);
436
+ if (rbracket[1] != ':')
437
+ rb_raise(rb_eArgError, "':' not found in IPv6 addr=%s",
438
+ host_ptr);
439
+ colon = rbracket + 1;
440
+ host_ptr++;
441
+ *rbracket = 0;
442
+ inet->ss_family = af = AF_INET6;
443
+ dst = &in6->sin6_addr;
444
+ portdst = &in6->sin6_port;
445
+ } else { /* ipv4 */
446
+ struct sockaddr_in *in = (struct sockaddr_in *)inet;
447
+ colon = memchr(host_ptr, ':', host_len);
448
+ inet->ss_family = af = AF_INET;
449
+ dst = &in->sin_addr;
450
+ portdst = &in->sin_port;
451
+ }
246
452
 
247
- host_port = RSTRING_PTR(addr);
248
- colon = memchr(host_port, ':', RSTRING_LEN(addr));
249
453
  if (!colon)
250
- rb_raise(rb_eArgError, "port not found in: `%s'", host_port);
251
-
454
+ rb_raise(rb_eArgError, "port not found in: `%s'", host_ptr);
455
+ port = strtoul(colon + 1, &check, 10);
252
456
  *colon = 0;
253
- inet->addr = inet_addr(host_port);
457
+ rc = inet_pton(af, host_ptr, dst);
254
458
  *colon = ':';
255
- inet->port = htons(my_inet_port(colon + 1));
459
+ if (rbracket) *rbracket = ']';
460
+ if (*check || ((uint16_t)port != port))
461
+ rb_raise(rb_eArgError, "invalid port: %s", colon + 1);
462
+ if (rc != 1)
463
+ rb_raise(rb_eArgError, "inet_pton failed for: `%s' with %d",
464
+ host_ptr, rc);
465
+ *portdst = ntohs((uint16_t)port);
466
+ }
467
+
468
+ /* generates inet_diag bytecode to match all addrs */
469
+ static void gen_bytecode_all(struct iovec *iov)
470
+ {
471
+ struct inet_diag_bc_op *op;
472
+ struct inet_diag_hostcond *cond;
473
+
474
+ /* iov_len was already set and base allocated in a parent function */
475
+ assert(iov->iov_len == OPLEN && iov->iov_base && "iov invalid");
476
+ op = iov->iov_base;
477
+ op->code = INET_DIAG_BC_S_COND;
478
+ op->yes = OPLEN;
479
+ op->no = sizeof(struct inet_diag_bc_op) + OPLEN;
480
+ cond = (struct inet_diag_hostcond *)(op + 1);
481
+ cond->family = AF_UNSPEC;
482
+ cond->port = -1;
483
+ cond->prefix_len = 0;
256
484
  }
257
485
 
258
486
  /* generates inet_diag bytecode to match a single addr */
259
- static void gen_bytecode(struct iovec *iov, struct my_addr *inet)
487
+ static void gen_bytecode(struct iovec *iov, struct sockaddr_storage *inet)
260
488
  {
261
489
  struct inet_diag_bc_op *op;
262
490
  struct inet_diag_hostcond *cond;
@@ -269,47 +497,79 @@ static void gen_bytecode(struct iovec *iov, struct my_addr *inet)
269
497
  op->no = sizeof(struct inet_diag_bc_op) + OPLEN;
270
498
 
271
499
  cond = (struct inet_diag_hostcond *)(op + 1);
272
- cond->family = AF_INET;
273
- cond->port = ntohs(inet->port);
274
- cond->prefix_len = inet->addr == 0 ? 0 : sizeof(in_addr_t) * CHAR_BIT;
275
- *cond->addr = inet->addr;
500
+ cond->family = inet->ss_family;
501
+ switch (inet->ss_family) {
502
+ case AF_INET: {
503
+ struct sockaddr_in *in = (struct sockaddr_in *)inet;
504
+
505
+ cond->port = ntohs(in->sin_port);
506
+ cond->prefix_len = in->sin_addr.s_addr == 0 ? 0 :
507
+ sizeof(in->sin_addr.s_addr) * CHAR_BIT;
508
+ *cond->addr = in->sin_addr.s_addr;
509
+ }
510
+ break;
511
+ case AF_INET6: {
512
+ struct sockaddr_in6 *in6 = (struct sockaddr_in6 *)inet;
513
+
514
+ cond->port = ntohs(in6->sin6_port);
515
+ cond->prefix_len = memcmp(&in6addr_any, &in6->sin6_addr,
516
+ sizeof(struct in6_addr)) == 0 ?
517
+ 0 : sizeof(in6->sin6_addr) * CHAR_BIT;
518
+ memcpy(&cond->addr, &in6->sin6_addr, sizeof(struct in6_addr));
519
+ }
520
+ break;
521
+ default:
522
+ assert(0 && "unsupported address family, could that be IPv7?!");
523
+ }
276
524
  }
277
525
 
278
- static VALUE tcp_stats(struct nogvl_args *args, VALUE addr)
526
+ static void nl_errcheck(VALUE r)
279
527
  {
280
- const char *err;
281
- VALUE verr;
528
+ const char *err = (const char *)r;
282
529
 
283
- parse_addr(&args->query_addr, addr);
284
- gen_bytecode(&args->iov[2], &args->query_addr);
285
-
286
- verr = rb_thread_blocking_region(diag, args, RUBY_UBF_IO, 0);
287
- err = (const char *)verr;
288
530
  if (err) {
289
531
  if (err == err_nlmsg)
290
532
  rb_raise(rb_eRuntimeError, "NLMSG_ERROR");
291
533
  else
292
534
  rb_sys_fail(err);
293
535
  }
536
+ }
537
+
538
+ static VALUE tcp_stats(struct nogvl_args *args, VALUE addr)
539
+ {
540
+ struct sockaddr_storage query_addr;
541
+
542
+ parse_addr(&query_addr, addr);
543
+ gen_bytecode(&args->iov[2], &query_addr);
544
+
545
+ memset(&args->stats, 0, sizeof(struct listen_stats));
546
+ nl_errcheck(rb_thread_blocking_region(diag, args, 0, 0));
294
547
 
295
548
  return rb_listen_stats(&args->stats);
296
549
  }
297
550
 
298
551
  /*
299
552
  * call-seq:
553
+ * Raindrops::Linux.tcp_listener_stats([addrs[, sock]]) => hash
554
+ *
555
+ * If specified, +addr+ may be a string or array of strings representing
556
+ * listen addresses to filter for. Returns a hash with given addresses as
557
+ * keys and ListenStats objects as the values or a hash of all addresses.
558
+ *
300
559
  * addrs = %w(0.0.0.0:80 127.0.0.1:8080)
301
- * Raindrops::Linux.tcp_listener_stats(addrs) => hash
302
560
  *
303
- * Takes an array of strings representing listen addresses to filter for.
304
- * Returns a hash with given addresses as keys and ListenStats
305
- * objects as the values.
561
+ * If +addr+ is nil or not specified, all (IPv4) addresses are returned.
562
+ * If +sock+ is specified, it should be a Raindrops::InetDiagSock object.
306
563
  */
307
- static VALUE tcp_listener_stats(VALUE obj, VALUE addrs)
564
+ static VALUE tcp_listener_stats(int argc, VALUE *argv, VALUE self)
308
565
  {
309
566
  VALUE *ary;
310
567
  long i;
311
- VALUE rv;
568
+ VALUE rv = rb_hash_new();
312
569
  struct nogvl_args args;
570
+ VALUE addrs, sock;
571
+
572
+ rb_scan_args(argc, argv, "02", &addrs, &sock);
313
573
 
314
574
  /*
315
575
  * allocating page_size instead of OP_LEN since we'll reuse the
@@ -318,15 +578,45 @@ static VALUE tcp_listener_stats(VALUE obj, VALUE addrs)
318
578
  */
319
579
  args.iov[2].iov_len = OPLEN;
320
580
  args.iov[2].iov_base = alloca(page_size);
581
+ args.table = NULL;
582
+ if (NIL_P(sock))
583
+ sock = rb_funcall(cIDSock, id_new, 0);
584
+ args.fd = my_fileno(sock);
585
+
586
+ switch (TYPE(addrs)) {
587
+ case T_STRING:
588
+ rb_hash_aset(rv, addrs, tcp_stats(&args, addrs));
589
+ return rv;
590
+ case T_ARRAY:
591
+ ary = RARRAY_PTR(addrs);
592
+ i = RARRAY_LEN(addrs);
593
+ if (i == 1) {
594
+ rb_hash_aset(rv, *ary, tcp_stats(&args, *ary));
595
+ return rv;
596
+ }
597
+ for (; --i >= 0; ary++) {
598
+ struct sockaddr_storage check;
599
+
600
+ parse_addr(&check, *ary);
601
+ rb_hash_aset(rv, *ary, Qtrue);
602
+ }
603
+ /* fall through */
604
+ case T_NIL:
605
+ args.table = st_init_strtable();
606
+ gen_bytecode_all(&args.iov[2]);
607
+ break;
608
+ default:
609
+ rb_raise(rb_eArgError,
610
+ "addr must be an array of strings, a string, or nil");
611
+ }
321
612
 
322
- if (TYPE(addrs) != T_ARRAY)
323
- rb_raise(rb_eArgError, "addrs must be an Array or String");
613
+ nl_errcheck(rb_thread_blocking_region(diag, &args, NULL, 0));
324
614
 
325
- rv = rb_hash_new();
326
- ary = RARRAY_PTR(addrs);
327
- for (i = RARRAY_LEN(addrs); --i >= 0; ary++)
328
- rb_hash_aset(rv, *ary, tcp_stats(&args, *ary));
615
+ st_foreach(args.table, NIL_P(addrs) ? st_to_hash : st_AND_hash, rv);
616
+ st_free_table(args.table);
329
617
 
618
+ /* let GC deal with corner cases */
619
+ if (argc < 2) rb_io_close(sock);
330
620
  return rv;
331
621
  }
332
622
 
@@ -335,15 +625,23 @@ void Init_raindrops_linux_inet_diag(void)
335
625
  VALUE cRaindrops = rb_const_get(rb_cObject, rb_intern("Raindrops"));
336
626
  VALUE mLinux = rb_define_module_under(cRaindrops, "Linux");
337
627
 
628
+ rb_require("socket");
629
+ cIDSock = rb_const_get(rb_cObject, rb_intern("Socket"));
630
+ id_new = rb_intern("new");
631
+
632
+ /*
633
+ * Document-class: Raindrops::InetDiagSocket
634
+ *
635
+ * This is a subclass of +Socket+ specifically for talking
636
+ * to the inet_diag facility of Netlink.
637
+ */
638
+ cIDSock = rb_define_class_under(cRaindrops, "InetDiagSocket", cIDSock);
639
+ rb_define_singleton_method(cIDSock, "new", ids_s_new, 0);
640
+
338
641
  cListenStats = rb_const_get(cRaindrops, rb_intern("ListenStats"));
339
642
 
340
643
  rb_define_module_function(mLinux, "tcp_listener_stats",
341
- tcp_listener_stats, 1);
342
-
343
- #ifndef HAVE_RB_STRUCT_ALLOC_NOINIT
344
- id_new = rb_intern("new");
345
- #endif
346
- rb_require("raindrops/linux");
644
+ tcp_listener_stats, -1);
347
645
 
348
646
  page_size = getpagesize();
349
647