Fix capture cleanup ownership and manual conflict recovery

CHANGELOG.md

@@ -6,6 +6,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+- Allow manual profiling to recover after a conflicting capture rejects its start.
+- Stop captures once and retain ownership until results are collected.
+
 - Drop support for Ruby 3.2 and Rails 5.2 through 7.1.
 - Emit profiling response headers in lowercase for Rack 3 compatibility.
 

lib/app_profiler.rb

@@ -106,8 +106,7 @@ def start(*args, backend: nil, **kwargs)
     end
 
     def stop
-      profiler.stop
-      profiler.results.tap { clear }
+      profiler.stop_and_results { discard_profiler }
     end
 
     def running?
@@ -257,6 +256,10 @@ def profiler_backend
 
     def clear
       profiler.stop if profiler_backend.locked?
+      discard_profiler
+    end
+
+    def discard_profiler
       @profiler = nil
       @profiler_backend = nil
     end

lib/app_profiler/backend/base_backend.rb

@@ -13,8 +13,8 @@ def run(params = {})
 
         duration = Process.clock_gettime(Process::CLOCK_MONOTONIC) - start_time
 
-        stop
-        results_data = results
+        started = false
+        results_data = stop_and_results
 
         if results_data
           results_data.metadata[:duration] = duration
@@ -38,6 +38,15 @@ def results
         raise NotImplementedError
       end
 
+      def stop_and_results
+        @collecting_results_thread = Thread.current
+        stop
+        results.tap { yield if block_given? }
+      ensure
+        @collecting_results_thread = nil
+        release_run_lock
+      end
+
       def running?
         raise NotImplementedError
       end
@@ -63,6 +72,9 @@ def acquire_run_lock
       end
 
       def release_run_lock
+        # Keep capture ownership through result collection, including stop overrides.
+        return if @collecting_results_thread == Thread.current
+
         self.class.run_lock.unlock if self.class.run_lock.locked?
       rescue ThreadError
         AppProfiler.logger.warn("[AppProfiler] run lock not released as it was never acquired")

lib/app_profiler/server.rb

@@ -156,9 +156,7 @@ def start_running(stackprof_args)
         @semaphore.synchronize do
           return false if @profile_running
 
-          @profile_running = true
-
-          AppProfiler.start(**stackprof_args)
+          @profile_running = !!AppProfiler.start(**stackprof_args)
         end
       end
 

test/app_profiler/backend/stackprof_backend_test.rb

@@ -61,6 +61,59 @@ def teardown
         assert_equal(1000, profile[:interval])
       end
 
+      test ".run stops its capture only once" do
+        stop = StackProf.method(:stop)
+        StackProf.expects(:stop).once.with do
+          stop.call
+          true
+        end
+
+        AppProfiler.profiler.run(stackprof_profile) { :completed }
+      end
+
+      test ".run releases the lock when result collection raises" do
+        backend = StackprofBackend.new
+        backend.expects(:results).raises(RuntimeError, "collection failed")
+
+        error = assert_raises(RuntimeError) { backend.run(stackprof_profile) { :completed } }
+
+        assert_equal("collection failed", error.message)
+        refute_predicate(StackprofBackend, :locked?)
+        refute_predicate(backend, :running?)
+      ensure
+        StackProf.results
+      end
+
+      test ".run honors backend stop overrides" do
+        backend = StackprofBackend.new
+        stop = backend.method(:stop)
+        backend.expects(:stop).once.with do
+          stop.call
+          true
+        end
+
+        assert_instance_of(StackprofProfile, backend.run(stackprof_profile) { :completed })
+      ensure
+        StackProf.stop
+        StackProf.results
+      end
+
+      test ".run keeps the shared lock until results are collected" do
+        competitor = StackprofBackend.new
+        StackProf.expects(:results).twice.with do
+          assert_predicate(StackprofBackend, :locked?)
+          true
+        end.returns(nil, stackprof_profile)
+
+        assert_instance_of(StackprofProfile, AppProfiler.profiler.run(stackprof_profile) { :completed })
+        StackProf.unstub(:results)
+        assert(competitor.start(stackprof_profile))
+      ensure
+        StackProf.unstub(:results)
+        competitor.stop
+        competitor.results
+      end
+
       test ".run assigns metadata to profiles" do
         profile = AppProfiler.profiler.run(stackprof_profile(metadata: { id: "wowza", context: "bar" })) do
           sleep(0.1)
@@ -161,6 +214,44 @@ def teardown
         AppProfiler.stop
       end
 
+      test ".stop discards its cached profiler before a later capture starts" do
+        backend = AppProfiler.profiler
+        release = backend.method(:release_run_lock)
+        replacement = nil
+        params = stackprof_profile
+        AppProfiler.start(params)
+        backend.stubs(:release_run_lock).with do
+          release.call
+          unless StackprofBackend.locked?
+            replacement = AppProfiler.profiler
+            refute_same(backend, replacement)
+            assert(replacement.start(params))
+          end
+          true
+        end
+
+        assert_instance_of(StackprofProfile, AppProfiler.stop)
+        assert_same(replacement, AppProfiler.profiler)
+        assert_predicate(replacement, :running?)
+      ensure
+        replacement&.stop
+        replacement&.results
+      end
+
+      test ".stop keeps the shared lock until results are collected" do
+        AppProfiler.start(stackprof_profile)
+        StackProf.expects(:results).once.with do
+          assert_predicate(StackprofBackend, :locked?)
+          true
+        end.returns(stackprof_profile)
+
+        assert_instance_of(StackprofProfile, AppProfiler.stop)
+        refute_predicate(StackprofBackend, :locked?)
+      ensure
+        StackProf.unstub(:results)
+        AppProfiler.stop
+      end
+
       test ".results prints error when failed" do
         AppProfiler.profiler.expects(:backend_results).returns({})
         AppProfiler.logger.expects(:info).with { |value| value =~ /failed to obtain the profile/ }

test/app_profiler/backend/vernier_backend_test.rb

@@ -61,6 +61,17 @@ def teardown
         assert_equal(1, profile[:meta][:interval])
       end
 
+      test ".run keeps the shared lock until the profile is built" do
+        profile = BaseProfile.from_vernier(vernier_profile(meta: { mode: :wall }))
+        BaseProfile.expects(:from_vernier).with do |_data|
+          assert_predicate(VernierBackend, :locked?)
+          true
+        end.returns(profile)
+
+        assert_same(profile, AppProfiler.profiler.run(vernier_params) { :completed })
+        refute_predicate(VernierBackend, :locked?)
+      end
+
       test ".run assigns metadata to profiles" do
         profile = AppProfiler.profiler.run(
           vernier_params(metadata: {

test/app_profiler/profile_server_test.rb

@@ -356,6 +356,18 @@ def app
         AppProfiler.stop
       end
 
+      test "app can capture after an external profiling conflict clears" do
+        AppProfiler.start
+        get("/profile?duration=0")
+        assert_equal(409, last_response.status)
+
+        AppProfiler.stop
+        get("/profile?duration=0")
+        assert_equal(200, last_response.status)
+      ensure
+        AppProfiler.stop
+      end
+
       private
 
       def with_profiled_workload(workload, &block)