Merge pull request #2475 from markiz/ma/thread-local-storage

Thread-local storage for Tracer.state

lib/new_relic/agent.rb

@@ -31,6 +31,7 @@ module Agent
     require 'new_relic/noticed_error'
     require 'new_relic/agent/noticeable_error'
     require 'new_relic/supportability_helper'
+    require 'new_relic/thread_local_storage'
 
     require 'new_relic/agent/encoding_normalizer'
     require 'new_relic/agent/stats'

lib/new_relic/agent/configuration/default_source.rb

@@ -554,6 +554,13 @@ def self.enforce_fallback(allowed_values: nil, fallback: nil)
           :allowed_from_server => false,
           :description => 'When set to `true`, forces a synchronous connection to the New Relic [collector](/docs/using-new-relic/welcome-new-relic/get-started/glossary/#collector) during application startup. For very short-lived processes, this helps ensure the New Relic agent has time to report.'
         },
+        :thread_local_tracer_state => {
+          :default => false,
+          :public => true,
+          :type => Boolean,
+          :allowed_from_server => false,
+          :description => 'If `true`, tracer state storage is thread-local, otherwise, fiber-local'
+        },
         :timeout => {
           :default => 2 * 60, # 2 minutes
           :public => true,

lib/new_relic/agent/threading/agent_thread.rb

@@ -9,8 +9,7 @@ class AgentThread
         def self.create(label, &blk)
           ::NewRelic::Agent.logger.debug("Creating AgentThread: #{label}")
           wrapped_blk = proc do
-            if ::Thread.current[:newrelic_tracer_state] && Thread.current[:newrelic_tracer_state].current_transaction
-              txn = ::Thread.current[:newrelic_tracer_state].current_transaction
+            if (txn = ::NewRelic::ThreadLocalStorage[:newrelic_tracer_state]&.current_transaction)
               ::NewRelic::Agent.logger.warn("AgentThread created with current transaction #{txn.best_name}")
             end
             begin

lib/new_relic/agent/tracer.rb

@@ -392,11 +392,11 @@ def start_message_broker_segment(action:,
         #
         # If ever exposed, this requires additional synchronization
         def state_for(thread)
-          state = thread[:newrelic_tracer_state]
+          state = ThreadLocalStorage.get(thread, :newrelic_tracer_state)
 
           if state.nil?
             state = Tracer::State.new
-            thread[:newrelic_tracer_state] = state
+            ThreadLocalStorage.set(thread, :newrelic_tracer_state, state)
           end
 
           state
@@ -405,7 +405,7 @@ def state_for(thread)
         alias_method :tl_state_for, :state_for
 
         def clear_state
-          Thread.current[:newrelic_tracer_state] = nil
+          ThreadLocalStorage[:newrelic_tracer_state] = nil
         end
 
         alias_method :tl_clear, :clear_state
@@ -420,12 +420,12 @@ def thread_tracing_enabled?
 
         def thread_block_with_current_transaction(segment_name: nil, parent: nil, &block)
           parent ||= current_segment
-          current_txn = ::Thread.current[:newrelic_tracer_state]&.current_transaction if ::Thread.current[:newrelic_tracer_state]&.is_execution_traced?
+          current_txn = ThreadLocalStorage[:newrelic_tracer_state]&.current_transaction if ThreadLocalStorage[:newrelic_tracer_state]&.is_execution_traced?
           proc do |*args|
             begin
               if current_txn && !current_txn.finished?
                 NewRelic::Agent::Tracer.state.current_transaction = current_txn
-                ::Thread.current[:newrelic_thread_span_parent] = parent
+                ThreadLocalStorage[:newrelic_thread_span_parent] = parent
                 current_txn.async = true
                 segment_name = "#{segment_name}/Thread#{::Thread.current.object_id}/Fiber#{::Fiber.current.object_id}" if NewRelic::Agent.config[:'thread_ids_enabled']
                 segment = NewRelic::Agent::Tracer.start_segment(name: segment_name, parent: parent) if segment_name

lib/new_relic/agent/transaction/tracing.rb

@@ -41,11 +41,11 @@ def add_segment(segment, parent = nil)
 
         def thread_starting_span
           # if the previous current segment was in another thread, use the thread local parent
-          if ::Thread.current[:newrelic_thread_span_parent] &&
+          if ThreadLocalStorage[:newrelic_thread_span_parent] &&
               current_segment &&
               current_segment.starting_segment_key != NewRelic::Agent::Tracer.current_segment_key
 
-            ::Thread.current[:newrelic_thread_span_parent]
+            ThreadLocalStorage[:newrelic_thread_span_parent]
           end
         end
 

lib/new_relic/thread_local_storage.rb

@@ -0,0 +1,31 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+module NewRelic
+  module ThreadLocalStorage
+    def self.get(thread, key)
+      if Agent.config[:thread_local_tracer_state]
+        thread.thread_variable_get(key)
+      else
+        thread[key]
+      end
+    end
+
+    def self.set(thread, key, value)
+      if Agent.config[:thread_local_tracer_state]
+        thread.thread_variable_set(key, value)
+      else
+        thread[key] = value
+      end
+    end
+
+    def self.[](key)
+      get(::Thread.current, key)
+    end
+
+    def self.[]=(key, value)
+      set(::Thread.current, key, value)
+    end
+  end
+end

newrelic.yml

@@ -704,6 +704,9 @@ common: &default_settings
   # the New Relic agent has time to report.
   # sync_startup: false
 
+  # If true, tracer state storage is thread-local, otherwise, fiber-local
+  # thread_local_tracer_state: false
+
   # If true, enables use of the thread profiler.
   # thread_profiler.enabled: false
 

test/multiverse/suites/rake/instrumentation_test.rb

@@ -57,7 +57,7 @@ def test_invoke_with_tracing_with_exception
     NewRelic::Agent::Instrumentation::Rake.stub :should_trace?, true, [instance.name] do
       error = RuntimeError.new('expected')
       # produce the exception we want to have the method rescue
-      NewRelic::Agent.stub :config, -> { raise error } do
+      NewRelic::Agent.stub :instance, -> { raise error } do
         logger = MiniTest::Mock.new
         NewRelic::Agent.stub :logger, logger do
           logger.expect :error, nil, [/^Exception/, error]

test/new_relic/agent/tracer_test.rb

@@ -17,6 +17,16 @@ def test_tracer_aliases
         refute_nil state
       end
 
+      def test_tracer_state_in_fiber_with_thread_local_tracer_state
+        with_config(:thread_local_tracer_state => true) do
+          state = Tracer.state
+          fiber = Fiber.new do
+            assert_equal Tracer.state, state
+          end
+          fiber.resume
+        end
+      end
+
       def test_current_transaction_with_transaction
         in_transaction do |txn|
           assert_equal txn, Tracer.current_transaction

test/new_relic/thread_local_storage_test.rb

@@ -0,0 +1,70 @@
+# This file is distributed under New Relic's license terms.
+# See https://github.com/newrelic/newrelic-ruby-agent/blob/main/LICENSE for complete details.
+# frozen_string_literal: true
+
+require_relative '../test_helper'
+require 'new_relic/thread_local_storage'
+
+class NewRelic::ThreadLocalStorageTest < Minitest::Test
+  def test_basic_ops
+    assert_nil NewRelic::ThreadLocalStorage.get(Thread.current, :basic)
+    NewRelic::ThreadLocalStorage.set(Thread.current, :basic, 'foobar')
+
+    assert_equal('foobar', NewRelic::ThreadLocalStorage.get(Thread.current, :basic))
+    NewRelic::ThreadLocalStorage.set(Thread.current, :basic, 12345)
+
+    assert_equal(12345, NewRelic::ThreadLocalStorage.get(Thread.current, :basic))
+  end
+
+  def test_shortcut_ops
+    assert_nil NewRelic::ThreadLocalStorage[:shortcut]
+    NewRelic::ThreadLocalStorage[:shortcut] = 'baz'
+
+    assert_equal('baz', NewRelic::ThreadLocalStorage[:shortcut])
+    NewRelic::ThreadLocalStorage[:shortcut] = 98765
+
+    assert_equal(98765, NewRelic::ThreadLocalStorage[:shortcut])
+  end
+
+  def test_new_thread
+    NewRelic::ThreadLocalStorage[:new_thread] = :parent
+    thread = Thread.new do
+      Thread.current.abort_on_exception = true
+
+      assert_nil NewRelic::ThreadLocalStorage[:new_thread]
+      NewRelic::ThreadLocalStorage[:new_thread] = :child
+      sleep
+    end
+    sleep 0.2
+
+    assert_equal(:parent, NewRelic::ThreadLocalStorage[:new_thread])
+    assert_equal(:child, NewRelic::ThreadLocalStorage.get(thread, :new_thread))
+    thread.exit
+  end
+
+  def test_new_fiber_without_thread_local_tracer_state_flag
+    with_config(:thread_local_tracer_state => false) do
+      NewRelic::ThreadLocalStorage[:new_fiber_no_flag] = 1
+      fiber = Fiber.new do
+        assert_nil NewRelic::ThreadLocalStorage[:new_fiber_no_flag]
+        NewRelic::ThreadLocalStorage[:new_fiber_no_flag] = 2
+      end
+      fiber.resume
+
+      assert_equal(1, NewRelic::ThreadLocalStorage[:new_fiber_no_flag])
+    end
+  end
+
+  def test_new_fiber_with_thread_local_tracer_state_flag
+    with_config(:thread_local_tracer_state => true) do
+      NewRelic::ThreadLocalStorage[:new_fiber_with_flag] = 1
+      fiber = Fiber.new do
+        assert_equal(1, NewRelic::ThreadLocalStorage[:new_fiber_with_flag])
+        NewRelic::ThreadLocalStorage[:new_fiber_with_flag] = 2
+      end
+      fiber.resume
+
+      assert_equal(2, NewRelic::ThreadLocalStorage[:new_fiber_with_flag])
+    end
+  end
+end