Refactor span storage (#817)

* initial refactor with ETS for span storage

* add genserver back

* Fix updating root spans in ets

* remove comment

---------

Co-authored-by: Peter Solnica <[email protected]>

Author: savhappy

lib/sentry/application.ex

@@ -10,9 +10,6 @@ defmodule Sentry.Application do
     config = Config.validate!()
     :ok = Config.persist(config)
 
-    # Setup ETS tables for span storage
-    Sentry.Opentelemetry.SpanStorage.setup()
-
     http_client = Keyword.fetch!(config, :client)
 
     maybe_http_client_spec =
@@ -30,6 +27,7 @@ defmodule Sentry.Application do
         Sentry.Sources,
         Sentry.Dedupe,
         Sentry.ClientReport.Sender,
+        Sentry.Opentelemetry.SpanStorage,
         {Sentry.Integrations.CheckInIDMappings,
          [
            max_expected_check_in_time:

lib/sentry/opentelemetry/span_storage.ex

@@ -1,78 +1,79 @@
 defmodule Sentry.Opentelemetry.SpanStorage do
   @moduledoc false
+  use GenServer
 
-  @root_spans_table :sentry_root_spans
-  @child_spans_table :sentry_child_spans
+  @table :span_storage
 
-  def setup do
-    case :ets.whereis(@root_spans_table) do
-      :undefined ->
-        :ets.new(@root_spans_table, [:set, :public, :named_table])
+  @spec start_link(keyword()) :: GenServer.on_start()
+  def start_link(opts \\ []) do
+    name = Keyword.get(opts, :name, __MODULE__)
+    GenServer.start_link(__MODULE__, nil, name: name)
+  end
 
-      _ ->
-        :ok
-    end
+  @impl true
+  def init(nil) do
+    _table =
+      if :ets.whereis(@table) == :undefined do
+        :ets.new(@table, [:named_table, :public, :bag])
+      end
 
-    case :ets.whereis(@child_spans_table) do
-      :undefined ->
-        :ets.new(@child_spans_table, [:bag, :public, :named_table])
+    {:ok, :no_state}
+  end
 
-      _ ->
-        :ok
+  def store_span(span_data) when span_data.parent_span_id == nil do
+    case :ets.lookup(@table, {:root_span, span_data.span_id}) do
+      [] -> :ets.insert(@table, {{:root_span, span_data.span_id}, span_data})
+      _ -> :ok
     end
-
-    :ok
   end
 
   def store_span(span_data) do
-    if span_data.parent_span_id == nil do
-      :ets.insert(@root_spans_table, {span_data.span_id, span_data})
-    else
-      :ets.insert(@child_spans_table, {span_data.parent_span_id, span_data})
-    end
-
-    :ok
+    _ = :ets.insert(@table, {span_data.parent_span_id, span_data})
   end
 
   def get_root_span(span_id) do
-    case :ets.lookup(@root_spans_table, span_id) do
-      [{^span_id, span}] -> span
+    case :ets.lookup(@table, {:root_span, span_id}) do
+      [{{:root_span, ^span_id}, span}] -> span
       [] -> nil
     end
   end
 
   def get_child_spans(parent_span_id) do
-    :ets.lookup(@child_spans_table, parent_span_id)
+    :ets.lookup(@table, parent_span_id)
     |> Enum.map(fn {_parent_id, span} -> span end)
   end
 
   def update_span(span_data) do
     if span_data.parent_span_id == nil do
-      :ets.insert(@root_spans_table, {span_data.span_id, span_data})
+      case :ets.lookup(@table, {:root_span, span_data.span_id}) do
+        [] ->
+          :ets.insert(@table, {{:root_span, span_data.span_id}, span_data})
+
+        _ ->
+          :ets.delete(@table, {:root_span, span_data.span_id})
+          :ets.insert(@table, {{:root_span, span_data.span_id}, span_data})
+      end
     else
-      existing_spans = :ets.lookup(@child_spans_table, span_data.parent_span_id)
-
-      :ets.delete(@child_spans_table, span_data.parent_span_id)
+      existing_spans = :ets.lookup(@table, span_data.parent_span_id)
 
       Enum.each(existing_spans, fn {parent_id, span} ->
-        if span.span_id != span_data.span_id do
-          :ets.insert(@child_spans_table, {parent_id, span})
+        if span.span_id == span_data.span_id do
+          :ets.delete_object(@table, {parent_id, span})
+          :ets.insert(@table, {span_data.parent_span_id, span_data})
         end
       end)
-
-      :ets.insert(@child_spans_table, {span_data.parent_span_id, span_data})
     end
 
     :ok
   end
 
   def remove_span(span_id) do
-    :ets.delete(@root_spans_table, span_id)
+    :ets.delete(@table, {:root_span, span_id})
     :ok
   end
 
   def remove_child_spans(parent_span_id) do
-    :ets.delete(@child_spans_table, parent_span_id)
+    :ets.delete(@table, parent_span_id)
     :ok
   end
 end

test/sentry/opentelemetry/span_processor_test.exs

@@ -6,17 +6,10 @@ defmodule Sentry.Opentelemetry.SpanProcessorTest do
   alias Sentry.Opentelemetry.SpanStorage
 
   setup do
-    # Create tables
-    SpanStorage.setup()
-
     on_exit(fn ->
       # Only try to clean up tables if they exist
-      if :ets.whereis(:sentry_root_spans) != :undefined do
-        :ets.delete_all_objects(:sentry_root_spans)
-      end
-
-      if :ets.whereis(:sentry_child_spans) != :undefined do
-        :ets.delete_all_objects(:sentry_child_spans)
+      if :ets.whereis(:span_storage) != :undefined do
+        :ets.delete_all_objects(:span_storage)
       end
     end)
 
@@ -71,7 +64,6 @@ defmodule Sentry.Opentelemetry.SpanProcessorTest do
     TestEndpoint.instrumented_function()
 
     assert [%Sentry.Transaction{} = transaction] = Sentry.Test.pop_sentry_transactions()
-
     assert_valid_iso8601(transaction.timestamp)
     assert_valid_iso8601(transaction.start_timestamp)
     assert transaction.timestamp > transaction.start_timestamp