Start process_queue_thread when dead

Shopify · Dec 23, 2024 · d44d9bb · d44d9bb
1 parent 7bfbe6e
commit d44d9bb
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 5 deletions.
diff --git a/lib/app_profiler/storage/google_cloud_storage.rb b/lib/app_profiler/storage/google_cloud_storage.rb
@@ -37,7 +37,7 @@ def upload(profile, _params = {})
 
         def enqueue_upload(profile)
           mutex.synchronize do
-            process_queue_thread unless @process_queue_thread&.alive?
+            start_process_queue_thread
 
             @queue ||= init_queue
             begin
@@ -60,8 +60,10 @@ def init_queue
           @queue = SizedQueue.new(AppProfiler.upload_queue_max_length)
         end
 
-        def process_queue_thread
-          @process_queue_thread ||= Thread.new do
+        def start_process_queue_thread
+          return if @process_queue_thread&.alive?
+
+          @process_queue_thread = Thread.new do
             loop do
               process_queue
               sleep(AppProfiler.upload_queue_interval_secs)

diff --git a/test/app_profiler/storage/google_cloud_storage_test.rb b/test/app_profiler/storage/google_cloud_storage_test.rb
@@ -119,6 +119,22 @@ class GoogleCloudStorageTest < AppProfiler::TestCase
         assert(th.alive?)
       end
 
+      test "process_queue_thread is recreated on enqueue if stopped" do
+        GoogleCloudStorage.enqueue_upload(profile_from_stackprof)
+        th = GoogleCloudStorage.instance_variable_get(:@process_queue_thread)
+        assert(th.alive?)
+
+        th.kill
+        th.join
+
+        GoogleCloudStorage.enqueue_upload(profile_from_stackprof)
+        assert_not_predicate(th, :alive?)
+
+        GoogleCloudStorage.send(:start_process_queue_thread)
+        th = GoogleCloudStorage.instance_variable_get(:@process_queue_thread)
+        assert(th.alive?)
+      end
+
       private
 
       def with_forward_metadata_on_upload
@@ -131,10 +147,10 @@ def with_forward_metadata_on_upload
 
       def with_stubbed_process_queue_thread
         # Stub out the thread creation, so that tests are not flaky.
-        GoogleCloudStorage.stubs(:process_queue_thread)
+        GoogleCloudStorage.stubs(:start_process_queue_thread)
         yield
       ensure
-        GoogleCloudStorage.unstub(:process_queue_thread)
+        GoogleCloudStorage.unstub(:start_process_queue_thread)
       end
 
       def profile_from_stackprof