initial refactor with ETS for span storage

Author: savhappy

lib/sentry/opentelemetry/span_storage.ex

@@ -1,88 +1,82 @@
 defmodule Sentry.Opentelemetry.SpanStorage do
   use GenServer
 
-  def start_link(_opts) do
-    GenServer.start_link(__MODULE__, nil, name: __MODULE__)
-  end
+  @table :span_storage
 
-  def init(_) do
-    {:ok, %{root_spans: %{}, child_spans: %{}}}
+  @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
 
-  def store_span(span_data) do
-    GenServer.call(__MODULE__, {:store_span, span_data})
-  end
+  # start ets table on initialization
+  @impl true
+  def init(nil) do
+    _table =
+      if :ets.whereis(@table) == :undefined do
+        :ets.new(@table, [:named_table, :public, :bag])
+      end
 
-  def get_root_span(span_id) do
-    GenServer.call(__MODULE__, {:get_root_span, span_id})
+    {:ok, :no_state}
   end
 
-  def get_child_spans(parent_span_id) do
-    GenServer.call(__MODULE__, {:get_child_spans, parent_span_id})
+  def store_span(span_data) when span_data.parent_span_id == nil do
+    _ = :ets.insert(@table, {{:root_span, span_data.span_id}, span_data})
   end
 
-  def update_span(span_data) do
-    GenServer.call(__MODULE__, {:update_span, span_data})
+  def store_span(span_data) do
+    _ = :ets.insert(@table, {span_data.parent_span_id, span_data})
   end
 
-  def remove_span(span_id) do
-    GenServer.call(__MODULE__, {:remove_span, span_id})
+  def get_root_span(span_id) do
+    case :ets.lookup(@table, {:root_span, span_id}) do
+      [{{:root_span, ^span_id}, span}] -> span
+      # how do we handle if span cannot be found?
+      [] -> nil
+    end
   end
 
-  def remove_child_spans(parent_span_id) do
-    GenServer.call(__MODULE__, {:remove_child_spans, parent_span_id})
-  end
+  def get_child_spans(parent_span_id) do
+    child_spans = :ets.lookup(@table, parent_span_id)
 
-  def handle_call({:store_span, span_data}, _from, state) do
-    if span_data.parent_span_id == nil do
-      new_state = put_in(state, [:root_spans, span_data.span_id], span_data)
-      {:reply, :ok, new_state}
+    if child_spans == [] do
+      nil
     else
-      new_state =
-        update_in(state, [:child_spans, span_data.parent_span_id], fn spans ->
-          (spans || []) ++ [span_data]
-        end)
-
-      {:reply, :ok, new_state}
+      Enum.map(child_spans, fn {parent_span_id, span} -> span end)
     end
   end
 
-  def handle_call({:get_root_span, span_id}, _from, state) do
-    {:reply, state.root_spans[span_id], state}
-  end
-
-  def handle_call({:get_child_spans, parent_span_id}, _from, state) do
-    {:reply, state.child_spans[parent_span_id] || [], state}
-  end
+  def update_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})
 
-  def handle_call({:update_span, span_data}, _from, state) do
-    if span_data.parent_span_id == nil do
-      new_state = put_in(state, [:root_spans, span_data.span_id], span_data)
-      {:reply, :ok, new_state}
-    else
-      new_state =
-        update_in(state, [:child_spans, span_data.parent_span_id], fn spans ->
-          Enum.map(spans || [], fn span ->
-            if span.span_id == span_data.span_id, do: span_data, else: span
-          end)
-        end)
-
-      {:reply, :ok, new_state}
+      [{{:root_span, span_data.span_id}, span}] ->
+        :ets.delete(@table, {:root_span, span_data.span_id})
+        :ets.insert(@table, {{:root_span, span_data.span_id}, span_data})
     end
   end
 
-  def handle_call({:remove_span, span_id}, _from, state) do
-    new_state = %{
-      state
-      | root_spans: Map.delete(state.root_spans, span_id),
-        child_spans: Map.delete(state.child_spans, span_id)
-    }
+  def update_span(span_data) do
+    existing_spans = :ets.lookup(@table, span_data.parent_span_id)
+
+    :ets.delete(@table, span_data.parent_span_id)
+
+    Enum.each(existing_spans, fn {parent_span_id, span} ->
+      if span.span_id == span_data.span_id do
+        :ets.delete_object(@table, span)
+        :ets.insert(@table, {parent_span_id, span_data})
+      end
+    end)
+  end
 
-    {:reply, :ok, new_state}
+  def remove_span(span_id) do
+    :ets.delete(@table, {:root_span, span_id})
+    :ok
   end
 
-  def handle_call({:remove_child_spans, parent_span_id}, _from, state) do
-    new_state = %{state | child_spans: Map.delete(state.child_spans, parent_span_id)}
-    {:reply, :ok, new_state}
+  def remove_child_spans(parent_span_id) do
+    :ets.delete(@table, parent_span_id)
+    :ok
   end
 end