Created Retry.Annotation: can annotate a function with @Retry

Author: bbalser

lib/retry/annotation.ex

@@ -0,0 +1,133 @@
+defmodule Retry.Annotation do
+  defmacro __using__(_opts) do
+    quote do
+      import Retry.DelayStreams
+      import Stream, only: [take: 2]
+      require Retry
+      require Logger
+
+      Module.register_attribute(__MODULE__, :retry_funs, accumulate: true)
+
+      @on_definition {Retry.Annotation, :on_def}
+      @before_compile {Retry.Annotation, :before_compile}
+    end
+  end
+
+  def on_def(env, kind, name, args, guards, _body) do
+    retry_opts = Module.get_attribute(env.module, :retry, :no_retry)
+
+    unless retry_opts == :no_retry do
+      Module.put_attribute(env.module, :retry_funs, %{
+        kind: kind,
+        name: name,
+        args: Enum.map(args, &de_underscore_name/1),
+        guards: guards,
+        retry_opts: Keyword.update(retry_opts, :with, [], &Enum.to_list/1)
+      })
+
+      Module.delete_attribute(env.module, :retry)
+    end
+  end
+
+  defmacro before_compile(env) do
+    retry_funs = Module.get_attribute(env.module, :retry_funs, [])
+    Module.delete_attribute(env.module, :retry_funs)
+
+    override_list =
+      Enum.map(retry_funs, &gen_override_list/1)
+      |> List.flatten()
+
+    overrides =
+      quote location: :keep do
+        defoverridable unquote(override_list)
+      end
+
+    functions =
+      Enum.map(retry_funs, fn fun ->
+        body = gen_body(fun)
+        gen_function(fun, body)
+      end)
+
+    [overrides | functions]
+  end
+
+  defp gen_override_list(%{name: name, args: args}) do
+    no_default_args_length =
+      Enum.reduce(args, 0, fn
+        {:\\, _, _}, acc -> acc
+        _, acc -> acc + 1
+      end)
+
+    Enum.map(no_default_args_length..length(args), fn i -> {name, i} end)
+  end
+
+  defp gen_function(%{kind: :def, guards: [], name: name, args: args}, body) do
+    quote location: :keep do
+      def unquote(name)(unquote_splicing(args)) do
+        unquote(body)
+      end
+    end
+  end
+
+  defp gen_function(%{kind: :def, guards: guards, name: name, args: args}, body) do
+    quote location: :keep do
+      def unquote(name)(unquote_splicing(args)) when unquote_splicing(guards) do
+        unquote(body)
+      end
+    end
+  end
+
+  defp gen_function(%{kind: :defp, guards: [], name: name, args: args}, body) do
+    quote location: :keep do
+      defp unquote(name)(unquote_splicing(args)) do
+        unquote(body)
+      end
+    end
+  end
+
+  defp gen_function(%{kind: :defp, guards: guards, name: name, args: args}, body) do
+    quote location: :keep do
+      defp unquote(name)(unquote_splicing(args)) when unquote_splicing(guards) do
+        unquote(body)
+      end
+    end
+  end
+
+  defp gen_body(fun) do
+    args =
+      Enum.map(fun.args, fn
+        {:\\, _, [arg_name | _]} -> arg_name
+        arg -> arg
+      end)
+
+    quote location: :keep do
+      Retry.retry_while unquote(fun.retry_opts) do
+        case super(unquote_splicing(args)) do
+          {:error, reason} = error ->
+            Logger.info(fn ->
+              "#{__MODULE__}: Retrying function #{unquote(fun.name)}: #{inspect(reason)}"
+            end)
+
+            {:cont, error}
+
+          result ->
+            {:halt, result}
+        end
+      end
+    end
+  end
+
+  defp de_underscore_name({:\\, context, [{name, name_context, name_args} | t]} = arg) do
+    case to_string(name) do
+      "_" <> real_name -> {:\\, context, [{String.to_atom(real_name), name_context, name_args} | t]}
+      _ -> arg
+    end
+  end
+
+  defp de_underscore_name({name, context, args} = arg) do
+    case to_string(name) do
+      "_" <> real_name -> {String.to_atom(real_name), context, args}
+      _ -> arg
+    end
+  end
+end

test/retry/annotation_test.exs

@@ -0,0 +1,120 @@
+defmodule Retry.AnnotationTest do
+  use ExUnit.Case
+
+  defmodule Example do
+    use Retry.Annotation
+
+    @retry with: constant_backoff(100) |> take(10)
+    def default_params(x, _opts \\ []) do
+      {:ok, x}
+    end
+
+    @retry with: constant_backoff(100) |> take(10)
+    def process(pid) do
+      Agent.get_and_update(:test_store, fn s ->
+        {s, max(0, s - 1)}
+      end)
+      |> case do
+        0 ->
+          {:ok, 0}
+
+        n ->
+          send(pid, {:attempt, n})
+          {:error, "attempts remaining #{n}"}
+      end
+    end
+
+    @retry with: constant_backoff(100) |> take(10)
+    def process_with_guard(pid, x) when is_pid(pid) and is_binary(x) do
+      Agent.get_and_update(:test_store, fn s ->
+        {s, max(0, s - 1)}
+      end)
+      |> case do
+        0 ->
+          {:ok, x}
+
+        n ->
+          send(pid, {:attempt, n})
+          {:error, "attempts remaining #{n}"}
+      end
+    end
+
+    def no_retry(pid) do
+      send(pid, :no_retry)
+      {:error, "no_retry"}
+    end
+
+    def wrapper(pid) do
+      internal(pid)
+    end
+
+    @retry with: constant_backoff(100) |> take(10)
+    defp internal(pid) do
+      Agent.get_and_update(:test_store, fn s ->
+        {s, max(0, s - 1)}
+      end)
+      |> case do
+        0 ->
+          {:ok, 0}
+
+        n ->
+          send(pid, {:attempt, n})
+          {:error, "attempts remaining #{n}"}
+      end
+    end
+  end
+
+  setup do
+    {:ok, pid} = Agent.start_link(fn -> 0 end, name: :test_store)
+    on_exit(fn -> assert_down(pid) end)
+
+    :ok
+  end
+
+  test "will not retry function when ok tuple returned" do
+    assert {:ok, 0} = Example.process(self())
+    refute_receive {:attempt, _}
+  end
+
+  test "will retry on error tuple until ok tuple is received" do
+    Agent.update(:test_store, fn _ -> 6 end)
+    assert {:ok, 0} = Example.process(self())
+
+    Enum.each(1..6, fn i ->
+      assert_receive {:attempt, ^i}
+    end)
+  end
+
+  test "should not retry function that are not annotated" do
+    assert {:error, "no_retry"} = Example.no_retry(self())
+    assert_receive :no_retry
+    refute_receive :no_retry
+  end
+
+  test "guard clauses are still enforced on override function" do
+    Agent.update(:test_store, fn _ -> 4 end)
+    assert {:ok, "hello"} == Example.process_with_guard(self(), "hello")
+
+    Enum.each(1..4, fn i -> assert_receive {:attempt, ^i} end)
+
+    assert_raise FunctionClauseError, fn ->
+      Example.process_with_guard(self(), 1)
+    end
+  end
+
+  test "retries private functions as well" do
+    Agent.update(:test_store, fn _ -> 7 end)
+    assert {:ok, 0} = Example.wrapper(self())
+    Enum.each(1..7, fn i -> assert_receive {:attempt, ^i} end)
+
+    assert_raise UndefinedFunctionError, fn ->
+      Example.internal(self())
+    end
+  end
+
+  defp assert_down(pid) do
+    ref = Process.monitor(pid)
+    Process.exit(pid, :kill)
+    assert_receive {:DOWN, ^ref, _, _, _}
+  end
+end