Add execution-context generation token to DI snapshots

ext/libdatadog_api/di.c

@@ -39,6 +39,25 @@ static ID id_mesg;
 // method probes on Thread#[] / Thread#[]= cannot intercept guard reads/writes.
 static ID id_datadog_di_in_probe;
 
+// ID for the ivar on each Thread object that holds that thread's snapshot
+// generation token.
+static ID id_datadog_di_thread_generation;
+
+// Process-wide source of snapshot generation tokens. Incremented only inside
+// current_thread_generation, which runs with the GVL held and never releases
+// it, so the counter and the accompanying ivar write need no lock.
+static unsigned long long generation_counter;
+
+// Snapshot generation tokens for Thread objects frozen before their first
+// snapshot. rb_ivar_set raises FrozenError on a frozen Thread, so those
+// tokens are keyed by the Thread object here instead. The table is plain C
+// memory invisible to the GC and holds no reference to the Thread, so an
+// entry outlives its thread: a Thread allocated later at the same address
+// reuses the dead thread's token, and the table grows by one entry per
+// frozen thread that captured a snapshot. Guarded by the GVL like the
+// counter above.
+static st_table *frozen_thread_generations;
+
 // Returns the imemo type of an imemo object, that is `imemo_type()` from CRuby's internal/imemo.h.
 // The caller must have already checked the object is a T_IMEMO, otherwise the result is meaningless.
 static inline int ddtrace_imemo_type(VALUE imemo) {
@@ -151,6 +170,59 @@ static VALUE leave_probe(DDTRACE_UNUSED VALUE _self) {
   return Qnil;
 }
 
+/*
+ * Returns the snapshot generation token of a frozen Thread, assigning
+ * one from the process-wide counter on the first call for the thread.
+ * Only called for a Thread whose ivar write would raise FrozenError.
+ */
+static VALUE frozen_thread_generation(VALUE thread) {
+  st_data_t token;
+  if (st_lookup(frozen_thread_generations, (st_data_t)thread, &token)) {
+    return ULL2NUM((unsigned long long)token);
+  }
+  token = (st_data_t)++generation_counter;
+  st_insert(frozen_thread_generations, (st_data_t)thread, token);
+  return ULL2NUM((unsigned long long)token);
+}
+
+/*
+ * call-seq:
+ *   DI.current_thread_generation -> Integer
+ *
+ * Returns the snapshot generation token for the current thread, assigning
+ * one on the first call for the thread. Snapshots captured by the same
+ * thread share a token, and snapshots captured by different threads carry
+ * different tokens, so a consumer can tell which thread captured each
+ * snapshot. All fibers of a thread share that thread's token. Distinct
+ * Thread objects get distinct tokens, even when native thread ids are
+ * reused. Tokens are unique only within a runtime id.
+ *
+ * The token is a value of the process-wide counter above, stored in an
+ * ivar on the Thread object. The ivar is read via rb_attr_get, which
+ * returns nil for an unset ivar and, unlike rb_ivar_get, emits no
+ * verbose-mode warning on Ruby 2.6 and 2.7; it is written via rb_ivar_set.
+ * A Thread frozen before its first snapshot, where writing the ivar would
+ * raise FrozenError, keeps its token in a process-wide table keyed by the
+ * Thread object. Both accesses bypass Thread#thread_variable_get and
+ * Thread#thread_variable_set method dispatch, so a user-installed method
+ * probe on those methods cannot fire during DI snapshot building.
+ *
+ * @return [Integer] the generation token of the current thread
+ * @api private
+ */
+static VALUE current_thread_generation(DDTRACE_UNUSED VALUE _self) {
+  VALUE thread = rb_thread_current();
+  VALUE generation = rb_attr_get(thread, id_datadog_di_thread_generation);
+  if (NIL_P(generation)) {
+    if (OBJ_FROZEN(thread)) {
+      return frozen_thread_generation(thread);
+    }
+    generation = ULL2NUM(++generation_counter);
+    rb_ivar_set(thread, id_datadog_di_thread_generation, generation);
+  }
+  return generation;
+}
+
 /*
  * call-seq:
  *   DI.hash?(obj) -> true | false
@@ -224,6 +296,8 @@ void di_init(VALUE datadog_module) {
 
   id_mesg = rb_intern("mesg");
   id_datadog_di_in_probe = rb_intern("datadog_di_in_probe");
+  id_datadog_di_thread_generation = rb_intern("@datadog_di_thread_generation");
+  frozen_thread_generations = st_init_numtable();
 
   VALUE di_module = rb_define_module_under(datadog_module, "DI");
   rb_define_singleton_method(di_module, "all_iseqs", all_iseqs, 0);
@@ -232,6 +306,7 @@ void di_init(VALUE datadog_module) {
   rb_define_singleton_method(di_module, "enter_probe", enter_probe, 0);
   rb_define_singleton_method(di_module, "leave_probe", leave_probe, 0);
   rb_define_singleton_method(di_module, "hash?", is_hash, 1);
+  rb_define_singleton_method(di_module, "current_thread_generation", current_thread_generation, 0);
   #ifdef HAVE_RB_ISEQ_TYPE
     rb_define_singleton_method(di_module, "iseq_type", iseq_type, 1);
   #endif

lib/datadog/di/probe_notification_builder.rb

@@ -306,6 +306,14 @@ def format_backtrace_strings(backtrace)
         end
       end
 
+      # Builds the snapshot payload common to all probe snapshot notifications.
+      #
+      # @param context [Context] probe execution context for the snapshot
+      # @param evaluation_errors [Array<EvaluationFailure>] errors from
+      #   evaluating the probe template or capture expressions
+      # @param captures [Hash, nil] serialized captured values
+      # @param message [String, nil] rendered probe template message
+      # @return [Hash] the snapshot payload
       def build_snapshot_base(context, evaluation_errors: [], captures: nil, message: nil)
         probe = context.probe
 
@@ -379,6 +387,7 @@ def build_snapshot_base(context, evaluation_errors: [], captures: nil, message:
             # we can also determine which thread identifier to send
             # (Thread#native_thread_id or something else).
             thread_id: nil,
+            generation: DI.current_thread_generation,
             version: 2,
           },
           runtime_id: Core::Environment::Identity.id,

sig/datadog/di.rbs

@@ -27,6 +27,8 @@ module Datadog
     # args and return are untyped: variadic forwarding to a customer-supplied
     # Proc whose signature and return value are unknown to this primitive.
 
+    def self.current_thread_generation: () -> Integer
+
     def self.component: () -> Component
 
     def self.instrumented_count: (?instrumentation_kind? kind) -> Integer

spec/datadog/di/ext/thread_generation_spec.rb

@@ -0,0 +1,116 @@
+require "datadog/di/spec_helper"
+
+# Unit tests for the C-implemented snapshot generation token primitive:
+# Datadog::DI.current_thread_generation. These tests exercise the token's
+# semantics: a positive integer, stable for one thread, shared by all fibers
+# of a thread, distinct across threads, and still assigned to a Thread
+# frozen before its first snapshot.
+
+RSpec.describe "Datadog::DI snapshot generation token" do
+  di_test
+  describe ".current_thread_generation" do
+    context "for the calling thread" do
+      let(:token) { Datadog::DI.current_thread_generation }
+
+      before do
+        expect(token).to be_a(Integer)
+        expect(token).to be > 0
+      end
+
+      it "returns a positive integer" do
+        expect(token).to be_a(Integer)
+        expect(token).to be > 0
+      end
+
+      it "returns the same token across calls" do
+        expect(Datadog::DI.current_thread_generation).to eq(token)
+      end
+
+      it "returns the same token to all fibers of the thread" do
+        fiber_token = nil
+        Fiber.new { fiber_token = Datadog::DI.current_thread_generation }.resume
+
+        expect(fiber_token).to eq(token)
+      end
+    end
+
+    context "for a different thread" do
+      let(:calling_thread_token) { Datadog::DI.current_thread_generation }
+
+      before do
+        expect(calling_thread_token).to be_a(Integer)
+        expect(calling_thread_token).to be > 0
+      end
+
+      it "returns a token distinct from the calling thread's token" do
+        other_token = nil
+        thread = Thread.new { other_token = Datadog::DI.current_thread_generation }
+        raise "thread wait timeout" unless thread.join(5)
+
+        expect(other_token).to be_a(Integer)
+        expect(other_token).not_to eq(calling_thread_token)
+      end
+
+      it "returns the same token across calls within the thread" do
+        tokens = []
+        thread = Thread.new { 3.times { tokens << Datadog::DI.current_thread_generation } }
+        raise "thread wait timeout" unless thread.join(5)
+
+        expect(tokens.first).to be_a(Integer)
+        expect(tokens.uniq).to eq([tokens.first])
+      end
+    end
+
+    context "for a thread frozen before its first token" do
+      it "returns a positive integer" do
+        token = nil
+        thread = Thread.new do
+          Thread.current.freeze
+          token = Datadog::DI.current_thread_generation
+        end
+        raise "thread wait timeout" unless thread.join(5)
+
+        expect(token).to be_a(Integer)
+        expect(token).to be > 0
+      end
+
+      it "returns the same token across calls within the thread" do
+        tokens = []
+        thread = Thread.new do
+          Thread.current.freeze
+          3.times { tokens << Datadog::DI.current_thread_generation }
+        end
+        raise "thread wait timeout" unless thread.join(5)
+
+        expect(tokens.first).to be_a(Integer)
+        expect(tokens.uniq).to eq([tokens.first])
+      end
+
+      it "returns a token distinct from the calling thread's token" do
+        frozen_thread_token = nil
+        thread = Thread.new do
+          Thread.current.freeze
+          frozen_thread_token = Datadog::DI.current_thread_generation
+        end
+        raise "thread wait timeout" unless thread.join(5)
+
+        expect(frozen_thread_token).to be_a(Integer)
+        expect(frozen_thread_token).not_to eq(Datadog::DI.current_thread_generation)
+      end
+    end
+
+    context "for a thread frozen after its first token" do
+      it "keeps returning the token assigned before the freeze" do
+        tokens = []
+        thread = Thread.new do
+          tokens << Datadog::DI.current_thread_generation
+          Thread.current.freeze
+          tokens << Datadog::DI.current_thread_generation
+        end
+        raise "thread wait timeout" unless thread.join(5)
+
+        expect(tokens.uniq).to eq([tokens.first])
+      end
+    end
+  end
+end

spec/datadog/di/integration/everything_from_remote_config_spec.rb

@@ -207,6 +207,7 @@ def target_method
         method: "target_method",
         name: nil,
         thread_id: nil,
+        generation: Integer,
         thread_name: "Thread.main",
         version: 2,
       },
@@ -480,6 +481,7 @@ def assert_received_and_errored
             method: "target_method",
             name: nil,
             thread_id: nil,
+            generation: Integer,
             thread_name: "Thread.main",
             version: 2,
           },
@@ -659,6 +661,7 @@ def assert_received_and_errored
             method: nil,
             name: "instrumentation_integration_test_class.rb",
             thread_id: nil,
+            generation: Integer,
             thread_name: "Thread.main",
             version: 2,
           },
@@ -746,6 +749,7 @@ def assert_received_and_errored
               method: nil,
               name: "instrumentation_integration_test_class.rb",
               thread_id: nil,
+              generation: Integer,
               thread_name: "Thread.main",
               version: 2,
             },

spec/datadog/di/probe_notification_builder_spec.rb

@@ -269,6 +269,11 @@
       expect(payload[:runtime_id]).to eq(Datadog::Core::Environment::Identity.id)
     end
 
+    it "includes the thread generation token in the logger object" do
+      expect(payload[:logger][:generation]).to be_an(Integer)
+      expect(payload[:logger][:generation]).to eq(Datadog::DI.current_thread_generation)
+    end
+
     context "with template" do
       let(:probe) do
         Datadog::DI::Probe.new(id: "123", type: :log, file: "X", line_no: 1,
@@ -306,6 +311,7 @@
             method: nil,
             name: "X",
             thread_id: nil,
+            generation: Integer,
             thread_name: "Thread.main",
             version: 2,
           },
@@ -359,6 +365,7 @@
             method: nil,
             name: "X",
             thread_id: nil,
+            generation: Integer,
             thread_name: "Thread.main",
             version: 2,
           },
@@ -436,6 +443,7 @@
             method: nil,
             name: "X",
             thread_id: nil,
+            generation: Integer,
             thread_name: "Thread.main",
             version: 2,
           },

unreleased/6233-di-snapshot-generation-token.json

@@ -0,0 +1,6 @@
+{
+  "message": "Dynamic Instrumentation: snapshots now carry an execution-context generation token to disambiguate reused thread ids.",
+  "product": "Dynamic Instrumentation",
+  "pull_request": "https://github.com/DataDog/dd-trace-rb/pull/6233",
+  "type": "Added"
+}