From 554524db9bceb1f708f5529f4605a5797a80f7a9 Mon Sep 17 00:00:00 2001
From: Jordan Bracco <href@random.sh>
Date: Fri, 8 May 2020 01:55:49 +0200
Subject: [PATCH] Initial commit

---
 .formatter.exs           |  4 +++
 .gitignore               | 24 ++++++++++++++++
 lib/limiter.ex           | 47 +++++++++++++++++++++++++++++++
 mix.exs                  | 29 +++++++++++++++++++
 mix.lock                 |  4 +++
 test/limiter_test.exs    | 60 ++++++++++++++++++++++++++++++++++++++++
 test/samples/limiter.exs | 11 ++++++++
 test/test_helper.exs     |  1 +
 8 files changed, 180 insertions(+)
 create mode 100644 .formatter.exs
 create mode 100644 .gitignore
 create mode 100644 lib/limiter.ex
 create mode 100644 mix.exs
 create mode 100644 mix.lock
 create mode 100644 test/limiter_test.exs
 create mode 100644 test/samples/limiter.exs
 create mode 100644 test/test_helper.exs

diff --git a/.formatter.exs b/.formatter.exs
new file mode 100644
index 0000000..d2cda26
--- /dev/null
+++ b/.formatter.exs
@@ -0,0 +1,4 @@
+# Used by "mix format"
+[
+  inputs: ["{mix,.formatter}.exs", "{config,lib,test}/**/*.{ex,exs}"]
+]
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..97a97a3
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,24 @@
+# The directory Mix will write compiled artifacts to.
+/_build/
+
+# If you run "mix test --cover", coverage assets end up here.
+/cover/
+
+# The directory Mix downloads your dependencies sources to.
+/deps/
+
+# Where third-party dependencies like ExDoc output generated docs.
+/doc/
+
+# Ignore .fetch files in case you like to edit your project deps locally.
+/.fetch
+
+# If the VM crashes, it generates a dump, let's ignore it too.
+erl_crash.dump
+
+# Also ignore archive artifacts (built via "mix archive.build").
+*.ez
+
+# Ignore package tarball (built via "mix hex.build").
+limiter-*.tar
+
diff --git a/lib/limiter.ex b/lib/limiter.ex
new file mode 100644
index 0000000..84c601c
--- /dev/null
+++ b/lib/limiter.ex
@@ -0,0 +1,47 @@
+defmodule Limiter do
+  @ets __MODULE__.ETS
+
+  def new(name, max_running, max_waiting) do
+    name = atom_name(name)
+    :persistent_term.put(name, {max_running, max_waiting})
+    :ets.new(name, [:public, :named_table])
+    :ok
+  end
+
+  def limit(name, fun) do
+    {max_running, max_waiting} = :persistent_term.get(atom_name(name))
+    max = max_running + max_waiting
+    counter = inc(name)
+
+    cond do
+      counter <= max_running ->
+        fun.()
+
+      counter > max ->
+        {:error, :overload}
+
+      counter > max_running ->
+        wait(name, fun)
+    end
+  after
+    dec(name)
+  end
+
+  defp wait(name, fun) do
+    Process.sleep(150)
+    dec(name)
+    limit(name, fun)
+  end
+
+  defp inc(name) do
+    name = atom_name(name)
+    :ets.update_counter(name, name, {2, 1}, {name, 0})
+  end
+
+  def dec(name) do
+    name = atom_name(name)
+    :ets.update_counter(name, name, {2, -1}, {name, 0})
+  end
+
+  defp atom_name(suffix), do: Module.concat(@ets, suffix)
+end
diff --git a/mix.exs b/mix.exs
new file mode 100644
index 0000000..f3ff456
--- /dev/null
+++ b/mix.exs
@@ -0,0 +1,29 @@
+defmodule Limiter.MixProject do
+  use Mix.Project
+
+  def project do
+    [
+      app: :limiter,
+      version: "0.1.0",
+      elixir: "~> 1.10",
+      start_permanent: Mix.env() == :prod,
+      deps: deps()
+    ]
+  end
+
+  # Run "mix help compile.app" to learn about applications.
+  def application do
+    [
+      extra_applications: [:logger]
+    ]
+  end
+
+  # Run "mix help deps" to learn about dependencies.
+  defp deps do
+    [
+      {:benchee, "~> 1.0", only: [:dev, :test]}
+      # {:dep_from_hexpm, "~> 0.3.0"},
+      # {:dep_from_git, git: "https://github.com/elixir-lang/my_dep.git", tag: "0.1.0"}
+    ]
+  end
+end
diff --git a/mix.lock b/mix.lock
new file mode 100644
index 0000000..5aa28d0
--- /dev/null
+++ b/mix.lock
@@ -0,0 +1,4 @@
+%{
+  "benchee": {:hex, :benchee, "1.0.1", "66b211f9bfd84bd97e6d1beaddf8fc2312aaabe192f776e8931cb0c16f53a521", [:mix], [{:deep_merge, "~> 1.0", [hex: :deep_merge, repo: "hexpm", optional: false]}], "hexpm", "3ad58ae787e9c7c94dd7ceda3b587ec2c64604563e049b2a0e8baafae832addb"},
+  "deep_merge": {:hex, :deep_merge, "1.0.0", "b4aa1a0d1acac393bdf38b2291af38cb1d4a52806cf7a4906f718e1feb5ee961", [:mix], [], "hexpm", "ce708e5f094b9cd4e8f2be4f00d2f4250c4095be93f8cd6d018c753894885430"},
+}
diff --git a/test/limiter_test.exs b/test/limiter_test.exs
new file mode 100644
index 0000000..c60db9e
--- /dev/null
+++ b/test/limiter_test.exs
@@ -0,0 +1,60 @@
+defmodule LimiterTest do
+  use ExUnit.Case
+  doctest Limiter
+
+  defp test_ets(name, max, sleep, fun) do
+    count = :ets.update_counter(:limiter_test, name, {2, 1}, {name, 0})
+
+    if count <= max do
+      fun.({:ok, count})
+      Process.sleep(sleep)
+    else
+      fun.(:fail)
+    end
+  after
+    :ets.update_counter(:limiter_test, name, {2, -1}, {name, 1})
+  end
+
+  test "limits with ets" do
+    :ets.new(:limiter_test, [:public, :named_table])
+    ets = "test"
+    test = self()
+    spawn_link(fn -> test_ets(ets, 2, 500, fn result -> send(test, result) end) end)
+    spawn_link(fn -> test_ets(ets, 2, 750, fn result -> send(test, result) end) end)
+    spawn_link(fn -> test_ets(ets, 2, 500, fn result -> send(test, result) end) end)
+    assert_receive {:ok, 1}
+    assert_receive {:ok, 2}
+    assert_receive :fail
+    Process.sleep(500)
+    spawn_link(fn -> test_ets(ets, 2, 500, fn result -> send(test, result) end) end)
+    assert_receive {:ok, 2}
+  end
+
+  test "limiter" do
+    name = "test1"
+    self = self()
+    Limiter.set(name, 2, 2)
+
+    sleepy = fn sleep ->
+      case Limiter.limit(name, fn ->
+             send(self, :ok)
+             Process.sleep(sleep)
+             :ok
+           end) do
+        :ok -> :ok
+        other -> send(self, other)
+      end
+    end
+
+    spawn_link(fn -> sleepy.(500) end)
+    spawn_link(fn -> sleepy.(500) end)
+    spawn_link(fn -> sleepy.(500) end)
+    spawn_link(fn -> sleepy.(500) end)
+    spawn_link(fn -> sleepy.(500) end)
+    assert_receive :ok, 2000
+    assert_receive :ok, 2000
+    assert_receive {:error, :overload}, 2000
+    assert_receive :ok, 2000
+    assert_receive :ok, 2000
+  end
+end
diff --git a/test/samples/limiter.exs b/test/samples/limiter.exs
new file mode 100644
index 0000000..896056d
--- /dev/null
+++ b/test/samples/limiter.exs
@@ -0,0 +1,11 @@
+:ets.new(:limiter_bench, [:public, :named_table])
+Limiter.new(:bench, 1_000_000_000_000_000_000_000_000_000_000_000_000_000_000_000, 0)
+
+Benchee.run(%{
+  "update_counter" => fn ->
+    :ets.update_counter(:limiter_bench, "bench", {2, 1}, {"bench", 0})
+  end,
+  "limit" => fn ->
+    Limiter.limit(:bench, fn -> :ok end)
+  end
+})
diff --git a/test/test_helper.exs b/test/test_helper.exs
new file mode 100644
index 0000000..869559e
--- /dev/null
+++ b/test/test_helper.exs
@@ -0,0 +1 @@
+ExUnit.start()