[PATCH] implement SrcLoc::each to walk all allocations

[ew <e@80x24.org>]

From: Eric Wong <e@80x24.org>

Add some cautionary notes around exposing alloc_hdr.
---
 Mwrap.xs  | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++++--
 t/mwrap.t | 22 ++++++++++++++++++++++
 2 files changed, 74 insertions(+), 2 deletions(-)

diff --git a/Mwrap.xs b/Mwrap.xs
index 09d5581..6d3c6d2 100644
--- a/Mwrap.xs
+++ b/Mwrap.xs
@@ -188,7 +188,7 @@ static char *int2str(unsigned num, char *dst, size_t * size)
 	return NULL;
 }
 
-/* allocated via real_malloc/real_free */
+/* allocated via real_malloc, immortal for safety reasons */
 struct src_loc {
 	pthread_mutex_t *mtx;
 	size_t total;
@@ -207,7 +207,10 @@ struct src_loc {
  */
 typedef struct src_loc * Devel__Mwrap__SrcLoc;
 
-/* every allocation has this in the header, maintain alignment with malloc  */
+/*
+ * Every allocation has this in the header, maintain alignment with malloc
+ * Do not expose this to Perl code because of use-after-free concerns.
+ */
 struct alloc_hdr {
 	struct cds_list_head anode; /* <=> src_loc.allocs */
 	union {
@@ -940,3 +943,50 @@ OUTPUT:
 	RETVAL
 CLEANUP:
 	--locating;
+
+void
+src_loc_each(self, min, cb)
+	Devel::Mwrap::SrcLoc self;
+	size_t min;
+	SV *cb;
+PREINIT:
+	struct alloc_hdr *h;
+	bool err = false;
+CODE:
+	++locating;
+	rcu_read_lock();
+	cds_list_for_each_entry_rcu(h, &self->allocs, anode) {
+		size_t size = uatomic_read(&h->size);
+
+		if (size > min) {
+			dSP;
+			ENTER;
+			SAVETMPS;
+
+			PUSHMARK(SP);
+
+			/*
+			 * note: we MUST NOT expose alloc_hdr to Perl code
+			 * since that opens us up to use-after-free
+			 */
+			XPUSHs(sv_2mortal(newSVuv(size)));
+			PUTBACK;
+
+			call_sv(cb, G_DISCARD|G_EVAL);
+
+			SPAGAIN;
+			if (SvTRUE(ERRSV))
+				err = true;
+
+			FREETMPS;
+			LEAVE;
+			if (err)
+				break;
+			assert(rcu_read_ongoing());
+		}
+	}
+	if (SvTRUE(ERRSV))
+		croak(NULL);
+CLEANUP:
+	rcu_read_unlock();
+	--locating;
diff --git a/t/mwrap.t b/t/mwrap.t
index 8806847..661a90a 100644
--- a/t/mwrap.t
+++ b/t/mwrap.t
@@ -88,6 +88,28 @@ Devel::Mwrap::each($nbytes, sub { push @keep, @_ });
 scalar(@keep) == 0 or die "::reset did not work";
 EOF
 
+mwrap_run('Devel::Mwrap::SrcLoc::each', {}, '-e', <<'EOF');
+open my $zero, '<', '/dev/zero' or die "open /dev/zero: $!";
+my @keep;
+my $nr = 10;
+my $nbytes = 1024 * 10;
+sub do_read () {
+	sysread($zero, my $buf, $nbytes);
+	# this forces us to allocate a new buf with every call
+	pop @keep;
+	push @keep, $buf;
+}
+for (1..$nr) { do_read() }
+my $loc = Devel::Mwrap::get('-e:6');
+$loc && $loc->total >= ($nbytes * $nr) or die "wrong line or bad stats";
+my @sl_each;
+$loc->each($nbytes, sub { push @sl_each, \@_ });
+my $n = scalar(@sl_each);
+$n == 1 or die "SrcLoc::each returned unexpected: $n";
+$sl_each[0]->[0] >= $nbytes or die "$sl_each[0]->[0] < $nbytes";
+EOF
+diag slurp($out);
+
 done_testing();
 
 sub slurp {