Make Coverage suspendable (#4856)

* Make Coverage suspendable

Add `Coverage.suspend`, `Coverage.resume` and some methods.

[Feature #18176] [ruby-core:105321]

6 files changed, 305 insertions, 16 deletions

diff --git a/ext/coverage/coverage.c b/ext/coverage/coverage.c
index 0af5579ffc..f948f62307 100644
--- a/ext/coverage/coverage.c
+++ b/ext/coverage/coverage.c
@@ -15,21 +15,38 @@
 #include "ruby.h"
 #include "vm_core.h"
 
+static enum {
+    IDLE,
+    SUSPENDED,
+    RUNNING
+} current_state = IDLE;
 static int current_mode;
 static VALUE me2counter = Qnil;
 
 /*
  * call-seq:
- *    Coverage.start  => nil
+ *    Coverage.setup                                              => nil
+ *    Coverage.setup(:all)                                        => nil
+ *    Coverage.setup(lines: bool, branches: bool, methods: bool)  => nil
+ *    Coverage.setup(oneshot_lines: true)                         => nil
  *
- * Enables coverage measurement.
+ * Set up the coverage measurement.
+ *
+ * Note that this method does not start the measurement itself.
+ * Use Coverage.resume to start the measurement.
+ *
+ * You may want to use Coverage.start to setup and then start the measurement.
  */
 static VALUE
-rb_coverage_start(int argc, VALUE *argv, VALUE klass)
+rb_coverage_setup(int argc, VALUE *argv, VALUE klass)
 {
     VALUE coverages, opt;
     int mode;
 
+    if (current_state != IDLE) {
+	rb_raise(rb_eRuntimeError, "coverage measurement is already setup");
+    }
+
     rb_scan_args(argc, argv, "01", &opt);
 
     if (argc == 0) {
@@ -70,10 +87,57 @@ rb_coverage_start(int argc, VALUE *argv, VALUE klass)
 	current_mode = mode;
 	if (mode == 0) mode = COVERAGE_TARGET_LINES;
 	rb_set_coverages(coverages, mode, me2counter);
+        current_state = SUSPENDED;
     }
     else if (current_mode != mode) {
 	rb_raise(rb_eRuntimeError, "cannot change the measuring target during coverage measurement");
     }
+
+
+    return Qnil;
+}
+
+/*
+ * call-seq:
+ *    Coverage.resume  => nil
+ *
+ * Start/resume the coverage measurement.
+ *
+ * Caveat: Currently, only process-global coverage measurement is supported.
+ * You cannot measure per-thread covearge. If your process has multiple thread,
+ * using Coverage.resume/suspend to capture code coverage executed from only
+ * a limited code block, may yield misleading results.
+ */
+VALUE
+rb_coverage_resume(VALUE klass)
+{
+    if (current_state == IDLE) {
+	rb_raise(rb_eRuntimeError, "coverage measurement is not set up yet");
+    }
+    if (current_state == RUNNING) {
+	rb_raise(rb_eRuntimeError, "coverage measurement is already running");
+    }
+    rb_resume_coverages();
+    current_state = RUNNING;
+    return Qnil;
+}
+
+/*
+ * call-seq:
+ *    Coverage.start                                              => nil
+ *    Coverage.start(:all)                                        => nil
+ *    Coverage.start(lines: bool, branches: bool, methods: bool)  => nil
+ *    Coverage.start(oneshot_lines: true)                         => nil
+ *
+ * Enables the coverage measurement.
+ * See the documentation of Coverage class in detail.
+ * This is equivalent to Coverage.setup and Coverage.resume.
+ */
+static VALUE
+rb_coverage_start(int argc, VALUE *argv, VALUE klass)
+{
+    rb_coverage_setup(argc, argv, klass);
+    rb_coverage_resume(klass);
     return Qnil;
 }
 
@@ -280,6 +344,24 @@ clear_me2counter_i(VALUE key, VALUE value, VALUE unused)
 }
 
 /*
+ * call-seq:
+ *    Coverage.suspend  => nil
+ *
+ * Suspend the coverage measurement.
+ * You can use Coverage.resumt to restart the measurement.
+ */
+VALUE
+rb_coverage_suspend(VALUE klass)
+{
+    if (current_state != RUNNING) {
+	rb_raise(rb_eRuntimeError, "coverage measurement is not running");
+    }
+    rb_suspend_coverages();
+    current_state = SUSPENDED;
+    return Qnil;
+}
+
+/*
  *  call-seq:
  *     Coverage.result(stop: true, clear: true)  => hash
  *
@@ -294,6 +376,10 @@ rb_coverage_result(int argc, VALUE *argv, VALUE klass)
     VALUE opt;
     int stop = 1, clear = 1;
 
+    if (current_state == IDLE) {
+	rb_raise(rb_eRuntimeError, "coverage measurement is not enabled");
+    }
+
     rb_scan_args(argc, argv, "01", &opt);
 
     if (argc == 1) {
@@ -312,8 +398,12 @@ rb_coverage_result(int argc, VALUE *argv, VALUE klass)
         if (!NIL_P(me2counter)) rb_hash_foreach(me2counter, clear_me2counter_i, Qnil);
     }
     if (stop) {
+        if (current_state == RUNNING) {
+            rb_coverage_suspend(klass);
+        }
         rb_reset_coverages();
         me2counter = Qnil;
+        current_state = IDLE;
     }
     return ncoverages;
 }
@@ -321,6 +411,23 @@ rb_coverage_result(int argc, VALUE *argv, VALUE klass)
 
 /*
  *  call-seq:
+ *     Coverage.state  => :idle, :suspended, :running
+ *
+ * Returns the state of the coverage measurement.
+ */
+static VALUE
+rb_coverage_state(VALUE klass)
+{
+    switch (current_state) {
+        case IDLE: return ID2SYM(rb_intern("idle"));
+        case SUSPENDED: return ID2SYM(rb_intern("suspended"));
+        case RUNNING: return ID2SYM(rb_intern("running"));
+    }
+    return Qnil;
+}
+
+/*
+ *  call-seq:
  *     Coverage.running?  => bool
  *
  * Returns true if coverage stats are currently being collected (after
@@ -329,13 +436,15 @@ rb_coverage_result(int argc, VALUE *argv, VALUE klass)
 static VALUE
 rb_coverage_running(VALUE klass)
 {
-    VALUE coverages = rb_get_coverages();
-    return RTEST(coverages) ? Qtrue : Qfalse;
+    return current_state == RUNNING ? Qtrue : Qfalse;
 }
 
 /* Coverage provides coverage measurement feature for Ruby.
  * This feature is experimental, so these APIs may be changed in future.
  *
+ * Caveat: Currently, only process-global coverage measurement is supported.
+ * You cannot measure per-thread covearge.
+ *
  * = Usage
  *
  * 1. require "coverage"
@@ -480,9 +589,13 @@ void
 Init_coverage(void)
 {
     VALUE rb_mCoverage = rb_define_module("Coverage");
+    rb_define_module_function(rb_mCoverage, "setup", rb_coverage_setup, -1);
     rb_define_module_function(rb_mCoverage, "start", rb_coverage_start, -1);
+    rb_define_module_function(rb_mCoverage, "resume", rb_coverage_resume, 0);
+    rb_define_module_function(rb_mCoverage, "suspend", rb_coverage_suspend, 0);
     rb_define_module_function(rb_mCoverage, "result", rb_coverage_result, -1);
     rb_define_module_function(rb_mCoverage, "peek_result", rb_coverage_peek_result, 0);
+    rb_define_module_function(rb_mCoverage, "state", rb_coverage_state, 0);
     rb_define_module_function(rb_mCoverage, "running?", rb_coverage_running, 0);
     rb_global_variable(&me2counter);
 }
diff --git a/spec/ruby/library/coverage/result_spec.rb b/spec/ruby/library/coverage/result_spec.rb
index 6cf5be1346..4cc43e8462 100644
--- a/spec/ruby/library/coverage/result_spec.rb
+++ b/spec/ruby/library/coverage/result_spec.rb
@@ -65,12 +65,25 @@ describe 'Coverage.result' do
     result.should == {}
   end
 
-  it 'second Coverage.start does nothing' do
-    Coverage.start
-    require @config_file.chomp('.rb')
-    result = Coverage.result
+  ruby_version_is ''...'3.1' do
+    it 'second Coverage.start does nothing' do
+      Coverage.start
+      require @config_file.chomp('.rb')
+      result = Coverage.result
+
+      result.should == { @config_file => [1, 1, 1] }
+    end
+  end
 
-    result.should == { @config_file => [1, 1, 1] }
+  ruby_version_is '3.1' do
+    it 'second Coverage.start give exception' do
+      Coverage.start
+      -> {
+        require @config_file.chomp('.rb')
+      }.should raise_error(RuntimeError, 'coverage measurement is already setup')
+    ensure
+      Coverage.result
+    end
   end
 
   it 'does not include the file starting coverage since it is not tracked' do
diff --git a/test/coverage/test_coverage.rb b/test/coverage/test_coverage.rb
index 22557bd9d8..882368363a 100644
--- a/test/coverage/test_coverage.rb
+++ b/test/coverage/test_coverage.rb
@@ -774,4 +774,150 @@ class TestCoverage < Test::Unit::TestCase
       end
     end;
   end
+
+  def test_coverage_suspendable
+    Dir.mktmpdir {|tmp|
+      Dir.chdir(tmp) {
+        File.open("test.rb", "w") do |f|
+          f.puts <<-EOS
+            def foo
+              :ok
+            end
+
+            def bar
+              :ok
+            end
+
+            def baz
+              :ok
+            end
+          EOS
+        end
+
+        cov1 = "[0, 0, nil, nil, 0, 1, nil, nil, 0, 0, nil]"
+        cov2 = "[0, 0, nil, nil, 0, 1, nil, nil, 0, 1, nil]"
+        assert_in_out_err(%w[-rcoverage], <<-"end;", [cov1, cov2], [])
+          Coverage.setup
+          tmp = Dir.pwd
+          require tmp + "/test.rb"
+          foo
+          Coverage.resume
+          bar
+          Coverage.suspend
+          baz
+          p Coverage.peek_result[tmp + "/test.rb"]
+          Coverage.resume
+          baz
+          p Coverage.result[tmp + "/test.rb"]
+        end;
+
+        cov1 = "{:lines=>[0, 0, nil, nil, 0, 1, nil, nil, 0, 0, nil], :branches=>{}, :methods=>{[Object, :baz, 9, 12, 11, 15]=>0, [Object, :bar, 5, 12, 7, 15]=>1, [Object, :foo, 1, 12, 3, 15]=>0}}"
+        cov2 = "{:lines=>[0, 0, nil, nil, 0, 1, nil, nil, 0, 1, nil], :branches=>{}, :methods=>{[Object, :baz, 9, 12, 11, 15]=>1, [Object, :bar, 5, 12, 7, 15]=>1, [Object, :foo, 1, 12, 3, 15]=>0}}"
+        assert_in_out_err(%w[-rcoverage], <<-"end;", [cov1, cov2], [])
+          Coverage.setup(:all)
+          tmp = Dir.pwd
+          require tmp + "/test.rb"
+          foo
+          Coverage.resume
+          bar
+          Coverage.suspend
+          baz
+          p Coverage.peek_result[tmp + "/test.rb"]
+          Coverage.resume
+          baz
+          p Coverage.result[tmp + "/test.rb"]
+        end;
+
+        cov1 = "{:oneshot_lines=>[6]}"
+        cov2 = "{:oneshot_lines=>[6, 10]}"
+        assert_in_out_err(%w[-rcoverage], <<-"end;", [cov1, cov2], [])
+          Coverage.setup(oneshot_lines: true)
+          tmp = Dir.pwd
+          require tmp + "/test.rb"
+          foo
+          Coverage.resume
+          bar
+          Coverage.suspend
+          baz
+          p Coverage.peek_result[tmp + "/test.rb"]
+          Coverage.resume
+          baz
+          p Coverage.result[tmp + "/test.rb"]
+        end;
+      }
+    }
+  end
+
+  def test_coverage_state
+    assert_in_out_err(%w[-rcoverage], <<-"end;", [":idle", ":running", ":running", ":idle"], [])
+      p Coverage.state
+      Coverage.start
+      p Coverage.state
+      Coverage.peek_result
+      p Coverage.state
+      Coverage.result
+      p Coverage.state
+    end;
+
+    assert_in_out_err(%w[-rcoverage], <<-"end;", [":idle", ":suspended", ":running", ":suspended", ":running", ":suspended", ":idle"], [])
+      p Coverage.state
+      Coverage.setup
+      p Coverage.state
+      Coverage.resume
+      p Coverage.state
+      Coverage.suspend
+      p Coverage.state
+      Coverage.resume
+      p Coverage.state
+      Coverage.suspend
+      p Coverage.state
+      Coverage.result
+      p Coverage.state
+    end;
+  end
+
+  def test_result_without_resume
+    assert_in_out_err(%w[-rcoverage], <<-"end;", ["{}"], [])
+      Coverage.setup
+      p Coverage.result
+    end;
+  end
+
+  def test_result_after_suspend
+    assert_in_out_err(%w[-rcoverage], <<-"end;", ["{}"], [])
+      Coverage.start
+      Coverage.suspend
+      p Coverage.result
+    end;
+  end
+
+  def test_resume_without_setup
+    assert_in_out_err(%w[-rcoverage], <<-"end;", [], /coverage measurement is not set up yet/)
+      Coverage.resume
+      p :NG
+    end;
+  end
+
+  def test_suspend_without_setup
+    assert_in_out_err(%w[-rcoverage], <<-"end;", [], /coverage measurement is not running/)
+      Coverage.suspend
+      p :NG
+    end;
+  end
+
+  def test_double_resume
+    assert_in_out_err(%w[-rcoverage], <<-"end;", [], /coverage measurement is already running/)
+      Coverage.start
+      Coverage.resume
+      p :NG
+    end;
+  end
+
+  def test_double_suspend
+    assert_in_out_err(%w[-rcoverage], <<-"end;", [], /coverage measurement is not running/)
+      Coverage.setup
+      Coverage.suspend
+      p :NG
+    end;
+  end
 end
diff --git a/thread.c b/thread.c
index 336ac7836c..17510b32f2 100644
--- a/thread.c
+++ b/thread.c
@@ -5746,7 +5746,15 @@ void
 rb_set_coverages(VALUE coverages, int mode, VALUE me2counter)
 {
     GET_VM()->coverages = coverages;
+    GET_VM()->me2counter = me2counter;
     GET_VM()->coverage_mode = mode;
+}
+
+void
+rb_resume_coverages()
+{
+    int mode = GET_VM()->coverage_mode;
+    VALUE me2counter = GET_VM()->me2counter;
     rb_add_event_hook2((rb_event_hook_func_t) update_line_coverage, RUBY_EVENT_COVERAGE_LINE, Qnil, RUBY_EVENT_HOOK_FLAG_SAFE | RUBY_EVENT_HOOK_FLAG_RAW_ARG);
     if (mode & COVERAGE_TARGET_BRANCHES) {
 	rb_add_event_hook2((rb_event_hook_func_t) update_branch_coverage, RUBY_EVENT_COVERAGE_BRANCH, Qnil, RUBY_EVENT_HOOK_FLAG_SAFE | RUBY_EVENT_HOOK_FLAG_RAW_ARG);
@@ -5756,13 +5764,9 @@ rb_set_coverages(VALUE coverages, int mode, VALUE me2counter)
     }
 }
 
-/* Make coverage arrays empty so old covered files are no longer tracked. */
 void
-rb_reset_coverages(void)
+rb_suspend_coverages()
 {
-    rb_clear_coverages();
-    rb_iseq_remove_coverage_all();
-    GET_VM()->coverages = Qfalse;
     rb_remove_event_hook((rb_event_hook_func_t) update_line_coverage);
     if (GET_VM()->coverage_mode & COVERAGE_TARGET_BRANCHES) {
 	rb_remove_event_hook((rb_event_hook_func_t) update_branch_coverage);
@@ -5772,6 +5776,15 @@ rb_reset_coverages(void)
     }
 }
 
+/* Make coverage arrays empty so old covered files are no longer tracked. */
+void
+rb_reset_coverages(void)
+{
+    rb_clear_coverages();
+    rb_iseq_remove_coverage_all();
+    GET_VM()->coverages = Qfalse;
+}
+
 VALUE
 rb_default_coverage(int n)
 {
diff --git a/vm.c b/vm.c
index a18b2140d3..7f3376ce68 100644
--- a/vm.c
+++ b/vm.c
@@ -2516,6 +2516,7 @@ rb_vm_update_references(void *ptr)
 
         if (vm->coverages) {
             vm->coverages = rb_gc_location(vm->coverages);
+            vm->me2counter = rb_gc_location(vm->me2counter);
         }
     }
 }
@@ -2602,6 +2603,7 @@ rb_vm_mark(void *ptr)
         rb_gc_mark_movable(vm->top_self);
         rb_gc_mark_movable(vm->orig_progname);
         RUBY_MARK_MOVABLE_UNLESS_NULL(vm->coverages);
+        RUBY_MARK_MOVABLE_UNLESS_NULL(vm->me2counter);
         /* Prevent classes from moving */
         rb_mark_tbl(vm->defined_module_hash);
 
diff --git a/vm_core.h b/vm_core.h
index 76d0025613..a696a1af39 100644
--- a/vm_core.h
+++ b/vm_core.h
@@ -667,7 +667,7 @@ typedef struct rb_vm_struct {
     rb_nativethread_lock_t workqueue_lock;
 
     VALUE orig_progname, progname;
-    VALUE coverages;
+    VALUE coverages, me2counter;
     int coverage_mode;
 
     st_table * defined_module_hash;
@@ -2060,6 +2060,8 @@ extern VALUE rb_get_coverages(void);
 extern void rb_set_coverages(VALUE, int, VALUE);
 extern void rb_clear_coverages(void);
 extern void rb_reset_coverages(void);
+extern void rb_resume_coverages(void);
+extern void rb_suspend_coverages(void);
 
 void rb_postponed_job_flush(rb_vm_t *vm);