feat: rate counter for failing auth due to PG connection errors (#1508)

* Change RateCounter to allow for not only average but sum limits
* Change Authorization to rate limit based on the amount of Connection
  errors we get when calling verifying channel authorisation

Author: edgurgel

README.md

@@ -239,8 +239,8 @@ This is the list of operational codes that can help you understand your deployme
 | ReconnectSubscribeToPostgres       | Postgres changes still waiting to be subscribed                                                                                                                                                       |
 | ChannelRateLimitReached            | The number of channels you can create has reached its limit                                                                                                                                           |
 | ConnectionRateLimitReached         | The number of connected clients as reached its limit                                                                                                                                                  |
-| ClientJoinRateLimitReached         | The rate of joins per second from your clients as reached the channel limits                                                                                                                          |
-| MessagePerSecondRateLimitReached   | The rate of messages per second from your clients as reached the channel limits                                                                                                                       |
+| ClientJoinRateLimitReached         | The rate of joins per second from your clients has reached the channel limits                                                                                                                         |
+| MessagePerSecondRateLimitReached   | The rate of messages per second from your clients has reached the channel limits                                                                                                                      |
 | RealtimeDisabledForTenant          | Realtime has been disabled for the tenant                                                                                                                                                             |
 | UnableToConnectToTenantDatabase    | Realtime was not able to connect to the tenant's database                                                                                                                                             |
 | DatabaseLackOfConnections          | Realtime was not able to connect to the tenant's database due to not having enough available connections                                                                                              |

lib/realtime/database.ex

@@ -186,6 +186,10 @@ defmodule Realtime.Database do
     e ->
       log_error("ErrorExecutingTransaction", e, metadata)
       {:error, e}
+  catch
+    :exit, reason ->
+      log_error("ErrorExecutingTransaction", reason, metadata)
+      {:error, {:exit, reason}}
   end
 
   @spec connect_db(__MODULE__.t()) :: {:ok, pid()} | {:error, any()}

lib/realtime/rate_counter/rate_counter.ex

@@ -28,40 +28,43 @@ defmodule Realtime.RateCounter do
 
   defstruct id: nil,
             avg: 0.0,
+            sum: 0,
             bucket: [],
             max_bucket_len: @max_bucket_len,
             tick: @tick,
             tick_ref: nil,
             idle_shutdown: @idle_shutdown,
             idle_shutdown_ref: nil,
             limit: %{log: false},
-            telemetry: %{
-              event_name: [@app_name] ++ [:rate_counter],
-              measurements: %{sum: 0},
-              metadata: %{}
-            }
+            telemetry: %{emit: false}
 
   @type t :: %__MODULE__{
           id: term(),
           avg: float(),
+          sum: non_neg_integer(),
           bucket: list(),
           max_bucket_len: integer(),
           tick: integer(),
-          tick_ref: reference(),
+          tick_ref: reference() | nil,
           idle_shutdown: integer() | :infinity,
-          idle_shutdown_ref: reference(),
-          limit: %{
-            log: boolean(),
-            value: integer(),
-            triggered: boolean(),
-            log_fn: (-> term())
-          },
-          telemetry: %{
-            emit: false,
-            event_name: :telemetry.event_name(),
-            measurements: :telemetry.event_measurements(),
-            metadata: :telemetry.event_metadata()
-          }
+          idle_shutdown_ref: reference() | nil,
+          limit:
+            %{log: false}
+            | %{
+                log: true,
+                value: integer(),
+                measurement: :sum | :avg,
+                triggered: boolean(),
+                log_fn: (-> term())
+              },
+          telemetry:
+            %{emit: false}
+            | %{
+                emit: true,
+                event_name: :telemetry.event_name(),
+                measurements: :telemetry.event_measurements(),
+                metadata: :telemetry.event_metadata()
+              }
         }
 
   @spec start_link([keyword()]) :: {:ok, pid()} | {:error, {:already_started, pid()}}
@@ -166,8 +169,9 @@ defmodule Realtime.RateCounter do
       if limit_opts do
         %{
           log: true,
-          value: limit_opts[:value],
-          log_fn: limit_opts[:log_fn],
+          value: Keyword.fetch!(limit_opts, :value),
+          measurement: Keyword.fetch!(limit_opts, :measurement),
+          log_fn: Keyword.fetch!(limit_opts, :log_fn),
           triggered: false
         }
       else
@@ -211,12 +215,10 @@ defmodule Realtime.RateCounter do
     bucket = [count | state.bucket] |> Enum.take(state.max_bucket_len)
     bucket_len = Enum.count(bucket)
 
-    avg =
-      bucket
-      |> Enum.sum()
-      |> Kernel./(bucket_len)
+    sum = Enum.sum(bucket)
+    avg = sum / bucket_len
 
-    state = %{state | bucket: bucket, avg: avg}
+    state = %{state | bucket: bucket, sum: sum, avg: avg}
 
     state = maybe_trigger_limit(state)
     tick(state.tick)
@@ -252,18 +254,18 @@ defmodule Realtime.RateCounter do
 
   defp maybe_trigger_limit(%{limit: %{log: false}} = state), do: state
 
-  defp maybe_trigger_limit(%{limit: %{triggered: true}} = state) do
+  defp maybe_trigger_limit(%{limit: %{triggered: true, measurement: measurement}} = state) do
     # Limit has been triggered, but we need to check if it is still above the limit
-    if state.avg < state.limit.value do
+    if Map.fetch!(state, measurement) < state.limit.value do
       %{state | limit: %{state.limit | triggered: false}}
     else
       # Limit is still above the threshold, so we keep the state as is
       state
     end
   end
 
-  defp maybe_trigger_limit(state) do
-    if state.avg >= state.limit.value do
+  defp maybe_trigger_limit(%{limit: %{measurement: measurement}} = state) do
+    if Map.fetch!(state, measurement) >= state.limit.value do
       state.limit.log_fn.()
 
       %{state | limit: %{state.limit | triggered: true}}

lib/realtime/tenants.ex

@@ -154,6 +154,7 @@ defmodule Realtime.Tenants do
       },
       limit: [
         value: max_joins_per_second,
+        measurement: :avg,
         log_fn: fn ->
           Logger.critical("ClientJoinRateLimitReached: Too many joins per second",
             external_id: tenant_id,
@@ -199,6 +200,7 @@ defmodule Realtime.Tenants do
       },
       limit: [
         value: max_events_per_second,
+        measurement: :avg,
         log: true,
         log_fn: fn ->
           Logger.error("MessagePerSecondRateLimitReached: Too many messages per second",
@@ -277,6 +279,7 @@ defmodule Realtime.Tenants do
       },
       limit: [
         value: max_presence_events_per_second,
+        measurement: :avg,
         log_fn: fn ->
           Logger.error("PresenceRateLimitReached: Too many presence events per second",
             external_id: tenant_id,
@@ -306,6 +309,31 @@ defmodule Realtime.Tenants do
     {:channel, :presence_events, tenant.external_id}
   end
 
+  @spec authorization_errors_per_second_rate(Tenant.t()) :: RateCounter.Args.t()
+  def authorization_errors_per_second_rate(%Tenant{external_id: external_id} = tenant) do
+    opts = [
+      max_bucket_len: 30,
+      limit: [
+        value: pool_size(tenant),
+        measurement: :sum,
+        log_fn: fn ->
+          Logger.critical("IncreaseConnectionPool: Too many database timeouts",
+            external_id: external_id,
+            project: external_id
+          )
+        end
+      ]
+    ]
+
+    %RateCounter.Args{id: {:channel, :authorization_errors, external_id}, opts: opts}
+  end
+
+  defp pool_size(%{extensions: [%{settings: settings} | _]}) do
+    Database.pool_size_by_application_name("realtime_connect", settings)
+  end
+
+  defp pool_size(_), do: 1
+
   @spec get_tenant_limits(Realtime.Api.Tenant.t(), maybe_improper_list) :: list
   def get_tenant_limits(%Tenant{} = tenant, keys) when is_list(keys) do
     nodes = [Node.self() | Node.list()]

lib/realtime/tenants/authorization.ex

@@ -13,9 +13,11 @@ defmodule Realtime.Tenants.Authorization do
 
   alias DBConnection.ConnectionError
   alias Realtime.Api.Message
+  alias Realtime.Api.Tenant
   alias Realtime.Database
-  alias Realtime.Repo
+  alias Realtime.GenCounter
   alias Realtime.GenRpc
+  alias Realtime.Repo
   alias Realtime.Tenants.Authorization.Policies
 
   defstruct [:tenant_id, :topic, :headers, :jwt, :claims, :role, :sub]
@@ -60,26 +62,42 @@ defmodule Realtime.Tenants.Authorization do
   @spec get_read_authorizations(Policies.t(), pid(), t()) ::
           {:ok, Policies.t()} | {:error, any()} | {:error, :rls_policy_error, any()}
   def get_read_authorizations(policies, db_conn, authorization_context) when node() == node(db_conn) do
-    case get_read_policies_for_connection(db_conn, authorization_context, policies) do
-      {:ok, %Policies{} = policies} -> {:ok, policies}
-      {:ok, {:error, %Postgrex.Error{} = error}} -> {:error, :rls_policy_error, error}
-      {:error, %ConnectionError{reason: :queue_timeout}} -> {:error, :increase_connection_pool}
-      {:error, error} -> {:error, error}
+    rate_counter = rate_counter(authorization_context.tenant_id)
+
+    if rate_counter.limit.triggered == false do
+      db_conn
+      |> get_read_policies_for_connection(authorization_context, policies)
+      |> handle_policies_result(rate_counter)
+    else
+      {:error, :increase_connection_pool}
     end
   end
 
   # Remote call
   def get_read_authorizations(policies, db_conn, authorization_context) do
-    case GenRpc.call(
-           node(db_conn),
-           __MODULE__,
-           :get_read_authorizations,
-           [policies, db_conn, authorization_context],
-           tenant_id: authorization_context.tenant_id,
-           key: authorization_context.tenant_id
-         ) do
-      {:error, :rpc_error, reason} -> {:error, reason}
-      response -> response
+    rate_counter = rate_counter(authorization_context.tenant_id)
+
+    if rate_counter.limit.triggered == false do
+      case GenRpc.call(
+             node(db_conn),
+             __MODULE__,
+             :get_read_authorizations,
+             [policies, db_conn, authorization_context],
+             tenant_id: authorization_context.tenant_id,
+             key: authorization_context.tenant_id
+           ) do
+        {:error, :increase_connection_pool} = error ->
+          GenCounter.add(rate_counter.id)
+          error
+
+        {:error, :rpc_error, reason} ->
+          {:error, reason}
+
+        response ->
+          response
+      end
+    else
+      {:error, :increase_connection_pool}
     end
   end
 
@@ -91,33 +109,70 @@ defmodule Realtime.Tenants.Authorization do
   @spec get_write_authorizations(Policies.t(), pid(), __MODULE__.t()) ::
           {:ok, Policies.t()} | {:error, any()} | {:error, :rls_policy_error, any()}
   def get_write_authorizations(policies, db_conn, authorization_context) when node() == node(db_conn) do
-    case get_write_policies_for_connection(db_conn, authorization_context, policies) do
-      {:ok, %Policies{} = policies} -> {:ok, policies}
-      {:ok, {:error, %Postgrex.Error{} = error}} -> {:error, :rls_policy_error, error}
-      {:error, %ConnectionError{reason: :queue_timeout}} -> {:error, :increase_connection_pool}
-      {:error, error} -> {:error, error}
+    rate_counter = rate_counter(authorization_context.tenant_id)
+
+    if rate_counter.limit.triggered == false do
+      db_conn
+      |> get_write_policies_for_connection(authorization_context, policies)
+      |> handle_policies_result(rate_counter)
+    else
+      {:error, :increase_connection_pool}
     end
   end
 
   # Remote call
   def get_write_authorizations(policies, db_conn, authorization_context) do
-    case GenRpc.call(
-           node(db_conn),
-           __MODULE__,
-           :get_write_authorizations,
-           [policies, db_conn, authorization_context],
-           tenant_id: authorization_context.tenant_id,
-           key: authorization_context.tenant_id
-         ) do
-      {:error, :rpc_error, reason} -> {:error, reason}
-      response -> response
+    rate_counter = rate_counter(authorization_context.tenant_id)
+
+    if rate_counter.limit.triggered == false do
+      case GenRpc.call(
+             node(db_conn),
+             __MODULE__,
+             :get_write_authorizations,
+             [policies, db_conn, authorization_context],
+             tenant_id: authorization_context.tenant_id,
+             key: authorization_context.tenant_id
+           ) do
+        {:error, :increase_connection_pool} = error ->
+          GenCounter.add(rate_counter.id)
+          error
+
+        {:error, :rpc_error, reason} ->
+          {:error, reason}
+
+        response ->
+          response
+      end
+    else
+      {:error, :increase_connection_pool}
     end
   end
 
   def get_write_authorizations(db_conn, authorization_context) do
     get_write_authorizations(%Policies{}, db_conn, authorization_context)
   end
 
+  defp handle_policies_result(result, rate_counter) do
+    case result do
+      {:ok, %Policies{} = policies} ->
+        {:ok, policies}
+
+      {:ok, {:error, %Postgrex.Error{} = error}} ->
+        {:error, :rls_policy_error, error}
+
+      {:error, %ConnectionError{reason: :queue_timeout}} ->
+        GenCounter.add(rate_counter.id)
+        {:error, :increase_connection_pool}
+
+      {:error, {:exit, _}} ->
+        GenCounter.add(rate_counter.id)
+        {:error, :increase_connection_pool}
+
+      {:error, error} ->
+        {:error, error}
+    end
+  end
+
   @doc """
   Sets the current connection configuration with the following config values:
   * role: The role of the user
@@ -282,4 +337,11 @@ defmodule Realtime.Tenants.Authorization do
         e
     end
   end
+
+  defp rate_counter(tenant_id) do
+    %Tenant{} = tenant = Realtime.Tenants.Cache.get_tenant_by_external_id(tenant_id)
+    rate_counter = Realtime.Tenants.authorization_errors_per_second_rate(tenant)
+    {:ok, rate_counter} = Realtime.RateCounter.get(rate_counter)
+    rate_counter
+  end
 end

lib/realtime_web/channels/realtime_channel.ex

@@ -205,12 +205,12 @@ defmodule RealtimeWeb.RealtimeChannel do
   end
 
   @impl true
-  def handle_info(:update_rate_counter, %{assigns: %{limits: %{max_events_per_second: max}}} = socket) do
+  def handle_info(:update_rate_counter, socket) do
     count(socket)
 
     {:ok, rate_counter} = RateCounter.get(socket.assigns.rate_counter)
 
-    if rate_counter.avg > max do
+    if rate_counter.limit.triggered do
       message = "Too many messages per second"
       shutdown_response(socket, message)
     else
@@ -490,11 +490,11 @@ defmodule RealtimeWeb.RealtimeChannel do
     RateCounter.new(rate_args)
 
     case RateCounter.get(rate_args) do
-      {:ok, %{avg: avg}} when avg < limits.max_joins_per_second ->
+      {:ok, %{limit: %{triggered: false}}} ->
         GenCounter.add(rate_args.id)
         :ok
 
-      {:ok, %{avg: _}} ->
+      {:ok, %{limit: %{triggered: true}}} ->
         {:error, :too_many_joins}
 
       error ->

mix.exs

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