Add sweeping of expired spans

Author: solnic

lib/sentry/opentelemetry/span_storage.ex

@@ -2,64 +2,76 @@ defmodule Sentry.OpenTelemetry.SpanStorage do
   @moduledoc false
   use GenServer
 
+  require Logger
+
   @table :span_storage
+  @cleanup_interval :timer.minutes(5)
+  @span_ttl :timer.minutes(30)
 
   @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)
+    GenServer.start_link(__MODULE__, opts, name: name)
   end
 
   @impl true
-  def init(nil) do
+  def init(opts) do
     _table =
       if :ets.whereis(@table) == :undefined do
         :ets.new(@table, [:named_table, :public, :bag])
       end
 
-    {:ok, :no_state}
+    cleanup_interval = Keyword.get(opts, :cleanup_interval, @cleanup_interval)
+    schedule_cleanup(cleanup_interval)
+
+    {:ok, %{cleanup_interval: cleanup_interval}}
   end
 
   def store_span(span_data) when span_data.parent_span_id == nil do
+    stored_at = System.system_time(:second)
+
     case :ets.lookup(@table, {:root_span, span_data.span_id}) do
-      [] -> :ets.insert(@table, {{:root_span, span_data.span_id}, span_data})
+      [] -> :ets.insert(@table, {{:root_span, span_data.span_id}, {span_data, stored_at}})
       _ -> :ok
     end
   end
 
   def store_span(span_data) do
-    _ = :ets.insert(@table, {span_data.parent_span_id, span_data})
+    stored_at = System.system_time(:second)
+    _ = :ets.insert(@table, {span_data.parent_span_id, {span_data, stored_at}})
   end
 
   def get_root_span(span_id) do
     case :ets.lookup(@table, {:root_span, span_id}) do
-      [{{:root_span, ^span_id}, span}] -> span
+      [{{:root_span, ^span_id}, {span, _stored_at}}] -> span
       [] -> nil
     end
   end
 
   def get_child_spans(parent_span_id) do
     :ets.lookup(@table, parent_span_id)
-    |> Enum.map(fn {_parent_id, span} -> span end)
+    |> Enum.map(fn {_parent_id, {span, _stored_at}} -> span end)
   end
 
   def update_span(span_data) do
+    stored_at = System.system_time(:second)
+
     if 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})
+          :ets.insert(@table, {{:root_span, span_data.span_id}, {span_data, stored_at}})
 
         _ ->
           :ets.delete(@table, {:root_span, span_data.span_id})
-          :ets.insert(@table, {{:root_span, span_data.span_id}, span_data})
+          :ets.insert(@table, {{:root_span, span_data.span_id}, {span_data, stored_at}})
       end
     else
       existing_spans = :ets.lookup(@table, span_data.parent_span_id)
 
-      Enum.each(existing_spans, fn {parent_id, span} ->
+      Enum.each(existing_spans, fn {parent_id, {span, stored_at}} ->
         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})
+          :ets.delete_object(@table, {parent_id, {span, stored_at}})
+          :ets.insert(@table, {span_data.parent_span_id, {span_data, stored_at}})
         end
       end)
     end
@@ -69,7 +81,9 @@ defmodule Sentry.OpenTelemetry.SpanStorage do
 
   def remove_span(span_id) do
     case get_root_span(span_id) do
-      nil -> :ok
+      nil ->
+        :ok
+
       _root_span ->
         :ets.delete(@table, {:root_span, span_id})
         remove_child_spans(span_id)
@@ -80,4 +94,44 @@ defmodule Sentry.OpenTelemetry.SpanStorage do
     :ets.delete(@table, parent_span_id)
     :ok
   end
+
+  @impl true
+  def handle_info(:cleanup_stale_spans, state) do
+    cleanup_stale_spans()
+    schedule_cleanup(state.cleanup_interval)
+    {:noreply, state}
+  end
+
+  defp schedule_cleanup(interval) do
+    Process.send_after(self(), :cleanup_stale_spans, interval)
+  end
+
+  defp cleanup_stale_spans do
+    now = System.system_time(:second)
+    cutoff_time = now - @span_ttl
+
+    :ets.match_object(@table, {{:root_span, :_}, {:_, :_}})
+    |> Enum.each(fn {{:root_span, span_id}, {_span, stored_at}} ->
+      if stored_at < cutoff_time do
+        Logger.debug("Cleaning up stale root span: #{span_id}")
+        remove_span(span_id)
+      end
+    end)
+
+    :ets.match_object(@table, {:_, {:_, :_}})
+    |> Enum.each(fn {parent_id, {span, stored_at}} = object ->
+      cond do
+        get_root_span(parent_id) != nil and stored_at < cutoff_time ->
+          Logger.debug("Cleaning up stale child span: #{span.span_id}")
+          :ets.delete_object(@table, object)
+
+        get_root_span(parent_id) == nil and stored_at < cutoff_time ->
+          Logger.debug("Cleaning up stale orphaned child span: #{span.span_id}")
+          :ets.delete_object(@table, object)
+
+        true ->
+          :ok
+      end
+    end)
+  end
 end

test/sentry/opentelemetry/span_storage_test.exs

@@ -95,14 +95,11 @@ defmodule Sentry.OpenTelemetry.SpanStorageTest do
       SpanStorage.store_span(child_span1)
       SpanStorage.store_span(child_span2)
 
-      # Verify initial state
       assert root_span == SpanStorage.get_root_span("root123")
       assert length(SpanStorage.get_child_spans("root123")) == 2
 
-      # Remove root span should remove everything
       SpanStorage.remove_span("root123")
 
-      # Verify everything is removed
       assert nil == SpanStorage.get_root_span("root123")
       assert [] == SpanStorage.get_child_spans("root123")
     end
@@ -218,4 +215,147 @@ defmodule Sentry.OpenTelemetry.SpanStorageTest do
     assert nil == SpanStorage.get_root_span("root123")
     assert [] == SpanStorage.get_child_spans("root123")
   end
+
+  describe "stale span cleanup" do
+    test "cleans up stale spans" do
+      start_supervised!({SpanStorage, cleanup_interval: 100, name: :cleanup_test})
+
+      root_span = %SpanRecord{
+        span_id: "stale_root",
+        parent_span_id: nil,
+        trace_id: "trace123",
+        name: "stale_root_span"
+      }
+
+      child_span = %SpanRecord{
+        span_id: "stale_child",
+        parent_span_id: "stale_root",
+        trace_id: "trace123",
+        name: "stale_child_span"
+      }
+
+      old_time = System.system_time(:second) - :timer.minutes(31)
+      :ets.insert(:span_storage, {{:root_span, "stale_root"}, {root_span, old_time}})
+      :ets.insert(:span_storage, {"stale_root", {child_span, old_time}})
+
+      fresh_root_span = %SpanRecord{
+        span_id: "fresh_root",
+        parent_span_id: nil,
+        trace_id: "trace123",
+        name: "fresh_root_span"
+      }
+
+      SpanStorage.store_span(fresh_root_span)
+
+      Process.sleep(200)
+
+      assert nil == SpanStorage.get_root_span("stale_root")
+      assert [] == SpanStorage.get_child_spans("stale_root")
+
+      assert SpanStorage.get_root_span("fresh_root")
+    end
+
+    test "cleans up orphaned child spans" do
+      start_supervised!({SpanStorage, cleanup_interval: 100, name: :cleanup_test})
+
+      child_span = %SpanRecord{
+        span_id: "stale_child",
+        parent_span_id: "non_existent_parent",
+        trace_id: "trace123",
+        name: "stale_child_span"
+      }
+
+      old_time = System.system_time(:second) - :timer.minutes(31)
+      :ets.insert(:span_storage, {"non_existent_parent", {child_span, old_time}})
+
+      Process.sleep(200)
+
+      assert [] == SpanStorage.get_child_spans("non_existent_parent")
+    end
+
+    test "cleans up expired root spans with all their children regardless of child timestamps" do
+      start_supervised!({SpanStorage, cleanup_interval: 100, name: :cleanup_test})
+
+      root_span = %SpanRecord{
+        span_id: "root123",
+        parent_span_id: nil,
+        trace_id: "trace123",
+        name: "root_span"
+      }
+
+      old_child = %SpanRecord{
+        span_id: "old_child",
+        parent_span_id: "root123",
+        trace_id: "trace123",
+        name: "old_child_span"
+      }
+
+      fresh_child = %SpanRecord{
+        span_id: "fresh_child",
+        parent_span_id: "root123",
+        trace_id: "trace123",
+        name: "fresh_child_span"
+      }
+
+      old_time = System.system_time(:second) - :timer.minutes(31)
+      :ets.insert(:span_storage, {{:root_span, "root123"}, {root_span, old_time}})
+
+      :ets.insert(:span_storage, {"root123", {old_child, old_time}})
+      SpanStorage.store_span(fresh_child)
+
+      Process.sleep(200)
+
+      assert nil == SpanStorage.get_root_span("root123")
+      assert [] == SpanStorage.get_child_spans("root123")
+    end
+
+    test "handles mixed expiration times in child spans" do
+      start_supervised!({SpanStorage, cleanup_interval: 100, name: :cleanup_test})
+
+      root_span = %SpanRecord{
+        span_id: "root123",
+        parent_span_id: nil,
+        trace_id: "trace123",
+        name: "root_span"
+      }
+
+      old_child1 = %SpanRecord{
+        span_id: "old_child1",
+        parent_span_id: "root123",
+        trace_id: "trace123",
+        name: "old_child_span_1"
+      }
+
+      old_child2 = %SpanRecord{
+        span_id: "old_child2",
+        parent_span_id: "root123",
+        trace_id: "trace123",
+        name: "old_child_span_2"
+      }
+
+      fresh_child = %SpanRecord{
+        span_id: "fresh_child",
+        parent_span_id: "root123",
+        trace_id: "trace123",
+        name: "fresh_child_span"
+      }
+
+      SpanStorage.store_span(root_span)
+
+      old_time = System.system_time(:second) - :timer.minutes(31)
+      :ets.insert(:span_storage, {"root123", {old_child1, old_time}})
+      :ets.insert(:span_storage, {"root123", {old_child2, old_time}})
+
+      SpanStorage.store_span(fresh_child)
+
+      Process.sleep(200)
+
+      assert root_span == SpanStorage.get_root_span("root123")
+      children = SpanStorage.get_child_spans("root123")
+      assert length(children) == 1
+      assert fresh_child in children
+      refute old_child1 in children
+      refute old_child2 in children
+    end
+  end
 end