fix: fix tenant cache handling (#1665)

Author: edgurgel

lib/realtime/api.ex

@@ -119,6 +119,14 @@ defmodule Realtime.Api do
       %Tenant{}
       |> Tenant.changeset(attrs)
       |> Repo.insert()
+      |> case do
+        {:ok, tenant} ->
+          Cache.global_cache_update(tenant)
+          {:ok, tenant}
+
+        error ->
+          error
+      end
     else
       call(:create_tenant, [attrs], tenant_id)
     end
@@ -144,7 +152,7 @@ defmodule Realtime.Api do
 
     case updated do
       {:ok, tenant} ->
-        maybe_invalidate_cache(changeset)
+        maybe_update_cache(tenant, changeset)
         maybe_trigger_disconnect(changeset)
         maybe_restart_db_connection(changeset)
         Logger.debug("Tenant updated: #{inspect(tenant, pretty: true)}")
@@ -216,7 +224,12 @@ defmodule Realtime.Api do
       tenant
       |> Tenant.changeset(%{migrations_ran: count})
       |> Repo.update()
-      |> tap(fn _ -> Cache.distributed_invalidate_tenant_cache(external_id) end)
+      |> tap(fn result ->
+        case result do
+          {:ok, tenant} -> Cache.global_cache_update(tenant)
+          _ -> :ok
+        end
+      end)
     else
       call(:update_migrations_ran, [external_id, count], external_id)
     end
@@ -241,12 +254,11 @@ defmodule Realtime.Api do
     |> Map.put(:events_per_second_now, current)
   end
 
-  defp maybe_invalidate_cache(%Changeset{changes: changes, valid?: true, data: %{external_id: external_id}})
-       when changes != %{} do
-    Tenants.Cache.distributed_invalidate_tenant_cache(external_id)
+  defp maybe_update_cache(tenant, %Changeset{changes: changes, valid?: true}) when changes != %{} do
+    Tenants.Cache.global_cache_update(tenant)
   end
 
-  defp maybe_invalidate_cache(_changeset), do: nil
+  defp maybe_update_cache(_tenant, _changeset), do: :ok
 
   defp maybe_trigger_disconnect(%Changeset{data: %{external_id: external_id}} = changeset)
        when requires_disconnect(changeset) do

lib/realtime/context_cache.ex

@@ -5,6 +5,7 @@ defmodule Realtime.Tenants.Cache do
   require Cachex.Spec
   require Logger
 
+  alias Realtime.GenRpc
   alias Realtime.Tenants
 
   def child_spec(_) do
@@ -16,32 +17,42 @@ defmodule Realtime.Tenants.Cache do
     }
   end
 
-  def get_tenant_by_external_id(keyword), do: apply_repo_fun(__ENV__.function, [keyword])
+  def get_tenant_by_external_id(tenant_id) do
+    case Cachex.fetch(__MODULE__, cache_key(tenant_id), fn _key ->
+           case Tenants.get_tenant_by_external_id(tenant_id) do
+             nil -> {:ignore, nil}
+             tenant -> {:commit, tenant}
+           end
+         end) do
+      {:commit, value} -> value
+      {:ok, value} -> value
+      {:ignore, value} -> value
+    end
+  end
+
+  defp cache_key(tenant_id), do: {:get_tenant_by_external_id, tenant_id}
 
   @doc """
   Invalidates the cache for a tenant in the local node
   """
-  def invalidate_tenant_cache(tenant_id), do: Cachex.del(__MODULE__, {{:get_tenant_by_external_id, 1}, [tenant_id]})
+  def invalidate_tenant_cache(tenant_id), do: Cachex.del(__MODULE__, cache_key(tenant_id))
+
+  def distributed_invalidate_tenant_cache(tenant_id) when is_binary(tenant_id) do
+    GenRpc.multicast(__MODULE__, :invalidate_tenant_cache, [tenant_id])
+  end
 
   @doc """
-  Broadcasts a message to invalidate the tenant cache to all connected nodes
+  Update the cache for a tenant
   """
-  @spec distributed_invalidate_tenant_cache(String.t()) :: boolean()
-  def distributed_invalidate_tenant_cache(tenant_id) when is_binary(tenant_id) do
-    nodes = [Node.self() | Node.list()]
-    results = :erpc.multicall(nodes, __MODULE__, :invalidate_tenant_cache, [tenant_id], 1000)
-
-    results
-    |> Enum.map(fn
-      {res, _} ->
-        res
-
-      exception ->
-        Logger.error("Failed to invalidate tenant cache: #{inspect(exception)}")
-        :error
-    end)
-    |> Enum.all?(&(&1 == :ok))
+  def update_cache(tenant) do
+    Cachex.put(__MODULE__, cache_key(tenant.external_id), tenant)
   end
 
-  defp apply_repo_fun(arg1, arg2), do: Realtime.ContextCache.apply_fun(Tenants, arg1, arg2)
+  @doc """
+  Update the cache for a tenant in all nodes
+  """
+  @spec global_cache_update(Realtime.Api.Tenant.t()) :: :ok
+  def global_cache_update(tenant) do
+    GenRpc.multicast(__MODULE__, :update_cache, [tenant])
+  end
 end

lib/realtime/tenants/cache.ex

@@ -195,7 +195,7 @@ defmodule RealtimeWeb.TenantController do
     with %Tenant{} = tenant <- Api.get_tenant_by_external_id(tenant_id, use_replica: false),
          _ <- Tenants.suspend_tenant_by_external_id(tenant_id),
          true <- Api.delete_tenant_by_external_id(tenant_id),
-         true <- Cache.distributed_invalidate_tenant_cache(tenant_id),
+         :ok <- Cache.distributed_invalidate_tenant_cache(tenant_id),
          :ok <- PostgresCdc.stop_all(tenant, stop_all_timeout),
          :ok <- Database.replication_slot_teardown(tenant) do
       send_resp(conn, 204, "")

lib/realtime_web/controllers/tenant_controller.ex

@@ -4,7 +4,7 @@ defmodule Realtime.MixProject do
   def project do
     [
       app: :realtime,
-      version: "2.69.1",
+      version: "2.69.2",
       elixir: "~> 1.18",
       elixirc_paths: elixirc_paths(Mix.env()),
       start_permanent: Mix.env() == :prod,

mix.exs

@@ -419,11 +419,7 @@ defmodule Realtime.Integration.RtChannelTest do
     test "broadcast to another tenant does not get mixed up", %{tenant: tenant, serializer: serializer} do
       other_tenant = Containers.checkout_tenant(run_migrations: true)
 
-      Cachex.put!(
-        Realtime.Tenants.Cache,
-        {{:get_tenant_by_external_id, 1}, [other_tenant.external_id]},
-        {:cached, other_tenant}
-      )
+      Realtime.Tenants.Cache.update_cache(other_tenant)
 
       {socket, _} = get_connection(tenant, serializer)
       config = %{broadcast: %{self: false}, private: false}
@@ -2406,7 +2402,7 @@ defmodule Realtime.Integration.RtChannelTest do
       |> Realtime.Api.Tenant.changeset(%{limit => value})
       |> Realtime.Repo.update!()
 
-    Cachex.put!(Realtime.Tenants.Cache, {{:get_tenant_by_external_id, 1}, [tenant.external_id]}, {:cached, tenant})
+    Realtime.Tenants.Cache.update_cache(tenant)
   end
 
   defp assert_process_down(pid, timeout \\ 1000) do

test/integration/rt_channel_test.exs

@@ -55,6 +55,10 @@ defmodule Realtime.ApiTest do
 
       external_id = random_string()
 
+      expect(Realtime.Tenants.Cache, :global_cache_update, fn tenant ->
+        assert tenant.external_id == external_id
+      end)
+
       valid_attrs = %{
         external_id: external_id,
         name: external_id,
@@ -89,6 +93,7 @@ defmodule Realtime.ApiTest do
     end
 
     test "invalid data returns error changeset" do
+      reject(&Realtime.Tenants.Cache.global_cache_update/1)
       assert {:error, %Ecto.Changeset{}} = Api.create_tenant(%{external_id: nil, jwt_secret: nil, name: nil})
     end
   end
@@ -197,10 +202,14 @@ defmodule Realtime.ApiTest do
 
     test "valid data and tenant data change will not restart the database connection" do
       tenant = Containers.checkout_tenant(run_migrations: true)
-      expect(Realtime.Tenants.Cache, :distributed_invalidate_tenant_cache, fn _ -> :ok end)
+
+      expect(Realtime.Tenants.Cache, :global_cache_update, fn tenant ->
+        assert tenant.max_concurrent_users == 101
+      end)
+
       {:ok, old_pid} = Connect.lookup_or_start_connection(tenant.external_id)
 
-      assert {:ok, %Tenant{}} = Api.update_tenant_by_external_id(tenant.external_id, %{max_concurrent_users: 100})
+      assert {:ok, %Tenant{}} = Api.update_tenant_by_external_id(tenant.external_id, %{max_concurrent_users: 101})
       refute_receive {:DOWN, _, :process, ^old_pid, :shutdown}, 500
       assert Process.alive?(old_pid)
       assert {:ok, new_pid} = Connect.lookup_or_start_connection(tenant.external_id)
@@ -241,12 +250,15 @@ defmodule Realtime.ApiTest do
     end
 
     test "valid data and change to tenant data will refresh cache", %{tenants: [tenant | _]} do
+      expect(Realtime.Tenants.Cache, :global_cache_update, fn tenant ->
+        assert tenant.name == "new_name"
+      end)
+
       assert {:ok, %Tenant{}} = Api.update_tenant_by_external_id(tenant.external_id, %{name: "new_name"})
-      assert %Tenant{name: "new_name"} = Realtime.Tenants.Cache.get_tenant_by_external_id(tenant.external_id)
     end
 
     test "valid data and no changes to tenant will not refresh cache", %{tenants: [tenant | _]} do
-      reject(&Realtime.Tenants.Cache.distributed_invalidate_tenant_cache/1)
+      reject(&Realtime.Tenants.Cache.global_cache_update/1)
       assert {:ok, %Tenant{}} = Api.update_tenant_by_external_id(tenant.external_id, %{name: tenant.name})
     end
   end
@@ -367,8 +379,13 @@ defmodule Realtime.ApiTest do
   describe "update_migrations_ran/1" do
     test "updates migrations_ran to the count of all migrations" do
       tenant = tenant_fixture(%{migrations_ran: 0})
-      Api.update_migrations_ran(tenant.external_id, 1)
-      tenant = Repo.reload!(tenant)
+
+      expect(Realtime.Tenants.Cache, :global_cache_update, fn tenant ->
+        assert tenant.migrations_ran == 1
+        :ok
+      end)
+
+      assert {:ok, tenant} = Api.update_migrations_ran(tenant.external_id, 1)
       assert tenant.migrations_ran == 1
     end
   end

test/realtime/api_test.exs

@@ -282,7 +282,7 @@ defmodule Realtime.Tenants.AuthorizationTest do
   def rls_context(context) do
     tenant = Containers.checkout_tenant(run_migrations: true)
     # Warm cache to avoid Cachex and Ecto.Sandbox ownership issues
-    Cachex.put!(Realtime.Tenants.Cache, {{:get_tenant_by_external_id, 1}, [tenant.external_id]}, {:cached, tenant})
+    Realtime.Tenants.Cache.update_cache(tenant)
 
     {:ok, db_conn} = Database.connect(tenant, "realtime_test", :stop)
     topic = context[:topic] || random_string()
@@ -323,6 +323,6 @@ defmodule Realtime.Tenants.AuthorizationTest do
     {:ok, tenant} = Realtime.Api.update_tenant_by_external_id(tenant.external_id, %{extensions: extensions})
 
     # Warm cache to avoid Cachex and Ecto.Sandbox ownership issues
-    Cachex.put!(Realtime.Tenants.Cache, {{:get_tenant_by_external_id, 1}, [tenant.external_id]}, {:cached, tenant})
+    Realtime.Tenants.Cache.update_cache(tenant)
   end
 end

test/realtime/tenants/authorization_test.exs

@@ -16,7 +16,7 @@ defmodule Realtime.Tenants.BatchBroadcastTest do
 
   setup do
     tenant = Containers.checkout_tenant(run_migrations: true)
-    Cachex.put!(Realtime.Tenants.Cache, {{:get_tenant_by_external_id, 1}, [tenant.external_id]}, {:cached, tenant})
+    Realtime.Tenants.Cache.update_cache(tenant)
     {:ok, tenant: tenant}
   end
 

test/realtime/tenants/batch_broadcast_test.exs

@@ -21,6 +21,12 @@ defmodule Realtime.Tenants.CacheTest do
       assert %Api.Tenant{name: "new name"} = Tenants.get_tenant_by_external_id(external_id)
       assert %Api.Tenant{name: "tenant"} = Cache.get_tenant_by_external_id(external_id)
     end
+
+    test "does not cache when tenant is not found" do
+      assert Cache.get_tenant_by_external_id("not found") == nil
+
+      assert Cachex.exists?(Cache, {:get_tenant_by_external_id, "not found"}) == {:ok, false}
+    end
   end
 
   describe "invalidate_tenant_cache/1" do
@@ -40,6 +46,18 @@ defmodule Realtime.Tenants.CacheTest do
     end
   end
 
+  describe "update_cache/1" do
+    test "updates the cache given a tenant", %{tenant: tenant} do
+      external_id = tenant.external_id
+      assert %Api.Tenant{name: "tenant"} = Cache.get_tenant_by_external_id(external_id)
+      # Update a tenant
+      updated_tenant = %{tenant | name: "updated name"}
+      # Update cache
+      Cache.update_cache(updated_tenant)
+      assert %Api.Tenant{name: "updated name"} = Cache.get_tenant_by_external_id(external_id)
+    end
+  end
+
   describe "distributed_invalidate_tenant_cache/1" do
     setup do
       {:ok, node} = Clustered.start()
@@ -53,30 +71,61 @@ defmodule Realtime.Tenants.CacheTest do
       dummy_name = random_string()
 
       # Ensure cache has the values
-      Cachex.put!(
-        Realtime.Tenants.Cache,
-        {{:get_tenant_by_external_id, 1}, [external_id]},
-        {:cached, %{tenant | name: dummy_name}}
-      )
-
-      Rpc.enhanced_call(node, Cachex, :put!, [
-        Realtime.Tenants.Cache,
-        {{:get_tenant_by_external_id, 1}, [external_id]},
-        {:cached, %{tenant | name: dummy_name}}
-      ])
+      Realtime.Tenants.Cache.update_cache(%{tenant | name: dummy_name})
+
+      Rpc.enhanced_call(node, Realtime.Tenants.Cache, :update_cache, [%{tenant | name: dummy_name}])
 
       # Cache showing old value
-      assert %Api.Tenant{name: ^dummy_name} = Cache.get_tenant_by_external_id(external_id)
-      assert %Api.Tenant{name: ^dummy_name} = Rpc.enhanced_call(node, Cache, :get_tenant_by_external_id, [external_id])
+      assert {:ok, %Api.Tenant{name: ^dummy_name}} = Cachex.get(Cache, {:get_tenant_by_external_id, external_id})
+
+      assert {:ok, %Api.Tenant{name: ^dummy_name}} =
+               Rpc.enhanced_call(node, Cachex, :get, [Cache, {:get_tenant_by_external_id, external_id}])
 
       # Invalidate cache
-      assert true = Cache.distributed_invalidate_tenant_cache(external_id)
+      assert :ok = Cache.distributed_invalidate_tenant_cache(external_id)
 
+      # wait for cache to be invalidated in both nodes
+      Process.sleep(200)
       # Cache showing new value
       assert %Api.Tenant{name: ^expected_name} = Cache.get_tenant_by_external_id(external_id)
 
       assert %Api.Tenant{name: ^expected_name} =
                Rpc.enhanced_call(node, Cache, :get_tenant_by_external_id, [external_id])
     end
   end
+
+  describe "global_cache_update/1" do
+    setup do
+      {:ok, node} = Clustered.start()
+      %{node: node}
+    end
+
+    test "update the cache given a tenant_id", %{node: node} do
+      external_id = "dev_tenant"
+      %Api.Tenant{name: expected_name} = tenant = Tenants.get_tenant_by_external_id(external_id)
+
+      dummy_name = random_string()
+
+      # Ensure cache has the values
+      Realtime.Tenants.Cache.update_cache(%{tenant | name: dummy_name})
+
+      Rpc.enhanced_call(node, Cache, :update_cache, [%{tenant | name: dummy_name}])
+
+      # Cache showing old value
+      assert %Api.Tenant{name: ^dummy_name} = Cache.get_tenant_by_external_id(external_id)
+      assert %Api.Tenant{name: ^dummy_name} = Rpc.enhanced_call(node, Cache, :get_tenant_by_external_id, [external_id])
+
+      # Update cache
+      assert :ok = Cache.global_cache_update(tenant)
+
+      # wait for cache to be updated in both nodes
+      Process.sleep(200)
+
+      # Cache showing new value
+      assert {:ok, %Api.Tenant{name: ^expected_name}} = Cachex.get(Cache, {:get_tenant_by_external_id, external_id})
+
+      assert {:ok, %Api.Tenant{name: ^expected_name}} =
+               Rpc.enhanced_call(node, Cachex, :get, [Cache, {:get_tenant_by_external_id, external_id}])
+    end
+  end
 end