fix: move rate limit message to happen inside the RateCounter (#1487)

Author: edgurgel

lib/realtime/rate_counter/rate_counter.ex

@@ -34,6 +34,7 @@ defmodule Realtime.RateCounter do
             tick_ref: nil,
             idle_shutdown: @idle_shutdown,
             idle_shutdown_ref: nil,
+            limit: %{log: false},
             telemetry: %{
               event_name: [@app_name] ++ [:rate_counter],
               measurements: %{sum: 0},
@@ -49,6 +50,12 @@ defmodule Realtime.RateCounter do
           tick_ref: reference(),
           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(),
@@ -133,6 +140,8 @@ defmodule Realtime.RateCounter do
     every = Keyword.get(args, :tick, @tick)
     max_bucket_len = Keyword.get(args, :max_bucket_len, @max_bucket_len)
     idle_shutdown_ms = Keyword.get(args, :idle_shutdown, @idle_shutdown)
+    limit_opts = Keyword.get(args, :limit)
+
     Logger.info("Starting #{__MODULE__} for: #{inspect(id)}")
 
     # Always reset the counter in case the counter had already accumulated without
@@ -153,6 +162,18 @@ defmodule Realtime.RateCounter do
         %{emit: false}
       end
 
+    limit =
+      if limit_opts do
+        %{
+          log: true,
+          value: limit_opts[:value],
+          log_fn: limit_opts[:log_fn],
+          triggered: false
+        }
+      else
+        %{log: false}
+      end
+
     ticker = tick(0)
 
     idle_shutdown_ref =
@@ -165,7 +186,8 @@ defmodule Realtime.RateCounter do
       max_bucket_len: max_bucket_len,
       idle_shutdown: idle_shutdown_ms,
       idle_shutdown_ref: idle_shutdown_ref,
-      telemetry: telemetry
+      telemetry: telemetry,
+      limit: limit
     }
 
     Cachex.put!(@cache, id, state)
@@ -195,6 +217,8 @@ defmodule Realtime.RateCounter do
       |> Kernel./(bucket_len)
 
     state = %{state | bucket: bucket, avg: avg}
+
+    state = maybe_trigger_limit(state)
     tick(state.tick)
 
     Cachex.put!(@cache, state.id, state)
@@ -226,6 +250,28 @@ defmodule Realtime.RateCounter do
     end
   end
 
+  defp maybe_trigger_limit(%{limit: %{log: false}} = state), do: state
+
+  defp maybe_trigger_limit(%{limit: %{triggered: true}} = 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
+      %{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
+      state.limit.log_fn.()
+
+      %{state | limit: %{state.limit | triggered: true}}
+    else
+      state
+    end
+  end
+
   defp tick(every) do
     Process.send_after(self(), :tick, every)
   end

lib/realtime/tenants.ex

@@ -151,7 +151,16 @@ defmodule Realtime.Tenants do
         event_name: [:channel, :joins],
         measurements: %{limit: max_joins_per_second},
         metadata: %{tenant: tenant_id}
-      }
+      },
+      limit: [
+        value: max_joins_per_second,
+        log_fn: fn ->
+          Logger.critical("ClientJoinRateLimitReached: Too many joins per second",
+            external_id: tenant_id,
+            project: tenant_id
+          )
+        end
+      ]
     ]
 
     %RateCounter.Args{id: joins_per_second_key(tenant_id), opts: opts}

lib/realtime_web/channels/realtime_channel.ex

@@ -137,7 +137,7 @@ defmodule RealtimeWeb.RealtimeChannel do
 
       {:error, :too_many_joins} ->
         msg = "Too many joins per second"
-        Logging.log_error_message(:error, "ClientJoinRateLimitReached", msg)
+        {:error, %{reason: msg}}
 
       {:error, :increase_connection_pool} ->
         msg = "Please increase your connection pool size"
@@ -504,10 +504,10 @@ defmodule RealtimeWeb.RealtimeChannel do
     rate_args = Tenants.joins_per_second_rate(tenant, limits.max_joins_per_second)
 
     RateCounter.new(rate_args)
-    GenCounter.add(rate_args.id)
 
     case RateCounter.get(rate_args) do
       {:ok, %{avg: avg}} when avg < limits.max_joins_per_second ->
+        GenCounter.add(rate_args.id)
         :ok
 
       {:ok, %{avg: _}} ->

mix.exs

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

test/integration/rt_channel_test.exs

@@ -1812,21 +1812,31 @@ defmodule Realtime.Integration.RtChannelTest do
       config = %{broadcast: %{self: true}, private: false}
       realtime_topic = "realtime:#{random_string()}"
 
-      for _ <- 1..1000 do
-        WebsocketClient.join(socket, realtime_topic, %{config: config})
-        1..5 |> Enum.random() |> Process.sleep()
-      end
+      log =
+        capture_log(fn ->
+          for _ <- 1..1000 do
+            WebsocketClient.join(socket, realtime_topic, %{config: config})
+            1..5 |> Enum.random() |> Process.sleep()
+          end
 
-      assert_receive %Message{
-                       event: "phx_reply",
-                       payload: %{
-                         "response" => %{"reason" => "Too many joins per second"},
-                         "status" => "error"
-                       }
-                     },
-                     2000
+          assert_receive %Message{
+                           event: "phx_reply",
+                           payload: %{
+                             "response" => %{"reason" => "Too many joins per second"},
+                             "status" => "error"
+                           }
+                         },
+                         2000
+
+          change_tenant_configuration(tenant, :max_joins_per_second, max_joins_per_second)
+        end)
+
+      assert log =~
+               "project=#{tenant.external_id} external_id=#{tenant.external_id} [critical] ClientJoinRateLimitReached: Too many joins per second"
 
-      change_tenant_configuration(tenant, :max_joins_per_second, max_joins_per_second)
+      # Only one log message should be emitted
+      # Splitting by the error message returns the error message and the rest of the log only
+      assert length(String.split(log, "ClientJoinRateLimitReached")) == 2
     end
   end
 

test/realtime/rate_counter/rate_counter_test.exs

@@ -1,10 +1,14 @@
 defmodule Realtime.RateCounterTest do
   use Realtime.DataCase, async: true
 
+  require Logger
+
   alias Realtime.RateCounter
   alias Realtime.RateCounter.Args
   alias Realtime.GenCounter
 
+  import ExUnit.CaptureLog
+
   describe "new/2" do
     test "starts a new rate counter without telemetry" do
       id = {:domain, :metric, Ecto.UUID.generate()}
@@ -14,13 +18,14 @@ defmodule Realtime.RateCounterTest do
       assert %Realtime.RateCounter{
                id: ^id,
                avg: +0.0,
-               bucket: [],
+               bucket: _,
                max_bucket_len: 60,
                tick: 1000,
                tick_ref: _,
                idle_shutdown: 900_000,
                idle_shutdown_ref: _,
-               telemetry: %{emit: false}
+               telemetry: %{emit: false},
+               limit: %{log: false}
              } = :sys.get_state(pid)
     end
 
@@ -39,6 +44,7 @@ defmodule Realtime.RateCounterTest do
       args = %Args{
         id: id,
         opts: [
+          tick: 10,
           telemetry: %{
             event_name: [:custom, :new_event],
             measurements: %{limit: 123},
@@ -54,7 +60,7 @@ defmodule Realtime.RateCounterTest do
                avg: +0.0,
                bucket: _,
                max_bucket_len: 60,
-               tick: 1000,
+               tick: 10,
                tick_ref: _,
                idle_shutdown: 900_000,
                idle_shutdown_ref: _,
@@ -75,6 +81,56 @@ defmodule Realtime.RateCounterTest do
       }
     end
 
+    test "starts a new rate counter with limit to log" do
+      id = {:domain, :metric, Ecto.UUID.generate()}
+
+      args = %Args{
+        id: id,
+        opts: [
+          tick: 100,
+          max_bucket_len: 10,
+          limit: [
+            value: 10,
+            log_fn: fn ->
+              Logger.error("ErrorMessage: Reason", external_id: "tenant123", project: "tenant123")
+            end
+          ]
+        ]
+      }
+
+      assert {:ok, pid} = RateCounter.new(args)
+
+      assert %RateCounter{
+               id: ^id,
+               avg: +0.0,
+               bucket: _,
+               max_bucket_len: 10,
+               telemetry: %{emit: false},
+               limit: %{
+                 log: true,
+                 value: 10,
+                 triggered: false
+               }
+             } = :sys.get_state(pid)
+
+      log =
+        capture_log(fn ->
+          GenCounter.add(args.id, 50)
+          Process.sleep(300)
+        end)
+
+      assert {:ok, %RateCounter{limit: %{triggered: true}}} = RateCounter.get(args)
+      assert log =~ "project=tenant123 external_id=tenant123 [error] ErrorMessage: Reason"
+
+      # Only one log message should be emitted
+      # Splitting by the error message returns the error message and the rest of the log only
+      assert length(String.split(log, "ErrorMessage: Reason")) == 2
+
+      Process.sleep(300)
+
+      assert {:ok, %RateCounter{limit: %{triggered: false}}} = RateCounter.get(args)
+    end
+
     test "reset counter if GenCounter already had something" do
       args = %Args{id: {:domain, :metric, Ecto.UUID.generate()}}
       GenCounter.add(args.id, 100)