Better span sampler

This updates OpenTelemetry.Sampler so that it uses
`traces_sample_rate` for sampling spans before they
get sent to the span processor.

This way we're not processing spans when they are
dropped during sampling, which was previously the
case as the Client was responsible for sampling
just before attemping to send a transaction.

Author: solnic

lib/sentry/client.ex

@@ -110,14 +110,12 @@ defmodule Sentry.Client do
   @spec send_transaction(Transaction.t(), keyword()) ::
           {:ok, transaction_id :: String.t()}
           | {:error, ClientError.t()}
-          | :unsampled
           | :excluded
   def send_transaction(%Transaction{} = transaction, opts \\ []) do
     opts = NimbleOptions.validate!(opts, Options.send_transaction_schema())
 
     result_type = Keyword.get_lazy(opts, :result, &Config.send_result/0)
     client = Keyword.get_lazy(opts, :client, &Config.client/0)
-    sample_rate = Keyword.get_lazy(opts, :sample_rate, &Config.traces_sample_rate/0)
     before_send = Keyword.get_lazy(opts, :before_send, &Config.before_send/0)
     after_send_event = Keyword.get_lazy(opts, :after_send_event, &Config.after_send_event/0)
 
@@ -126,16 +124,11 @@ defmodule Sentry.Client do
         Application.get_env(:sentry, :request_retries, Transport.default_retries())
       end)
 
-    with :ok <- sample_event(sample_rate),
-         {:ok, %Transaction{} = transaction} <- maybe_call_before_send(transaction, before_send) do
+    with {:ok, %Transaction{} = transaction} <- maybe_call_before_send(transaction, before_send) do
       send_result = encode_and_send(transaction, result_type, client, request_retries)
       _ignored = maybe_call_after_send(transaction, send_result, after_send_event)
       send_result
     else
-      :unsampled ->
-        ClientReport.Sender.record_discarded_events(:sample_rate, [transaction])
-        :unsampled
-
       :excluded ->
         :excluded
     end

lib/sentry/opentelemetry/sampler.ex

@@ -2,19 +2,28 @@ if Code.ensure_loaded?(:otel_sampler) do
   defmodule Sentry.OpenTelemetry.Sampler do
     @moduledoc false
 
+    alias OpenTelemetry.{Span, Tracer}
+
     @behaviour :otel_sampler
 
+    @sentry_sample_rate_key "sentry-sample_rate"
+    @sentry_sample_rand_key "sentry-sample_rand"
+    @sentry_sampled_key "sentry-sampled"
+
+    @impl true
     def setup(config) do
       config
     end
 
+    @impl true
     def description(_) do
       "SentrySampler"
     end
 
+    @impl true
     def should_sample(
-          _ctx,
-          _trace_id,
+          ctx,
+          trace_id,
           _links,
           span_name,
           _span_kind,
@@ -24,7 +33,78 @@ if Code.ensure_loaded?(:otel_sampler) do
       if span_name in config[:drop] do
         {:drop, [], []}
       else
-        {:record_and_sample, [], []}
+        sample_rate = Sentry.Config.traces_sample_rate()
+
+        case get_parent_sampling_decision(ctx, trace_id) do
+          {:inherit, parent_sampled, tracestate} ->
+            decision = if parent_sampled, do: :record_and_sample, else: :drop
+            {decision, [], tracestate}
+
+          :no_parent ->
+            make_sampling_decision(sample_rate)
+        end
+      end
+    end
+
+    defp get_parent_sampling_decision(ctx, trace_id) do
+      case Tracer.current_span_ctx(ctx) do
+        :undefined ->
+          :no_parent
+
+        span_ctx ->
+          parent_trace_id = Span.trace_id(span_ctx)
+
+          if parent_trace_id == trace_id do
+            tracestate = Span.tracestate(span_ctx)
+            parent_sampled = get_tracestate_value(tracestate, @sentry_sampled_key)
+
+            case parent_sampled do
+              "true" -> {:inherit, true, tracestate}
+              "false" -> {:inherit, false, tracestate}
+              nil -> :no_parent
+            end
+          else
+            :no_parent
+          end
+      end
+    end
+
+    defp make_sampling_decision(sample_rate) do
+      cond do
+        sample_rate == 0.0 ->
+          tracestate = build_tracestate(sample_rate, 1.0, false)
+          {:drop, [], tracestate}
+
+        sample_rate == 1.0 ->
+          tracestate = build_tracestate(sample_rate, 0.0, true)
+          {:record_and_sample, [], tracestate}
+
+        true ->
+          random_value = :rand.uniform()
+          sampled = random_value < sample_rate
+
+          tracestate = build_tracestate(sample_rate, random_value, sampled)
+          decision = if sampled, do: :record_and_sample, else: :drop
+          {decision, [], tracestate}
+      end
+    end
+
+    defp build_tracestate(sample_rate, random_value, sampled) do
+      [
+        {@sentry_sample_rate_key, Float.to_string(sample_rate)},
+        {@sentry_sample_rand_key, Float.to_string(random_value)},
+        {@sentry_sampled_key, to_string(sampled)}
+      ]
+    end
+
+    defp get_tracestate_value({:tracestate, tracestate}, key) do
+      get_tracestate_value(tracestate, key)
+    end
+
+    defp get_tracestate_value(tracestate, key) when is_list(tracestate) do
+      case List.keyfind(tracestate, key, 0) do
+        {^key, value} -> value
+        nil -> nil
       end
     end
   end

lib/sentry/opentelemetry/span_processor.ex

@@ -39,6 +39,9 @@ if Code.ensure_loaded?(OpenTelemetry) do
             :ignored ->
               true
 
+            :excluded ->
+              true
+
             {:error, error} ->
               Logger.warning("Failed to send transaction to Sentry: #{inspect(error)}")
               {:error, :invalid_span}

test/sentry/opentelemetry/sampler_test.exs

@@ -3,17 +3,103 @@ defmodule Sentry.Opentelemetry.SamplerTest do
 
   alias Sentry.OpenTelemetry.Sampler
 
-  test "drops spans with the given name" do
-    assert {:drop, [], []} =
-             Sampler.should_sample(nil, nil, nil, "Elixir.Oban.Stager process", nil, nil,
-               drop: ["Elixir.Oban.Stager process"]
-             )
+  setup do
+    original_rate = Sentry.Config.traces_sample_rate()
+
+    on_exit(fn ->
+      Sentry.Config.put_config(:traces_sample_rate, original_rate)
+    end)
+
+    :ok
+  end
+
+  describe "span name dropping" do
+    test "drops spans with the given name" do
+      assert {:drop, [], []} =
+               Sampler.should_sample(nil, nil, nil, "Elixir.Oban.Stager process", nil, nil,
+                 drop: ["Elixir.Oban.Stager process"]
+               )
+    end
+
+    test "records and samples spans not in drop list" do
+      Sentry.Config.put_config(:traces_sample_rate, 1.0)
+
+      assert {:record_and_sample, [], tracestate} =
+               Sampler.should_sample(nil, 123, nil, "Elixir.Oban.Worker process", nil, nil,
+                 drop: []
+               )
+
+      assert is_list(tracestate)
+      assert {"sentry-sample_rate", "1.0"} in tracestate
+      assert {"sentry-sampled", "true"} in tracestate
+    end
   end
 
-  test "records and samples spans with the given name" do
-    assert {:record_and_sample, [], []} =
-             Sampler.should_sample(nil, nil, nil, "Elixir.Oban.Worker process", nil, nil,
-               drop: []
-             )
+  describe "sampling based on traces_sample_rate" do
+    test "always drops when sample rate is 0.0" do
+      Sentry.Config.put_config(:traces_sample_rate, 0.0)
+
+      assert {:drop, [], tracestate} =
+               Sampler.should_sample(nil, 123, nil, "test span", nil, nil, drop: [])
+
+      assert {"sentry-sample_rate", "0.0"} in tracestate
+      assert {"sentry-sampled", "false"} in tracestate
+    end
+
+    test "always samples when sample rate is 1.0" do
+      Sentry.Config.put_config(:traces_sample_rate, 1.0)
+
+      assert {:record_and_sample, [], tracestate} =
+               Sampler.should_sample(nil, 123, nil, "test span", nil, nil, drop: [])
+
+      assert {"sentry-sample_rate", "1.0"} in tracestate
+      assert {"sentry-sampled", "true"} in tracestate
+    end
+
+    test "different trace_ids produce different sampling decisions" do
+      Sentry.Config.put_config(:traces_sample_rate, 0.5)
+
+      trace_ids = Enum.to_list(1..100)
+
+      results =
+        Enum.map(trace_ids, fn trace_id ->
+          {decision, [], _tracestate} =
+            Sampler.should_sample(nil, trace_id, nil, "test span", nil, nil, drop: [])
+
+          decision == :record_and_sample
+        end)
+
+      sampled_count = Enum.count(results, & &1)
+
+      assert sampled_count > 30 and sampled_count < 70
+    end
+  end
+
+  describe "parent span inheritance" do
+    test "inherits sampling decision from parent span with same trace_id" do
+      Sentry.Config.put_config(:traces_sample_rate, 1.0)
+
+      assert {:record_and_sample, [], _tracestate} =
+               Sampler.should_sample(nil, 123, nil, "test span", nil, nil, drop: [])
+    end
+  end
+
+  describe "tracestate management" do
+    test "builds tracestate with correct format" do
+      Sentry.Config.put_config(:traces_sample_rate, 0.75)
+
+      {_decision, [], tracestate} =
+        Sampler.should_sample(nil, 123, nil, "test span", nil, nil, drop: [])
+
+      assert List.keyfind(tracestate, "sentry-sample_rate", 0)
+      assert List.keyfind(tracestate, "sentry-sample_rand", 0)
+      assert List.keyfind(tracestate, "sentry-sampled", 0)
+
+      {"sentry-sample_rate", rate_str} = List.keyfind(tracestate, "sentry-sample_rate", 0)
+      assert rate_str == "0.75"
+
+      {"sentry-sampled", sampled_str} = List.keyfind(tracestate, "sentry-sampled", 0)
+      assert sampled_str in ["true", "false"]
+    end
   end
 end