diff --git a/lib/pleroma/plugs/rate_limiter/limiter_supervisor.ex b/lib/pleroma/plugs/rate_limiter/limiter_supervisor.ex index 187582ede..884268d96 100644 --- a/lib/pleroma/plugs/rate_limiter/limiter_supervisor.ex +++ b/lib/pleroma/plugs/rate_limiter/limiter_supervisor.ex @@ -7,8 +7,8 @@ defmodule Pleroma.Plugs.RateLimiter.LimiterSupervisor do DynamicSupervisor.start_link(__MODULE__, init_arg, name: __MODULE__) end - def add_limiter(limiter_name, expiration) do - {:ok, _pid} = + def add_or_return_limiter(limiter_name, expiration) do + result = DynamicSupervisor.start_child( __MODULE__, %{ @@ -28,6 +28,12 @@ defmodule Pleroma.Plugs.RateLimiter.LimiterSupervisor do ]} } ) + + case result do + {:ok, _pid} = result -> result + {:error, {:already_started, pid}} -> {:ok, pid} + _ -> result + end end @impl true diff --git a/lib/pleroma/plugs/rate_limiter/rate_limiter.ex b/lib/pleroma/plugs/rate_limiter/rate_limiter.ex index 9c362a392..b9cbe9716 100644 --- a/lib/pleroma/plugs/rate_limiter/rate_limiter.ex +++ b/lib/pleroma/plugs/rate_limiter/rate_limiter.ex @@ -171,7 +171,7 @@ defmodule Pleroma.Plugs.RateLimiter do {:error, value} {:error, :no_cache} -> - initialize_buckets(action_settings) + initialize_buckets!(action_settings) check_rate(action_settings) end end @@ -250,11 +250,16 @@ defmodule Pleroma.Plugs.RateLimiter do |> String.replace_leading(":", "") end - defp initialize_buckets(%{name: _name, limits: nil}), do: :ok + defp initialize_buckets!(%{name: _name, limits: nil}), do: :ok - defp initialize_buckets(%{name: name, limits: limits}) do - LimiterSupervisor.add_limiter(anon_bucket_name(name), get_scale(:anon, limits)) - LimiterSupervisor.add_limiter(user_bucket_name(name), get_scale(:user, limits)) + defp initialize_buckets!(%{name: name, limits: limits}) do + {:ok, _pid} = + LimiterSupervisor.add_or_return_limiter(anon_bucket_name(name), get_scale(:anon, limits)) + + {:ok, _pid} = + LimiterSupervisor.add_or_return_limiter(user_bucket_name(name), get_scale(:user, limits)) + + :ok end defp attach_identity(base, %{mode: :user, conn_info: conn_info}), diff --git a/test/plugs/rate_limiter_test.exs b/test/plugs/rate_limiter_test.exs index 104d67611..8cdc8d1a2 100644 --- a/test/plugs/rate_limiter_test.exs +++ b/test/plugs/rate_limiter_test.exs @@ -242,4 +242,35 @@ defmodule Pleroma.Plugs.RateLimiterTest do refute conn_2.halted end end + + test "doesn't crash due to a race condition when multiple requests are made at the same time and the bucket is not yet initialized" do + limiter_name = :test_race_condition + Pleroma.Config.put([:rate_limit, limiter_name], {1000, 5}) + Pleroma.Config.put([Pleroma.Web.Endpoint, :http, :ip], {8, 8, 8, 8}) + + opts = RateLimiter.init(name: limiter_name) + + conn = conn(:get, "/") + conn_2 = conn(:get, "/") + + %Task{pid: pid1} = + task1 = + Task.async(fn -> + receive do + :process2_up -> + RateLimiter.call(conn, opts) + end + end) + + task2 = + Task.async(fn -> + send(pid1, :process2_up) + RateLimiter.call(conn_2, opts) + end) + + Task.await(task1) + Task.await(task2) + + refute {:err, :not_found} == RateLimiter.inspect_bucket(conn, limiter_name, opts) + end end