fix(flags): track snapshot access with atomics instead of a process

.sampo/changesets/evaluations-atomics-tracker.md

@@ -0,0 +1,10 @@
+---
+hex/posthog: patch
+---
+
+Fix a process leak in feature flag snapshots. Snapshots no longer start a process to track accessed flags, so they need no cleanup and can be used from any process.
+
+Behaviour changes:
+
+- `PostHog.FeatureFlags.Evaluations`'s `accessed_pid` field is deprecated and always `nil`. Remove any `Agent.stop/1` call on it.
+- `PostHog.FeatureFlags.Evaluations.accessed/1` no longer returns keys that were read but are missing from the snapshot.

lib/posthog/feature_flags/evaluations.ex

@@ -7,11 +7,10 @@ defmodule PostHog.FeatureFlags.Evaluations do
   `PostHog.FeatureFlags.evaluate_flags/2` and lets you branch on multiple flags
   and enrich captured events from the same frozen result set.
 
-  Each snapshot owns a small `Agent` linked to the calling process that tracks
-  which flags were accessed via `enabled?/2` and `get_flag/2`. The Agent exits
-  with the calling process, whatever the exit reason — no manual cleanup is
-  required. A snapshot is therefore only usable while the process that created
-  it is alive; don't hand it to a process that outlives its creator.
+  Each snapshot tracks which flags were accessed via `enabled?/2` and
+  `get_flag/2` in an `:atomics` array that is garbage-collected together with
+  the snapshot. No process is started and no manual cleanup is required, and a
+  snapshot can be shared with other processes.
 
   ## Querying
 
@@ -69,7 +68,14 @@ defmodule PostHog.FeatureFlags.Evaluations do
     `errorsWhileComputingFlags`. When `true`, every event fired from this
     snapshot includes `errors_while_computing_flags` in its
     `$feature_flag_error` property.
-  - `:accessed_pid` - pid of the Agent tracking accessed keys.
+  - `:accessed` - internal. Its contents may change in any release without
+    notice. Don't read or match on it; use `accessed/1`.
+  - `:accessed_pid` - internal and deprecated. Always `nil`. It held the pid of
+    an internal Agent that tracked accessed keys and was never meant to be
+    used directly. Snapshots no longer start a process, so there is nothing to
+    stop or clean up, and the process leak this Agent caused no longer
+    happens. Kept only so existing code that reads or matches the field still
+    compiles. Don't use it; it will be removed in the next major version.
   """
   @type t :: %__MODULE__{
           supervisor_name: PostHog.supervisor_name(),
@@ -79,14 +85,21 @@ defmodule PostHog.FeatureFlags.Evaluations do
           request_id: String.t() | nil,
           evaluated_at: integer() | nil,
           errors_while_computing: boolean(),
-          accessed_pid: pid()
+          accessed: accessed(),
+          accessed_pid: nil
         }
 
-  @enforce_keys [:supervisor_name, :distinct_id, :flags, :accessed_pid]
+  @typedoc false
+  @opaque accessed :: {%{String.t() => pos_integer()}, :atomics.atomics_ref() | nil}
+
+  @enforce_keys [:supervisor_name, :distinct_id, :flags, :accessed]
   defstruct [
     :supervisor_name,
     :distinct_id,
     :flags,
+    :accessed,
+    # Deprecated and always nil. Kept for compatibility with code written for
+    # the old Agent-based tracker.
     :accessed_pid,
     :request_id,
     :evaluated_at,
@@ -111,7 +124,7 @@ defmodule PostHog.FeatureFlags.Evaluations do
       evaluated_at: Map.get(body, "evaluatedAt"),
       errors_while_computing: Map.get(body, "errorsWhileComputingFlags") == true,
       groups: normalize_groups(Map.get(body, :groups, Map.get(body, "groups", %{}))),
-      accessed_pid: start_accessed_agent()
+      accessed: new_accessed(flags)
     }
   end
 
@@ -134,7 +147,7 @@ defmodule PostHog.FeatureFlags.Evaluations do
         Map.get(metadata, :errors_while_computing, Map.get(metadata, "errorsWhileComputingFlags")) ==
           true,
       groups: normalize_groups(Map.get(metadata, :groups, Map.get(metadata, "groups", %{}))),
-      accessed_pid: start_accessed_agent()
+      accessed: new_accessed(results)
     }
   end
 
@@ -145,7 +158,7 @@ defmodule PostHog.FeatureFlags.Evaluations do
       supervisor_name: supervisor_name,
       distinct_id: "",
       flags: %{},
-      accessed_pid: start_accessed_agent()
+      accessed: new_accessed(%{})
     }
   end
 
@@ -206,14 +219,19 @@ defmodule PostHog.FeatureFlags.Evaluations do
   def keys(%__MODULE__{flags: flags}), do: flags |> Map.keys() |> Enum.sort()
 
   @doc """
-  Returns the sorted list of keys accessed via `enabled?/2` or `get_flag/2` on
-  this snapshot.
+  Returns the sorted list of keys present in this snapshot that were accessed
+  via `enabled?/2` or `get_flag/2`.
 
-  Includes keys that were accessed but absent from the snapshot.
+  Keys that were accessed but are absent from the snapshot are not included.
   """
   @spec accessed(t()) :: [String.t()]
-  def accessed(%__MODULE__{accessed_pid: pid}) do
-    pid |> Agent.get(& &1) |> MapSet.to_list() |> Enum.sort()
+  def accessed(%__MODULE__{accessed: {_indexes, nil}}), do: []
+
+  def accessed(%__MODULE__{accessed: {indexes, ref}}) do
+    indexes
+    |> Enum.filter(fn {_key, index} -> :atomics.get(ref, index) == 1 end)
+    |> Enum.map(fn {key, _index} -> key end)
+    |> Enum.sort()
   end
 
   @doc """
@@ -229,9 +247,7 @@ defmodule PostHog.FeatureFlags.Evaluations do
   """
   @spec only_accessed(t()) :: t()
   def only_accessed(%__MODULE__{flags: flags} = snapshot) do
-    accessed_set = MapSet.new(accessed(snapshot))
-    keep = flags |> Map.keys() |> Enum.filter(&MapSet.member?(accessed_set, &1))
-    clone_with(snapshot, Map.take(flags, keep))
+    clone_with(snapshot, Map.take(flags, accessed(snapshot)))
   end
 
   @doc """
@@ -275,25 +291,24 @@ defmodule PostHog.FeatureFlags.Evaluations do
     end
   end
 
-  defp start_accessed_agent do
-    # Trap exits so the Agent also stops when the caller exits with :normal.
-    # Without it, a :normal exit signal is ignored and the Agent is orphaned.
-    {:ok, pid} =
-      Agent.start_link(fn ->
-        Logger.metadata(posthog_skip_capture: true)
-        Process.flag(:trap_exit, true)
-        MapSet.new()
-      end)
+  # :atomics are garbage-collected with the snapshot, unlike a process-based
+  # tracker, which would live until its owner exits.
+  defp new_accessed(flags) when map_size(flags) == 0, do: {%{}, nil}
 
-    pid
+  defp new_accessed(flags) do
+    indexes = flags |> Map.keys() |> Enum.with_index(1) |> Map.new()
+    {indexes, :atomics.new(map_size(indexes), signed: false)}
   end
 
-  defp record_access(%__MODULE__{accessed_pid: pid}, key) do
-    Agent.update(pid, &MapSet.put(&1, key))
+  defp record_access(%__MODULE__{accessed: {indexes, ref}}, key) do
+    case Map.fetch(indexes, key) do
+      {:ok, index} -> :atomics.put(ref, index, 1)
+      :error -> :ok
+    end
   end
 
   defp clone_with(%__MODULE__{} = snapshot, flags) do
-    %{snapshot | flags: flags, accessed_pid: start_accessed_agent()}
+    %{snapshot | flags: flags, accessed: new_accessed(flags)}
   end
 
   defp fetch_and_log(%__MODULE__{flags: flags} = snapshot, key) do

lib/posthog/open_feature.ex

@@ -167,11 +167,7 @@ if Code.ensure_loaded?(OpenFeature.Provider) do
 
           case FeatureFlags.evaluate_flags(provider.supervisor_name, body) do
             {:ok, snapshot} ->
-              try do
-                fetch_result(provider, snapshot, key)
-              after
-                Agent.stop(snapshot.accessed_pid)
-              end
+              fetch_result(provider, snapshot, key)
 
             {:error, reason} ->
               {:error, :unexpected_error, to_exception(reason)}

test/posthog/feature_flags/evaluations_test.exs

@@ -52,27 +52,62 @@ defmodule PostHog.FeatureFlags.EvaluationsTest do
     }
   end
 
-  describe "access tracker lifetime" do
-    for reason <- [:normal, :shutdown, {:shutdown, :closed}] do
-      @reason reason
-      test "stops when the creating process exits with #{inspect(reason)}" do
-        test_pid = self()
+  describe "access tracking resources" do
+    test "snapshots do not start or link processes to the caller" do
+      expect(API.Mock, :request, fn _client, :post, "/flags", _opts ->
+        {:ok, stub_flags_response()}
+      end)
 
-        {creator, creator_ref} =
-          spawn_monitor(fn ->
-            send(test_pid, {:snapshot, Evaluations.empty(PostHog)})
-            receive do: (:exit -> exit(@reason))
-          end)
+      {:links, links_before} = Process.info(self(), :links)
 
-        assert_receive {:snapshot, %Evaluations{accessed_pid: agent}}
-        agent_ref = Process.monitor(agent)
+      {:ok, snapshot} = FeatureFlags.evaluate_flags("foo")
+      Evaluations.enabled?(snapshot, "boolean-flag")
+      _ = Evaluations.only(snapshot, ["variant-flag"])
+      _ = Evaluations.only_accessed(snapshot)
+      _ = Evaluations.empty(PostHog)
 
-        send(creator, :exit)
+      assert Process.info(self(), :links) == {:links, links_before}
+    end
 
-        assert_receive {:DOWN, ^creator_ref, :process, ^creator, _}
-        assert_receive {:DOWN, ^agent_ref, :process, ^agent, _}
+    test "keeps the deprecated accessed_pid field as nil" do
+      expect(API.Mock, :request, fn _client, :post, "/flags", _opts ->
+        {:ok, stub_flags_response()}
+      end)
+
+      {:ok, snapshot} = FeatureFlags.evaluate_flags("foo")
+
+      for snapshot <- [
+            snapshot,
+            Evaluations.only(snapshot, ["boolean-flag"]),
+            Evaluations.only_accessed(snapshot),
+            Evaluations.empty(PostHog)
+          ] do
+        assert %Evaluations{accessed_pid: nil} = snapshot
       end
     end
+
+    test "snapshots stay usable after the creating process exits" do
+      expect(API.Mock, :request, fn _client, :post, "/flags", _opts ->
+        {:ok, stub_flags_response()}
+      end)
+
+      test_pid = self()
+
+      {creator, creator_ref} =
+        spawn_monitor(fn ->
+          receive do: (:go -> :ok)
+          {:ok, snapshot} = FeatureFlags.evaluate_flags("foo")
+          send(test_pid, {:snapshot, snapshot})
+        end)
+
+      allow(API.Mock, test_pid, creator)
+      send(creator, :go)
+      assert_receive {:snapshot, snapshot}
+      assert_receive {:DOWN, ^creator_ref, :process, ^creator, :normal}
+
+      assert Evaluations.enabled?(snapshot, "boolean-flag")
+      assert Evaluations.accessed(snapshot) == ["boolean-flag"]
+    end
   end
 
   describe "evaluate_flags/2" do
@@ -433,6 +468,8 @@ defmodule PostHog.FeatureFlags.EvaluationsTest do
       Evaluations.enabled?(snapshot, "boolean-flag")
       Evaluations.enabled?(snapshot, "missing-flag")
 
+      assert Evaluations.accessed(snapshot) == ["boolean-flag"]
+
       narrowed = Evaluations.only_accessed(snapshot)
       assert Evaluations.keys(narrowed) == ["boolean-flag"]
     end

test/posthog/handler_test.exs

@@ -307,50 +307,6 @@ defmodule PostHog.HandlerTest do
     assert [] = all_captured(supervisor_name)
   end
 
-  test "snapshot tracker exits do not duplicate their creator's exception", %{
-    handler_id: handler_id,
-    handler_ref: ref,
-    config: %{supervisor_name: supervisor_name}
-  } do
-    test_pid = self()
-
-    {creator, creator_ref} =
-      spawn_monitor(fn ->
-        {:ok, snapshot} =
-          PostHog.FeatureFlags.evaluate_flags(supervisor_name, %{
-            distinct_id: "snapshot-user",
-            flag_keys: []
-          })
-
-        filtered = PostHog.FeatureFlags.Evaluations.only(snapshot, [])
-        send(test_pid, {:trackers, [snapshot.accessed_pid, filtered.accessed_pid]})
-        receive do: (:raise -> raise "snapshot creator failed")
-      end)
-
-    assert_receive {:trackers, trackers}
-
-    for pid <- [creator | trackers] do
-      LoggerHandlerKit.Arrange.allow(self(), pid, handler_id)
-      PostHog.Test.allow(supervisor_name, self(), pid)
-    end
-
-    tracker_refs = for pid <- trackers, do: {pid, Process.monitor(pid)}
-    send(creator, :raise)
-
-    assert_receive {:DOWN, ^creator_ref, :process, ^creator, _}
-
-    for {pid, tracker_ref} <- tracker_refs do
-      assert_receive {:DOWN, ^tracker_ref, :process, ^pid, _}
-    end
-
-    for _ <- 1..3, do: LoggerHandlerKit.Assert.assert_logged(ref)
-
-    assert [%{event: "$exception", properties: properties}] = all_captured(supervisor_name)
-
-    assert [%{type: "RuntimeError", value: "snapshot creator failed"} | _] =
-             properties[:"$exception_list"]
-  end
-
   @tag config: [capture_level: :warning]
   test "logs with crash reason always captured", %{
     handler_ref: ref,

test/posthog/open_feature_test.exs

@@ -38,45 +38,21 @@ defmodule PostHog.OpenFeature.ProviderTest do
     end
   end
 
-  describe "snapshot cleanup" do
-    for events <- [true, false], present <- [true, false] do
-      @events events
-      @present present
-      test "releases access trackers with events=#{events} and flag present=#{present}" do
-        flags = if @present, do: %{"flag" => flag(%{"key" => "flag"})}, else: %{}
-
-        expect(API.Mock, :request, 10, fn _client, :post, "/flags", _opts ->
-          {:ok, %{status: 200, body: %{"flags" => flags}}}
-        end)
+  describe "snapshot cleanup in a process trapping exits" do
+    test "does not deliver tracker exit messages to the caller" do
+      expect(API.Mock, :request, 10, fn _client, :post, "/flags", _opts ->
+        {:ok, %{status: 200, body: %{"flags" => %{"flag" => flag(%{"key" => "flag"})}}}}
+      end)
 
-        provider = provider(send_feature_flag_events: @events)
-        {:links, before} = Process.info(self(), :links)
-
-        for _ <- 1..10 do
-          result = Provider.resolve_boolean_value(provider, "flag", false, @context)
-
-          if @present do
-            assert {:ok, %ResolutionDetails{value: true}} = result
-          else
-            assert {:error, :flag_not_found} = result
-          end
-        end
-
-        {:links, after_calls} = Process.info(self(), :links)
-        retained = after_calls -- before
-
-        retained
-        |> Enum.filter(fn pid ->
-          is_pid(pid) and
-            match?(
-              {PostHog.FeatureFlags.Evaluations, _, _},
-              :proc_lib.translate_initial_call(pid)
-            )
-        end)
-        |> Enum.each(&Agent.stop/1)
+      provider = provider()
+      Process.flag(:trap_exit, true)
 
-        assert retained == []
+      for _ <- 1..10 do
+        assert {:ok, %ResolutionDetails{value: true}} =
+                 Provider.resolve_boolean_value(provider, "flag", false, @context)
       end
+
+      refute_receive {:EXIT, _pid, _reason}, 100
     end
   end