[PATCH] allow dump_heap: mask via MWRAP env

[Eric Wong <e@80x24.org>]

This is useful for people who do not want to modify existing
Ruby programs.
---
 ext/mwrap/mwrap.c  | 73 +++++++++++++++++++++++++++++++++++++++++-----
 test/test_mwrap.rb |  7 +++++
 2 files changed, 72 insertions(+), 8 deletions(-)

diff --git a/ext/mwrap/mwrap.c b/ext/mwrap/mwrap.c
index 986ca70..138384d 100644
--- a/ext/mwrap/mwrap.c
+++ b/ext/mwrap/mwrap.c
@@ -345,14 +345,20 @@ acc_mean(const struct acc *acc)
 	return DBL2NUM(acc->nr ? acc->mean : HUGE_VAL);
 }
 
-static VALUE
-acc_stddev(const struct acc *acc)
+static double
+acc_stddev_dbl(const struct acc *acc)
 {
 	if (acc->nr > 1) {
 		double variance = acc->m2 / (acc->nr - 1);
-		DBL2NUM(sqrt(variance));
+		return sqrt(variance);
 	}
-	return INT2NUM(0);
+	return 0.0;
+}
+
+static VALUE
+acc_stddev(const struct acc *acc)
+{
+	return DBL2NUM(acc_stddev_dbl(acc));
 }
 
 static struct src_loc *totals_add_rcu(struct src_loc *k)
@@ -1276,9 +1282,15 @@ static VALUE hpb_stat(int argc, VALUE *argv, VALUE hpb)
  * * dump_fd: a writable FD to dump to
  * * dump_path: a path to dump to, the file is opened in O_APPEND mode
  * * dump_min: the minimum allocation size (total) to dump
+ * * dump_heap: mask of heap_page_body statistics to dump
  *
  * If both `dump_fd' and `dump_path' are specified, dump_path takes
  * precedence.
+ *
+ * dump_heap bitmask
+ * * 0x01 - summary stats (same info as HeapPageBody.stat)
+ * * 0x02 - all live heaps (similar to HeapPageBody.each)
+ * * 0x04 - skip non-heap_page_body-related output
  */
 void Init_mwrap(void)
 {
@@ -1328,6 +1340,44 @@ void Init_mwrap(void)
 	--locating;
 }
 
+enum {
+	DUMP_HPB_STATS = 0x1,
+	DUMP_HPB_EACH = 0x2,
+	DUMP_HPB_EXCL = 0x4,
+};
+
+static void dump_hpb(FILE *fp, unsigned flags)
+{
+	if (flags & DUMP_HPB_STATS) {
+		fprintf(fp,
+			"lifespan_max: %zu\n"
+			"lifespan_min: %zu\n"
+			"lifespan_mean: %0.3f\n"
+			"lifespan_stddev: %0.3f\n"
+			"deathspan_max: %zu\n"
+			"deathspan_min: %zu\n"
+			"deathspan_mean: %0.3f\n"
+			"deathspan_stddev: %0.3f\n",
+			hpb_stats.alive.max,
+			hpb_stats.alive.min,
+			hpb_stats.alive.mean,
+			acc_stddev_dbl(&hpb_stats.alive),
+			hpb_stats.reborn.max,
+			hpb_stats.reborn.min,
+			hpb_stats.reborn.mean,
+			acc_stddev_dbl(&hpb_stats.reborn));
+	}
+	if (flags & DUMP_HPB_EACH) {
+		struct alloc_hdr *h;
+
+		cds_list_for_each_entry(h, &hpb_stats.bodies, anode) {
+			void *addr = hdr2ptr(h);
+
+			fprintf(fp, "%p\t%zu\n", addr, h->as.live.gen);
+		}
+	}
+}
+
 /* rb_cloexec_open isn't usable by non-Ruby processes */
 #ifndef O_CLOEXEC
 #  define O_CLOEXEC 0
@@ -1338,10 +1388,12 @@ static void mwrap_dump_destructor(void)
 {
 	const char *opt = getenv("MWRAP");
 	const char *modes[] = { "a", "a+", "w", "w+", "r+" };
-	struct dump_arg a;
+	struct dump_arg a = { .min = 0 };
 	size_t i;
 	int dump_fd;
+	unsigned dump_heap = 0;
 	char *dump_path;
+	char *s;
 
 	if (!opt)
 		return;
@@ -1368,8 +1420,11 @@ static void mwrap_dump_destructor(void)
 	else if (!sscanf(opt, "dump_fd:%d", &dump_fd))
 		goto out;
 
-	if (!sscanf(opt, "dump_min:%zu", &a.min))
-		a.min = 0;
+	if ((s = strstr(opt, "dump_min:")))
+		sscanf(s, "dump_min:%zu", &a.min);
+
+	if ((s = strstr(opt, "dump_heap:")))
+		sscanf(s, "dump_heap:%u", &dump_heap);
 
 	switch (dump_fd) {
 	case 0: goto out;
@@ -1390,7 +1445,9 @@ static void mwrap_dump_destructor(void)
 		}
 		/* we'll leak some memory here, but this is a destructor */
 	}
-	dump_to_file(&a);
+	if ((dump_heap & DUMP_HPB_EXCL) == 0)
+		dump_to_file(&a);
+	dump_hpb(a.fp, dump_heap);
 out:
 	--locating;
 }
diff --git a/test/test_mwrap.rb b/test/test_mwrap.rb
index 9c6141d..48fba23 100644
--- a/test/test_mwrap.rb
+++ b/test/test_mwrap.rb
@@ -58,6 +58,13 @@ class TestMwrap < Test::Unit::TestCase
       res = system(env, *cmd)
       assert res, $?.inspect
       assert_match(/\b10001\s+1\s+-e:1$/, tmp.read)
+
+      tmp.rewind
+      tmp.truncate(0)
+      env['MWRAP'] = "dump_path:#{tmp.path},dump_heap:5"
+      res = system(env, *cmd)
+      assert res, $?.inspect
+      assert_match %r{lifespan_stddev}, tmp.read
     end
   end
 
-- 
EW