support Mwrap.quiet to temporarily disable allocation tracking

Tracking memory used for monitoring code itself is noise,
so give users the power to omit it.

2 files changed, 42 insertions, 0 deletions

diff --git a/ext/mwrap/mwrap.c b/ext/mwrap/mwrap.c
index 2752598..c1a59e7 100644
--- a/ext/mwrap/mwrap.c
+++ b/ext/mwrap/mwrap.c
@@ -30,6 +30,7 @@ extern void * __attribute__((weak)) ruby_current_execution_context_ptr;
 extern void * __attribute__((weak)) ruby_current_vm_ptr; /* for rb_gc_count */
 extern size_t __attribute__((weak)) rb_gc_count(void);
 extern VALUE __attribute__((weak)) rb_cObject;
+extern VALUE __attribute__((weak)) rb_yield(VALUE);
 
 /* true for glibc/dlmalloc/ptmalloc, not sure about jemalloc */
 #define ASSUMED_MALLOC_ALIGNMENT (sizeof(void *) * 2)
@@ -997,6 +998,26 @@ static VALUE src_loc_name(VALUE self)
         return ret;
 }
 
+static VALUE reset_locating(VALUE ign) { --locating; return Qfalse; }
+
+/*
+ * call-seq:
+ *
+ *        Mwrap.quiet do |depth|
+ *          # expensive sort/calculate/emitting results of Mwrap.each
+ *          # affecting statistics of the rest of the app
+ *        end
+ *
+ * Stops allocation tracking inside the block.  This is useful for
+ * monitoring code which calls other Mwrap (or ObjectSpace/GC)
+ * functions which unavoidably allocate memory.
+ */
+static VALUE mwrap_quiet(VALUE mod)
+{
+        size_t cur = ++locating;
+        return rb_ensure(rb_yield, SIZET2NUM(cur), reset_locating, 0);
+}
+
 /*
  * Document-module: Mwrap
  *
@@ -1040,6 +1061,7 @@ void Init_mwrap(void)
         rb_define_singleton_method(mod, "clear", mwrap_clear, 0);
         rb_define_singleton_method(mod, "each", mwrap_each, -1);
         rb_define_singleton_method(mod, "[]", mwrap_aref, 1);
+        rb_define_singleton_method(mod, "quiet", mwrap_quiet, 0);
         rb_define_method(cSrcLoc, "each", src_loc_each, 0);
         rb_define_method(cSrcLoc, "frees", src_loc_frees, 0);
         rb_define_method(cSrcLoc, "allocations", src_loc_allocations, 0);
diff --git a/test/test_mwrap.rb b/test/test_mwrap.rb
index bc8694e..8425c35 100644
--- a/test/test_mwrap.rb
+++ b/test/test_mwrap.rb
@@ -252,4 +252,24 @@ class TestMwrap < Test::Unit::TestCase
       loc.name.frozen? or abort 'SourceLocation#name not frozen'
     end;
   end
+
+  def test_quiet
+    assert_separately(+"#{<<~"begin;"}\n#{<<~'end;'}")
+    begin;
+      require 'mwrap'
+      before = __LINE__
+      res = Mwrap.quiet do |depth|
+        depth == 1 or abort 'depth is not 1'
+        ('a' * 10000).clear
+        Mwrap.quiet { |d| d == 2 or abort 'depth is not 2' }
+        :foo
+      end
+      after = __LINE__ - 1
+      (before..after).each do |lineno|
+        Mwrap["#{__FILE__}:#{lineno}"] and
+          abort "unexpectedly tracked allocation at line #{lineno}"
+      end
+      res == :foo or abort 'Mwrap.quiet did not return block result'
+    end;
+  end
 end