fix: rate limit Connect.lookup_or_start_connection on error only (#1549)

Author: edgurgel

lib/realtime/tenants.ex

@@ -328,18 +328,18 @@ defmodule Realtime.Tenants do
     %RateCounter.Args{id: {:channel, :authorization_errors, external_id}, opts: opts}
   end
 
-  @connect_per_second_default 10
+  @connect_errors_per_second_default 10
   @doc "RateCounter arguments for counting connect per second."
-  @spec connect_per_second_rate(Tenant.t() | String.t()) :: RateCounter.Args.t()
-  def connect_per_second_rate(%Tenant{external_id: external_id}) do
-    connect_per_second_rate(external_id)
+  @spec connect_errors_per_second_rate(Tenant.t() | String.t()) :: RateCounter.Args.t()
+  def connect_errors_per_second_rate(%Tenant{external_id: external_id}) do
+    connect_errors_per_second_rate(external_id)
   end
 
-  def connect_per_second_rate(tenant_id) do
+  def connect_errors_per_second_rate(tenant_id) do
     opts = [
       max_bucket_len: 10,
       limit: [
-        value: @connect_per_second_default,
+        value: @connect_errors_per_second_default,
         measurement: :sum,
         log_fn: fn ->
           Logger.critical(

lib/realtime/tenants/connect.ex

@@ -57,7 +57,7 @@ defmodule Realtime.Tenants.Connect do
           | {:error, :connect_rate_limit_reached}
           | {:error, :rpc_error, term()}
   def lookup_or_start_connection(tenant_id, opts \\ []) when is_binary(tenant_id) do
-    rate_args = Tenants.connect_per_second_rate(tenant_id)
+    rate_args = Tenants.connect_errors_per_second_rate(tenant_id)
     RateCounter.new(rate_args)
 
     with {:ok, %{limit: %{triggered: false}}} <- RateCounter.get(rate_args),
@@ -68,8 +68,14 @@ defmodule Realtime.Tenants.Connect do
         {:error, :connect_rate_limit_reached}
 
       {:error, :tenant_database_connection_initializing} ->
-        GenCounter.add(rate_args.id)
-        call_external_node(tenant_id, opts)
+        case call_external_node(tenant_id, opts) do
+          {:ok, pid} ->
+            {:ok, pid}
+
+          error ->
+            GenCounter.add(rate_args.id)
+            error
+        end
 
       {:error, :initializing} ->
         {:error, :tenant_database_unavailable}

mix.exs

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

test/realtime/tenants/connect_test.exs

@@ -51,6 +51,36 @@ defmodule Realtime.Tenants.ConnectTest do
   end
 
   describe "handle cold start" do
+    test "multiple processes connecting calling Connect.connect", %{tenant: tenant} do
+      parent = self()
+
+      # Let's slow down Connect.connect so that multiple RPC calls are executed
+      stub(Connect, :connect, fn x, y, z ->
+        :timer.sleep(1000)
+        call_original(Connect, :connect, [x, y, z])
+      end)
+
+      connect = fn -> send(parent, Connect.lookup_or_start_connection(tenant.external_id)) end
+      # Let's call enough times to potentially trigger the Connect RateCounter
+
+      for _ <- 1..50, do: spawn(connect)
+
+      assert_receive({:ok, pid}, 1100)
+
+      for _ <- 1..49, do: assert_receive({:ok, ^pid})
+
+      # Does not trigger rate limit as connections eventually succeeded
+
+      {:ok, rate_counter} =
+        tenant.external_id
+        |> Tenants.connect_errors_per_second_rate()
+        |> Realtime.RateCounter.get()
+
+      assert rate_counter.sum == 0
+      assert rate_counter.avg == 0.0
+      assert rate_counter.limit.triggered == false
+    end
+
     test "multiple proccesses succeed together", %{tenant: tenant} do
       parent = self()
 

test/realtime_web/controllers/broadcast_controller_test.exs

@@ -272,7 +272,6 @@ defmodule RealtimeWeb.BroadcastControllerTest do
     } do
       request_events_key = Tenants.requests_per_second_key(tenant)
       broadcast_events_key = Tenants.events_per_second_key(tenant)
-      connect_events_key = Tenants.connect_per_second_rate(tenant).id
       expect(TenantBroadcaster, :pubsub_broadcast, 5, fn _, _, _, _ -> :ok end)
 
       messages_to_send =
@@ -291,10 +290,7 @@ defmodule RealtimeWeb.BroadcastControllerTest do
 
       GenCounter
       |> expect(:add, fn ^request_events_key -> :ok end)
-      |> expect(:add, length(messages), fn
-        ^broadcast_events_key -> :ok
-        ^connect_events_key -> :ok
-      end)
+      |> expect(:add, length(messages), fn ^broadcast_events_key -> :ok end)
 
       conn = post(conn, Routes.broadcast_path(conn, :broadcast), %{"messages" => messages})
 
@@ -330,7 +326,6 @@ defmodule RealtimeWeb.BroadcastControllerTest do
     } do
       request_events_key = Tenants.requests_per_second_key(tenant)
       broadcast_events_key = Tenants.events_per_second_key(tenant)
-      connect_events_key = Tenants.connect_per_second_rate(tenant).id
       expect(TenantBroadcaster, :pubsub_broadcast, 6, fn _, _, _, _ -> :ok end)
 
       channels =
@@ -359,10 +354,7 @@ defmodule RealtimeWeb.BroadcastControllerTest do
 
       GenCounter
       |> expect(:add, fn ^request_events_key -> :ok end)
-      |> expect(:add, length(messages), fn
-        ^broadcast_events_key -> :ok
-        ^connect_events_key -> :ok
-      end)
+      |> expect(:add, length(messages), fn ^broadcast_events_key -> :ok end)
 
       conn = post(conn, Routes.broadcast_path(conn, :broadcast), %{"messages" => messages})
 
@@ -416,7 +408,6 @@ defmodule RealtimeWeb.BroadcastControllerTest do
     } do
       request_events_key = Tenants.requests_per_second_key(tenant)
       broadcast_events_key = Tenants.events_per_second_key(tenant)
-      connect_events_key = Tenants.connect_per_second_rate(tenant).id
       expect(TenantBroadcaster, :pubsub_broadcast, 5, fn _, _, _, _ -> :ok end)
 
       messages_to_send =
@@ -438,7 +429,6 @@ defmodule RealtimeWeb.BroadcastControllerTest do
       GenCounter
       |> expect(:add, fn ^request_events_key -> :ok end)
       # remove the one message that won't be broadcasted for this user
-      |> expect(:add, 1, fn ^connect_events_key -> :ok end)
       |> expect(:add, length(messages) - 1, fn ^broadcast_events_key -> :ok end)
 
       conn = post(conn, Routes.broadcast_path(conn, :broadcast), %{"messages" => messages})