Support for traces_sampler option (#910)

* Add support for `traces_sampler` config option

* Update docs

* Remove unnecessary attribute merging

* Call traces_sampler only for root spans

* Reuse sampling decision logic

* Add test for handling invalid traces_sampler return values

* Add tests for SamplingContext access functions

Author: solnic

.dialyzer_ignore.exs

@@ -1,4 +1,5 @@
 [
   {"test/support/example_plug_application.ex"},
-  {"test/support/test_helpers.ex"}
+  {"test/support/test_helpers.ex"},
+  {"lib/sentry/opentelemetry/sampler.ex", :pattern_match, 1}
 ]

lib/sentry/config.ex

@@ -1,6 +1,13 @@
 defmodule Sentry.Config do
   @moduledoc false
 
+  @typedoc """
+  A function that determines the sample rate for transaction events.
+
+  The function receives a sampling context map and should return a boolean or a float between `0.0` and `1.0`.
+  """
+  @type traces_sampler_function :: (map() -> boolean() | float()) | {module(), atom()}
+
   integrations_schema = [
     max_expected_check_in_time: [
       type: :integer,
@@ -158,6 +165,34 @@ defmodule Sentry.Config do
       for guides on how to set it up.
       """
     ],
+    traces_sampler: [
+      type: {:custom, __MODULE__, :__validate_traces_sampler__, []},
+      default: nil,
+      type_doc: "`t:traces_sampler_function/0` or `nil`",
+      doc: """
+      A function that determines the sample rate for transaction events. This function
+      receives a sampling context struct and should return a boolean or a float between `0.0` and `1.0`.
+
+      The sampling context contains:
+      - `:parent_sampled` - boolean indicating if the parent trace span was sampled (nil if no parent)
+      - `:transaction_context` - map with transaction information (name, op, etc.)
+
+      If both `:traces_sampler` and `:traces_sample_rate` are configured, `:traces_sampler` takes precedence.
+
+      Example:
+      ```elixir
+      traces_sampler: fn sampling_context ->
+        case sampling_context.transaction_context.op do
+          "http.server" -> 0.1  # Sample 10% of HTTP requests
+          "db.query" -> 0.01    # Sample 1% of database queries
+          _ -> false            # Don't sample other operations
+        end
+      end
+      ```
+
+      This value is also used to determine if tracing is enabled: if it's not `nil`, tracing is enabled.
+      """
+    ],
     included_environments: [
       type: {:or, [{:in, [:all]}, {:list, {:or, [:atom, :string]}}]},
       deprecated: "Use :dsn to control whether to send events to Sentry.",
@@ -625,6 +660,9 @@ defmodule Sentry.Config do
   @spec traces_sample_rate() :: nil | float()
   def traces_sample_rate, do: fetch!(:traces_sample_rate)
 
+  @spec traces_sampler() :: traces_sampler_function() | nil
+  def traces_sampler, do: get(:traces_sampler)
+
   @spec hackney_opts() :: keyword()
   def hackney_opts, do: fetch!(:hackney_opts)
 
@@ -663,7 +701,7 @@ defmodule Sentry.Config do
   def integrations, do: fetch!(:integrations)
 
   @spec tracing?() :: boolean()
-  def tracing?, do: not is_nil(fetch!(:traces_sample_rate))
+  def tracing?, do: not is_nil(fetch!(:traces_sample_rate)) or not is_nil(get(:traces_sampler))
 
   @spec put_config(atom(), term()) :: :ok
   def put_config(key, value) when is_atom(key) do
@@ -773,6 +811,26 @@ defmodule Sentry.Config do
     end
   end
 
+  def __validate_traces_sampler__(nil), do: {:ok, nil}
+
+  def __validate_traces_sampler__(fun) when is_function(fun, 1) do
+    {:ok, fun}
+  end
+
+  def __validate_traces_sampler__({module, function})
+      when is_atom(module) and is_atom(function) do
+    if function_exported?(module, function, 1) do
+      {:ok, {module, function}}
+    else
+      {:error, "function #{module}.#{function}/1 is not exported"}
+    end
+  end
+
+  def __validate_traces_sampler__(other) do
+    {:error,
+     "expected :traces_sampler to be nil, a function with arity 1, or a {module, function} tuple, got: #{inspect(other)}"}
+  end
+
   def __validate_json_library__(nil) do
     {:error, "nil is not a valid value for the :json_library option"}
   end

lib/sentry/opentelemetry/sampler.ex

@@ -4,6 +4,9 @@ if Code.ensure_loaded?(:otel_sampler) do
 
     alias OpenTelemetry.{Span, Tracer}
     alias Sentry.ClientReport
+    alias SamplingContext
+
+    require Logger
 
     @behaviour :otel_sampler
 
@@ -24,27 +27,34 @@ if Code.ensure_loaded?(:otel_sampler) do
     @impl true
     def should_sample(
           ctx,
-          _trace_id,
+          trace_id,
           _links,
           span_name,
-          _span_kind,
-          _attributes,
+          span_kind,
+          attributes,
           config
         ) do
       result =
         if span_name in config[:drop] do
           {:drop, [], []}
         else
-          sample_rate = Sentry.Config.traces_sample_rate()
+          traces_sampler = Sentry.Config.traces_sampler()
+          traces_sample_rate = Sentry.Config.traces_sample_rate()
 
           case get_trace_sampling_decision(ctx) do
             {:inherit, trace_sampled, tracestate} ->
               decision = if trace_sampled, do: :record_and_sample, else: :drop
-
               {decision, [], tracestate}
 
             :no_trace ->
-              make_sampling_decision(sample_rate)
+              if traces_sampler do
+                sampling_context =
+                  build_sampling_context(nil, span_name, span_kind, attributes, trace_id)
+
+                make_sampler_decision(traces_sampler, sampling_context)
+              else
+                make_sampling_decision(traces_sample_rate)
+              end
           end
         end
 
@@ -121,6 +131,56 @@ if Code.ensure_loaded?(:otel_sampler) do
       end
     end
 
+    defp build_sampling_context(parent_sampled, span_name, _span_kind, attributes, trace_id) do
+      transaction_context = %{
+        name: span_name,
+        op: span_name,
+        trace_id: trace_id,
+        attributes: attributes
+      }
+
+      sampling_context = %SamplingContext{
+        transaction_context: transaction_context,
+        parent_sampled: parent_sampled
+      }
+
+      sampling_context
+    end
+
+    defp make_sampler_decision(traces_sampler, sampling_context) do
+      try do
+        result = call_traces_sampler(traces_sampler, sampling_context)
+        sample_rate = normalize_sampler_result(result)
+
+        if is_float(sample_rate) and sample_rate >= 0.0 and sample_rate <= 1.0 do
+          make_sampling_decision(sample_rate)
+        else
+          Logger.warning(
+            "traces_sampler function returned an invalid sample rate: #{inspect(sample_rate)}"
+          )
+
+          make_sampling_decision(0.0)
+        end
+      rescue
+        error ->
+          Logger.warning("traces_sampler function failed: #{inspect(error)}")
+
+          make_sampling_decision(0.0)
+      end
+    end
+
+    defp call_traces_sampler(fun, sampling_context) when is_function(fun, 1) do
+      fun.(sampling_context)
+    end
+
+    defp call_traces_sampler({module, function}, sampling_context) do
+      apply(module, function, [sampling_context])
+    end
+
+    defp normalize_sampler_result(true), do: 1.0
+    defp normalize_sampler_result(false), do: 0.0
+    defp normalize_sampler_result(rate), do: rate
+
     defp record_discarded_transaction() do
       ClientReport.Sender.record_discarded_events(:sample_rate, "transaction")
     end

lib/sentry/sampling_context.ex

@@ -0,0 +1,57 @@
+defmodule SamplingContext do
+  @moduledoc """
+  The struct for the **sampling_context** that is passed to `traces_sampler`.
+
+  This is set up via `Sentry.OpenTelemetry.Sampler`.
+
+  See also <https://develop.sentry.dev/sdk/telemetry/traces/#sampling-context>.
+  """
+
+  @moduledoc since: "11.0.0"
+
+  @typedoc """
+  The sampling context struct that contains information needed for sampling decisions.
+
+  This matches the structure used in the Python SDK's create_sampling_context function.
+  """
+  @type t :: %__MODULE__{
+          transaction_context: %{
+            name: String.t() | nil,
+            op: String.t(),
+            trace_id: String.t(),
+            attributes: map()
+          },
+          parent_sampled: boolean() | nil
+        }
+
+  @enforce_keys [:transaction_context, :parent_sampled]
+  defstruct [:transaction_context, :parent_sampled]
+
+  @behaviour Access
+
+  @impl Access
+  def fetch(struct, key) do
+    case Map.fetch(struct, key) do
+      {:ok, value} -> {:ok, value}
+      :error -> :error
+    end
+  end
+
+  @impl Access
+  def get_and_update(struct, key, function) do
+    current_value = Map.get(struct, key)
+
+    case function.(current_value) do
+      {get_value, update_value} ->
+        {get_value, Map.put(struct, key, update_value)}
+
+      :pop ->
+        {current_value, Map.delete(struct, key)}
+    end
+  end
+
+  @impl Access
+  def pop(struct, key) do
+    {Map.get(struct, key), Map.delete(struct, key)}
+  end
+end

test/sentry/config_traces_sampler_test.exs

@@ -0,0 +1,81 @@
+defmodule Sentry.ConfigTracesSamplerTest do
+  use ExUnit.Case, async: true
+
+  import Sentry.TestHelpers
+
+  describe "traces_sampler configuration validation" do
+    defmodule TestSampler do
+      def sample(_context), do: 0.5
+    end
+
+    test "accepts nil" do
+      assert :ok = put_test_config(traces_sampler: nil)
+      assert Sentry.Config.traces_sampler() == nil
+    end
+
+    test "accepts function with arity 1" do
+      fun = fn _context -> 0.5 end
+      assert :ok = put_test_config(traces_sampler: fun)
+      assert Sentry.Config.traces_sampler() == fun
+    end
+
+    test "accepts MFA tuple with exported function" do
+      assert :ok = put_test_config(traces_sampler: {TestSampler, :sample})
+      assert Sentry.Config.traces_sampler() == {TestSampler, :sample}
+    end
+
+    test "rejects MFA tuple with non-exported function" do
+      assert_raise ArgumentError, ~r/function.*is not exported/, fn ->
+        put_test_config(traces_sampler: {TestSampler, :non_existent})
+      end
+    end
+
+    test "rejects function with wrong arity" do
+      fun = fn -> 0.5 end
+
+      assert_raise ArgumentError, ~r/expected :traces_sampler to be/, fn ->
+        put_test_config(traces_sampler: fun)
+      end
+    end
+
+    test "rejects invalid types" do
+      assert_raise ArgumentError, ~r/expected :traces_sampler to be/, fn ->
+        put_test_config(traces_sampler: "invalid")
+      end
+
+      assert_raise ArgumentError, ~r/expected :traces_sampler to be/, fn ->
+        put_test_config(traces_sampler: 123)
+      end
+
+      assert_raise ArgumentError, ~r/expected :traces_sampler to be/, fn ->
+        put_test_config(traces_sampler: [])
+      end
+    end
+  end
+
+  describe "tracing? function" do
+    test "returns true when traces_sample_rate is set" do
+      put_test_config(traces_sample_rate: 0.5, traces_sampler: nil)
+
+      assert Sentry.Config.tracing?()
+    end
+
+    test "returns true when traces_sampler is set" do
+      put_test_config(traces_sample_rate: nil, traces_sampler: fn _ -> 0.5 end)
+
+      assert Sentry.Config.tracing?()
+    end
+
+    test "returns true when both are set" do
+      put_test_config(traces_sample_rate: 0.5, traces_sampler: fn _ -> 0.5 end)
+
+      assert Sentry.Config.tracing?()
+    end
+
+    test "returns false when neither is set" do
+      put_test_config(traces_sample_rate: nil, traces_sampler: nil)
+
+      refute Sentry.Config.tracing?()
+    end
+  end
+end