Merge pull request tmm1#128 from imjching-shopify/support-custom-meta…

…data

Support custom metadata option in stackprof

README.md

@@ -330,6 +330,7 @@ Option      | Meaning
 `interval`  | mode-relative sample rate [c.f.](#sampling)
 `aggregate` | defaults: `true` - if `false` disables [aggregation](#aggregation)
 `raw`       | defaults `false` - if `true` collects the extra data required by the `--flamegraph` and `--stackcollapse` report types
+`metadata`  | defaults to `{}`. Must be a `Hash`. metadata associated with this profile
 `save_every`| (rack middleware only) write the target file after this many requests
 
 ## Todo

ext/stackprof/stackprof.c

@@ -33,6 +33,7 @@ static struct {
     VALUE mode;
     VALUE interval;
     VALUE out;
+    VALUE metadata;
 
     VALUE *raw_samples;
     size_t raw_samples_len;
@@ -59,7 +60,7 @@ static struct {
 
 static VALUE sym_object, sym_wall, sym_cpu, sym_custom, sym_name, sym_file, sym_line;
 static VALUE sym_samples, sym_total_samples, sym_missed_samples, sym_edges, sym_lines;
-static VALUE sym_version, sym_mode, sym_interval, sym_raw, sym_frames, sym_out, sym_aggregate, sym_raw_timestamp_deltas;
+static VALUE sym_version, sym_mode, sym_interval, sym_raw, sym_metadata, sym_frames, sym_out, sym_aggregate, sym_raw_timestamp_deltas;
 static VALUE sym_gc_samples, objtracer;
 static VALUE gc_hook;
 static VALUE rb_mStackProf;
@@ -72,7 +73,7 @@ stackprof_start(int argc, VALUE *argv, VALUE self)
 {
     struct sigaction sa;
     struct itimerval timer;
-    VALUE opts = Qnil, mode = Qnil, interval = Qnil, out = Qfalse;
+    VALUE opts = Qnil, mode = Qnil, interval = Qnil, metadata = rb_hash_new(), out = Qfalse;
     int raw = 0, aggregate = 1;
 
     if (_stackprof.running)
@@ -85,6 +86,14 @@ stackprof_start(int argc, VALUE *argv, VALUE self)
 	interval = rb_hash_aref(opts, sym_interval);
 	out = rb_hash_aref(opts, sym_out);
 
+	VALUE metadata_val = rb_hash_aref(opts, sym_metadata);
+	if (RTEST(metadata_val)) {
+	    if (!RB_TYPE_P(metadata_val, T_HASH))
+		rb_raise(rb_eArgError, "metadata should be a hash");
+
+	    metadata = metadata_val;
+	}
+
 	if (RTEST(rb_hash_aref(opts, sym_raw)))
 	    raw = 1;
 	if (rb_hash_lookup2(opts, sym_aggregate, Qundef) == Qfalse)
@@ -128,6 +137,7 @@ stackprof_start(int argc, VALUE *argv, VALUE self)
     _stackprof.aggregate = aggregate;
     _stackprof.mode = mode;
     _stackprof.interval = interval;
+    _stackprof.metadata = metadata;
     _stackprof.out = out;
 
     if (raw) {
@@ -258,6 +268,7 @@ stackprof_results(int argc, VALUE *argv, VALUE self)
     rb_hash_aset(results, sym_samples, SIZET2NUM(_stackprof.overall_samples));
     rb_hash_aset(results, sym_gc_samples, SIZET2NUM(_stackprof.during_gc));
     rb_hash_aset(results, sym_missed_samples, SIZET2NUM(_stackprof.overall_signals - _stackprof.overall_samples));
+    rb_hash_aset(results, sym_metadata, _stackprof.metadata);
 
     frames = rb_hash_new();
     rb_hash_aset(results, sym_frames, frames);
@@ -662,6 +673,7 @@ Init_stackprof(void)
     S(raw);
     S(raw_timestamp_deltas);
     S(out);
+    S(metadata);
     S(frames);
     S(aggregate);
 #undef S

lib/stackprof/middleware.rb

@@ -13,12 +13,18 @@ def initialize(app, options = {})
       Middleware.enabled  = options[:enabled]
       options[:path]      = 'tmp/' if options[:path].to_s.empty?
       Middleware.path     = options[:path]
+      Middleware.metadata = options[:metadata] || {}
       at_exit{ Middleware.save } if options[:save_at_exit]
     end
 
     def call(env)
       enabled = Middleware.enabled?(env)
-      StackProf.start(mode: Middleware.mode, interval: Middleware.interval, raw: Middleware.raw) if enabled
+      StackProf.start(
+        mode:     Middleware.mode,
+        interval: Middleware.interval,
+        raw:      Middleware.raw,
+        metadata: Middleware.metadata,
+      ) if enabled
       @app.call(env)
     ensure
       if enabled
@@ -31,7 +37,7 @@ def call(env)
     end
 
     class << self
-      attr_accessor :enabled, :mode, :interval, :raw, :path
+      attr_accessor :enabled, :mode, :interval, :raw, :path, :metadata
 
       def enabled?(env)
         if enabled.respond_to?(:call)

test/test_middleware.rb

@@ -64,4 +64,10 @@ def test_raw
     StackProf::Middleware.new(Object.new, raw: true)
     assert StackProf::Middleware.raw
   end
+
+  def test_metadata
+    metadata = { key: 'value' }
+    StackProf::Middleware.new(Object.new, metadata: metadata)
+    assert_equal metadata, StackProf::Middleware.metadata
+  end
 end

test/test_stackprof.rb

@@ -109,6 +109,36 @@ def test_raw
     assert_equal 10, profile[:raw_timestamp_deltas].size
   end
 
+  def test_metadata
+    metadata = {
+      path: '/foo/bar',
+      revision: '5c0b01f1522ae8c194510977ae29377296dd236b',
+    }
+    profile = StackProf.run(mode: :cpu, metadata: metadata) do
+      math
+    end
+
+    assert_equal metadata, profile[:metadata]
+  end
+
+  def test_empty_metadata
+    profile = StackProf.run(mode: :cpu) do
+      math
+    end
+
+    assert_equal({}, profile[:metadata])
+  end
+
+  def test_raises_if_metadata_is_not_a_hash
+    exception = assert_raises ArgumentError do
+      StackProf.run(mode: :cpu, metadata: 'foobar') do
+        math
+      end
+    end
+
+    assert_equal 'metadata should be a hash', exception.message
+  end
+
   def test_fork
     StackProf.run do
       pid = fork do