Move capacity and refill to config

Author: jyeshe

.iex.exs

@@ -1,5 +1,5 @@
 if Code.loaded?(Ecto.Query) do
-  Kernel.SpecialForms.import(Ecto.Query)
+  import Ecto.Query
 end
 
 alias Lightning.Repo

config/config.exs

@@ -25,6 +25,11 @@ config :lightning, LightningWeb.Endpoint,
   pubsub_server: Lightning.PubSub,
   live_view: [signing_salt: "EfrmuOUr"]
 
+config :lightning, Lightning.WebhookRateLimiter,
+  start: false,
+  capacity: 10,
+  refill_per_second: 2
+
 config :lightning, Lightning.Extensions,
   rate_limiter: Lightning.Extensions.RateLimiter,
   usage_limiter: Lightning.Extensions.UsageLimiter,

config/test.exs

@@ -53,6 +53,8 @@ config :lightning, LightningWeb.Endpoint,
     "/8zedVJLxvmGGFoRExE3e870g7CGZZQ1Vq11A5MbQGPKOpK57MahVsPW6Wkkv61n",
   server: true
 
+config :lightning, Lightning.WebhookRateLimiter, start: true
+
 config :lightning, Lightning.Runtime.RuntimeManager,
   ws_url: "ws://localhost:4002/worker"
 

lib/lightning/application.ex

@@ -7,6 +7,11 @@ defmodule Lightning.Application do
 
   require Logger
 
+  @rate_limiter_opts Application.compile_env!(
+                       :lightning,
+                       Lightning.WebhookRateLimiter
+                     )
+
   @impl true
   def start(_type, _args) do
     # Initialize ETS table for adapter lookup
@@ -181,10 +186,12 @@ defmodule Lightning.Application do
   end
 
   def start_phase(:init_rate_limiter, :normal, _args) do
-    Horde.DynamicSupervisor.start_child(
-      Lightning.DistributedSupervisor,
-      Lightning.WebhookRateLimiter
-    )
+    if @rate_limiter_opts[:start] do
+      Horde.DynamicSupervisor.start_child(
+        Lightning.DistributedSupervisor,
+        {Lightning.WebhookRateLimiter, @rate_limiter_opts}
+      )
+    end
 
     :ok
   end

lib/lightning/webhook_rate_limiter.ex

@@ -2,9 +2,6 @@ defmodule Lightning.WebhookRateLimiter do
   @moduledoc false
   use GenServer
 
-  @capacity 10
-  @refill_per_sec 2
-
   require Logger
 
   def child_spec(opts) do
@@ -17,31 +14,36 @@ defmodule Lightning.WebhookRateLimiter do
 
     %{
       id: id,
-      start: {__MODULE__, :start_link, [name]},
+      start: {__MODULE__, :start_link, [Keyword.put(opts, :name, name)]},
       shutdown: 10_000,
       restart: :transient
     }
   end
 
-  def start_link(name) do
+  def start_link(opts) do
+    name = Keyword.fetch!(opts, :name)
+
     with {:error, {:already_started, pid}} <-
-           GenServer.start_link(__MODULE__, [], name: via_tuple(name)) do
+           GenServer.start_link(__MODULE__, opts, name: via_tuple(name)) do
       Logger.info("already started at #{inspect(pid)}, returning :ignore")
       :ignore
     end
   end
 
   @impl true
-  def init([]) do
+  def init(opts) do
     Process.flag(:trap_exit, true)
 
-    {:ok, %{table: :ets.new(:table, [:set])}}
+    capacity = Keyword.fetch!(opts, :capacity)
+    refill = Keyword.fetch!(opts, :refill_per_second)
+
+    {:ok, %{table: :ets.new(:table, [:set]), capacity: capacity, refill_per_second: refill}}
   end
 
-  def check_rate(bucket, cost \\ 1, name \\ __MODULE__) do
+  def check_rate(bucket, capacity \\ nil, refill \\ nil, name \\ __MODULE__) do
     name
     |> via_tuple()
-    |> GenServer.call({:check_rate, bucket, cost})
+    |> GenServer.call({:check_rate, bucket, capacity, refill})
   end
 
   def inspect_table(name \\ __MODULE__) do
@@ -51,8 +53,8 @@ defmodule Lightning.WebhookRateLimiter do
   end
 
   @impl true
-  def handle_call({:check_rate, bucket, cost}, _from, %{table: table} = state) do
-    {:reply, do_check_rate(table, bucket, cost), state}
+  def handle_call({:check_rate, bucket, capacity, refill}, _from, state) do
+    {:reply, do_check_rate(state, bucket, capacity, refill), state}
   end
 
   @impl true
@@ -72,17 +74,24 @@ defmodule Lightning.WebhookRateLimiter do
     {:stop, :normal, state}
   end
 
-  def do_check_rate(table, bucket, cost) do
+  def do_check_rate(
+        %{table: table} = config,
+        bucket,
+        capacity,
+        refill_per_sec
+      ) do
     now = System.monotonic_time(:millisecond)
+    capacity = capacity || config[:capacity]
+    refill_per_sec = refill_per_sec || config[:refill_per_second]
 
-    :ets.insert_new(table, {bucket, {@capacity, now}})
+    :ets.insert_new(table, {bucket, {capacity, now}})
     [{^bucket, {level, updated}}] = :ets.lookup(table, bucket)
 
-    refilled = div(now - updated, 1_000) * @refill_per_sec
-    current = min(@capacity, level + refilled)
+    refilled = div(now - updated, 1_000) * refill_per_sec
+    current = min(capacity, level + refilled)
 
-    if current >= cost do
-      level = current - cost
+    if current >= 1 do
+      level = current - 1
       :ets.insert(table, {bucket, {level, now}})
 
       {:allow, level}
@@ -92,9 +101,6 @@ defmodule Lightning.WebhookRateLimiter do
     end
   end
 
-  def capacity, do: @capacity
-  def refill_per_second, do: @refill_per_sec
-
   def via_tuple(name),
     do: {:via, Horde.Registry, {Lightning.HordeRegistry, name}}
 end