Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

use backend #5

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from
Draft

use backend #5

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4e347df
wip
ruslandoga May 4, 2024
986e0e6
wip
ruslandoga May 5, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
273 changes: 89 additions & 184 deletions lib/hammer.ex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,14 +4,77 @@ defmodule Hammer do

This is the main API for the Hammer rate-limiter. This module assumes a
backend pool has been started, most likely by the Hammer application.

defmodule MyApp.RateLimit do
use Hammer, backend: Hammer.Backend.ETS
end

MyApp.RateLimit.start_link(expiry_ms: :timer.minutes(10), cleanup_interval_ms: :timer.minutes(1))
{:allow, 1] = MyApp.RateLimit.hit("some-key", _scale = :timer.seconds(60), _limit = 10)
{:allow, 1} = MyApp.RateLimit.get("some-key", _scale = :timer.seconds(60), _limit = 10)
MyApp.RateLimit.wait("some-key", _scale = :timer.seconds(60), _limit = 10)

backend = Hammer.backend(Hammer.Backend.ETS, opts)
Hammer.check_rate(backend, )

"""

alias Hammer.Utils
@doc false
defmacro __using__(opts) do
quote bind_quoted: [opts: opts] do
backend = Keyword.fetch!(opts, :backend)

unless is_atom(backend) do
raise ArgumentError,
"expected :backend to be a module name like Hammer.Backend.ETS, got: #{inspect(backend)}"
end

@behaviour Hammer
@backend backend

def __backend__ do
@backend
end

def start_link(opts) do
@backend.start_link(opts)
end

def child_spec(opts) do
@backend.child_spec(opts)
end

def check_rate(id, scale_ms, limit) do
@backend.check_rate_inc(id, scale_ms, limit, 1)
end

def check_rate_inc(id, scale_ms, limit, increment) do
@backend.check_rate_inc(id, scale_ms, limit, increment)
end

def inspect_bucket(id, scale_ms, limit) do
{stamp, key} = Hammer.Utils.stamp_key(id, scale_ms)
ms_to_next_bucket = elem(key, 0) * scale_ms + scale_ms - stamp

case @backend.get_bucket(key) do
{:ok, nil} ->
{:ok, {0, limit, ms_to_next_bucket, nil, nil}}

{:ok, {_, count, created_at, updated_at}} ->
count_remaining = if limit > count, do: limit - count, else: 0
{:ok, {count, count_remaining, ms_to_next_bucket, created_at, updated_at}}

{:error, reason} ->
{:error, reason}
end
end

def delete_buckets(id) do
@backend.delete_buckets(id)
end
end
end

@spec check_rate(id :: String.t(), scale_ms :: integer, limit :: integer) ::
{:allow, count :: integer}
| {:deny, limit :: integer}
| {:error, reason :: any}
@doc """
Check if the action you wish to perform is within the bounds of the rate-limit.

Expand All @@ -27,92 +90,34 @@ defmodule Hammer do
Example:

user_id = 42076
case check_rate("file_upload:\#{user_id}", 60_000, 5) do
case MyHammer.check_rate("file_upload:\#{user_id}", 60_000, 5) do
{:allow, _count} ->
# do the file upload
{:deny, _limit} ->
# render an error page or something
end
"""
def check_rate(id, scale_ms, limit) do
check_rate(:single, id, scale_ms, limit)
end

@spec check_rate(backend :: atom, id :: String.t(), scale_ms :: integer, limit :: integer) ::
{:allow, count :: integer}
| {:deny, limit :: integer}
| {:error, reason :: any}
@doc """
Same as `check_rate/3`, but allows specifying a backend.
"""
def check_rate(backend, id, scale_ms, limit) do
{stamp, key} = Utils.stamp_key(id, scale_ms)

case call_backend(backend, :count_hit, [key, stamp]) do
{:ok, count} ->
if count > limit do
{:deny, limit}
else
{:allow, count}
end
@callback check_rate(id :: String.t(), scale_ms :: integer, limit :: integer) ::
{:allow, count :: integer}
| {:deny, limit :: integer}
| {:error, reason :: any}

{:error, reason} ->
{:error, reason}
end
end

@spec check_rate_inc(
id :: String.t(),
scale_ms :: integer,
limit :: integer,
increment :: integer
) ::
{:allow, count :: integer}
| {:deny, limit :: integer}
| {:error, reason :: any}
@doc """
Same as check_rate/3, but allows the increment number to be specified.
This is useful for limiting apis which have some idea of 'cost', where the cost
of each hit can be specified.
"""
def check_rate_inc(id, scale_ms, limit, increment) do
check_rate_inc(:single, id, scale_ms, limit, increment)
end
@callback check_rate_inc(
id :: String.t(),
scale_ms :: integer,
limit :: integer,
increment :: integer
) ::
{:allow, count :: integer}
| {:deny, limit :: integer}
| {:error, reason :: any}

@spec check_rate_inc(
backend :: atom,
id :: String.t(),
scale_ms :: integer,
limit :: integer,
increment :: integer
) ::
{:allow, count :: integer}
| {:deny, limit :: integer}
| {:error, reason :: any}
@doc """
Same as check_rate_inc/4, but allows specifying a backend.
"""
def check_rate_inc(backend, id, scale_ms, limit, increment) do
{stamp, key} = Utils.stamp_key(id, scale_ms)

case call_backend(backend, :count_hit, [key, stamp, increment]) do
{:ok, count} ->
if count > limit do
{:deny, limit}
else
{:allow, count}
end

{:error, reason} ->
{:error, reason}
end
end

@spec inspect_bucket(id :: String.t(), scale_ms :: integer, limit :: integer) ::
{:ok,
{count :: integer, count_remaining :: integer, ms_to_next_bucket :: integer,
created_at :: integer | nil, updated_at :: integer | nil}}
| {:error, reason :: any}
@doc """
Inspect bucket to get count, count_remaining, ms_to_next_bucket, created_at,
updated_at. This function is free of side-effects and should be called with
Expand All @@ -137,38 +142,12 @@ defmodule Hammer do
{:ok, {1, 2499, 29381612, 1450281014468, 1450281014468}}

"""
def inspect_bucket(id, scale_ms, limit) do
inspect_bucket(:single, id, scale_ms, limit)
end

@spec inspect_bucket(backend :: atom, id :: String.t(), scale_ms :: integer, limit :: integer) ::
{:ok,
{count :: integer, count_remaining :: integer, ms_to_next_bucket :: integer,
created_at :: integer | nil, updated_at :: integer | nil}}
| {:error, reason :: any}
@doc """
Same as inspect_bucket/3, but allows specifying a backend
"""
def inspect_bucket(backend, id, scale_ms, limit) do
{stamp, key} = Utils.stamp_key(id, scale_ms)
ms_to_next_bucket = elem(key, 0) * scale_ms + scale_ms - stamp

case call_backend(backend, :get_bucket, [key]) do
{:ok, nil} ->
{:ok, {0, limit, ms_to_next_bucket, nil, nil}}

{:ok, {_, count, created_at, updated_at}} ->
count_remaining = if limit > count, do: limit - count, else: 0
{:ok, {count, count_remaining, ms_to_next_bucket, created_at, updated_at}}

{:error, reason} ->
{:error, reason}
end
end
@callback inspect_bucket(id :: String.t(), scale_ms :: integer, limit :: integer) ::
{:ok,
{count :: integer, count_remaining :: integer, ms_to_next_bucket :: integer,
created_at :: integer | nil, updated_at :: integer | nil}}
| {:error, reason :: any}

@spec delete_buckets(id :: String.t()) ::
{:ok, count :: integer}
| {:error, reason :: any}
@doc """
Delete all buckets belonging to the provided id, including the current one.
Effectively resets the rate-limit for the id.
Expand All @@ -183,82 +162,8 @@ defmodule Hammer do
Example:

user_id = 2406
{:ok, _count} = delete_buckets("file_uploads:\#{user_id}")

"""
def delete_buckets(id) do
delete_buckets(:single, id)
end

@spec delete_buckets(backend :: atom, id :: String.t()) ::
{:ok, count :: integer}
| {:error, reason :: any}
@doc """
Same as delete_buckets/1, but allows specifying a backend
"""
def delete_buckets(backend, id) do
call_backend(backend, :delete_buckets, [id])
end

@spec make_rate_checker(id_prefix :: String.t(), scale_ms :: integer, limit :: integer) ::
(id :: String.t() ->
{:allow, count :: integer}
| {:deny, limit :: integer}
| {:error, reason :: any})
@doc """
Make a rate-checker function, with the given `id` prefix, scale_ms and limit.
{:ok, _count} = MyHammer.delete_buckets("file_uploads:\#{user_id}")

Arguments:

- `id_prefix`: String prefix to the `id`
- `scale_ms`: Integer indicating size of bucket in milliseconds
- `limit`: Integer maximum count of actions within the bucket

Returns a function which accepts an `id` suffix, which will be combined with
the `id_prefix`. Calling this returned function is equivalent to:
`Hammer.check_rate("\#{id_prefix}\#{id}", scale_ms, limit)`

Example:

chat_rate_limiter = make_rate_checker("send_chat_message:", 60_000, 20)
user_id = 203517
case chat_rate_limiter.(user_id) do
{:allow, _count} ->
# allow chat message
{:deny, _limit} ->
# deny
end
"""
def make_rate_checker(id_prefix, scale_ms, limit) do
make_rate_checker(:single, id_prefix, scale_ms, limit)
end

@spec make_rate_checker(
backend :: atom,
id_prefix :: String.t(),
scale_ms :: integer,
limit :: integer
) ::
(id :: String.t() ->
{:allow, count :: integer}
| {:deny, limit :: integer}
| {:error, reason :: any})
@doc """
"""
def make_rate_checker(backend, id_prefix, scale_ms, limit) do
fn id ->
check_rate(backend, "#{id_prefix}#{id}", scale_ms, limit)
end
end

defp call_backend(which, function, args) do
pool = Utils.pool_name(which)
backend = Utils.get_backend_module(which)

:poolboy.transaction(
pool,
fn pid -> apply(backend, function, [pid | args]) end,
60_000
)
end
@callback delete_buckets(id :: String.t()) :: {:ok, count :: integer} | {:error, reason :: any}
end
5 changes: 2 additions & 3 deletions lib/hammer/application.ex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,16 +53,15 @@ defmodule Hammer.Application do
"""

use Application
require Logger

def start(_type, _args) do
config =
children =
Application.get_env(
:hammer,
:backend,
{Hammer.Backend.ETS, []}
)

Hammer.Supervisor.start_link(config, name: Hammer.Supervisor)
Hammer.Supervisor.start_link(children, name: Hammer.Supervisor)
end
end
35 changes: 6 additions & 29 deletions lib/hammer/backend.ex
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,35 +7,12 @@ defmodule Hammer.Backend do
@type bucket_info ::
{key :: bucket_key, count :: integer, created :: integer, updated :: integer}

@callback count_hit(
pid :: pid(),
key :: bucket_key,
now :: integer
) ::
{:ok, count :: integer}
| {:error, reason :: any}
@callback count_hit(key :: bucket_key, now :: integer, increment :: integer, opts :: keyword) ::
{:ok, count :: integer} | {:error, reason :: any}

@callback count_hit(
pid :: pid(),
key :: bucket_key,
now :: integer,
increment :: integer
) ::
{:ok, count :: integer}
| {:error, reason :: any}
@callback get_bucket(key :: bucket_key, opts :: keyword) ::
{:ok, info :: bucket_info | nil} | {:error, reason :: any}

@callback get_bucket(
pid :: pid(),
key :: bucket_key
) ::
{:ok, info :: bucket_info}
| {:ok, nil}
| {:error, reason :: any}

@callback delete_buckets(
pid :: pid(),
id :: String.t()
) ::
{:ok, count_deleted :: integer}
| {:error, reason :: any}
@callback delete_buckets(id :: String.t(), opts :: keyword) ::
{:ok, count_deleted :: integer} | {:error, reason :: any}
end
Loading
Loading