Add ThreadLocalStorage abstraction to store thread-local data

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/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,23 @@
+# 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)
+      thread[key]
+    end
+
+    def self.set(thread, key, value)
+      thread[key] = value
+    end
+
+    def self.[](key)
+      get(::Thread.current, key)
+    end
+
+    def self.[]=(key, value)
+      set(::Thread.current, key, value)
+    end
+  end
+end

test/new_relic/thread_local_storage_test.rb

@@ -0,0 +1,42 @@
+# 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
+      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
+end