feat: add support for should_retry with arity of 2 to Retry middleware (#608)

* Add support for should_retry with arity of 2 to Retry middleware

* Format Middleware

* Change should_retry option to be of arity 1 or 3 instead

Author: mdlkxzmcp

lib/tesla/middleware.ex

@@ -14,13 +14,13 @@ defmodule Tesla.Middleware do
   or inside tuple in case of dynamic middleware (`Tesla.client/1`):
 
       Tesla.client([{Tesla.Middleware.BaseUrl, "https://example.com"}])
-      
+
   ## Ordering
-  
+
   The order in which middleware is defined matters. Note that the order when _sending_ the request
   matches the order the middleware was defined in, but the order when _receiving_ the response
   is reversed.
-  
+
   For example, `Tesla.Middleware.DecompressResponse` must come _after_ `Tesla.Middleware.JSON`,
   otherwise the response isn't decompressed before it reaches the JSON parser.
 

lib/tesla/middleware/retry.ex

@@ -34,6 +34,14 @@ defmodule Tesla.Middleware.Retry do
         {:ok, _} -> false
         {:error, _} -> true
       end
+      # or
+      plug Tesla.Middleware.Retry, should_retry: fn
+        {:ok, %{status: status}}, _env, _context when status in [400, 500] -> true
+        {:ok, _reason}, _env, _context -> false
+        {:error, _reason}, %Tesla.Env{method: :post}, _context -> false
+        {:error, _reason}, %Tesla.Env{method: :put}, %{retries: 2} -> false
+        {:error, _reason}, _env, _context -> true
+      end
   end
   ```
 
@@ -42,7 +50,9 @@ defmodule Tesla.Middleware.Retry do
   - `:delay` - The base delay in milliseconds (positive integer, defaults to 50)
   - `:max_retries` - maximum number of retries (non-negative integer, defaults to 5)
   - `:max_delay` - maximum delay in milliseconds (positive integer, defaults to 5000)
-  - `:should_retry` - function to determine if request should be retried
+  - `:should_retry` - function with an arity of 1 or 3 used to determine if the request should
+      be retried the first argument is the result, the second is the env and the third is
+      the context: options + `:retries` (defaults to a match on `{:error, _reason}`)
   - `:jitter_factor` - additive noise proportionality constant
       (float between 0 and 1, defaults to 0.2)
   """
@@ -65,7 +75,7 @@ defmodule Tesla.Middleware.Retry do
       delay: integer_opt!(opts, :delay, 1),
       max_retries: integer_opt!(opts, :max_retries, 0),
       max_delay: integer_opt!(opts, :max_delay, 1),
-      should_retry: Keyword.get(opts, :should_retry, &match?({:error, _}, &1)),
+      should_retry: should_retry_opt!(opts),
       jitter_factor: float_opt!(opts, :jitter_factor, 0, 1)
     }
 
@@ -84,15 +94,26 @@ defmodule Tesla.Middleware.Retry do
   defp retry(env, next, context) do
     res = Tesla.run(env, next)
 
-    if context.should_retry.(res) do
-      backoff(context.max_delay, context.delay, context.retries, context.jitter_factor)
-      context = update_in(context, [:retries], &(&1 + 1))
-      retry(env, next, context)
-    else
-      res
+    {:arity, should_retry_arity} = :erlang.fun_info(context.should_retry, :arity)
+
+    cond do
+      should_retry_arity == 1 and context.should_retry.(res) ->
+        do_retry(env, next, context)
+
+      should_retry_arity == 3 and context.should_retry.(res, env, context) ->
+        do_retry(env, next, context)
+
+      true ->
+        res
     end
   end
 
+  defp do_retry(env, next, context) do
+    backoff(context.max_delay, context.delay, context.retries, context.jitter_factor)
+    context = update_in(context, [:retries], &(&1 + 1))
+    retry(env, next, context)
+  end
+
   # Exponential backoff with jitter
   defp backoff(cap, base, attempt, jitter_factor) do
     factor = Bitwise.bsl(1, attempt)
@@ -124,6 +145,19 @@ defmodule Tesla.Middleware.Retry do
     end
   end
 
+  defp should_retry_opt!(opts) do
+    case Keyword.get(opts, :should_retry, &match?({:error, _}, &1)) do
+      should_retry_fun when is_function(should_retry_fun, 1) ->
+        should_retry_fun
+
+      should_retry_fun when is_function(should_retry_fun, 3) ->
+        should_retry_fun
+
+      value ->
+        invalid_should_retry_fun(value)
+    end
+  end
+
   defp invalid_integer(key, value, min) do
     raise(ArgumentError, "expected :#{key} to be an integer >= #{min}, got #{inspect(value)}")
   end
@@ -134,4 +168,11 @@ defmodule Tesla.Middleware.Retry do
       "expected :#{key} to be a float >= #{min} and <= #{max}, got #{inspect(value)}"
     )
   end
+
+  defp invalid_should_retry_fun(value) do
+    raise(
+      ArgumentError,
+      "expected :should_retry to be a function with arity of 1 or 3, got #{inspect(value)}"
+    )
+  end
 end

test/tesla/middleware/retry_test.exs

@@ -9,6 +9,7 @@ defmodule Tesla.Middleware.RetryTest do
         response =
           case env.url do
             "/ok" -> {:ok, env}
+            "/maybe" when retries == 2 -> {:error, :nxdomain}
             "/maybe" when retries < 5 -> {:error, :econnrefused}
             "/maybe" -> {:ok, env}
             "/nope" -> {:error, :econnrefused}
@@ -39,9 +40,20 @@ defmodule Tesla.Middleware.RetryTest do
       delay: 10,
       max_retries: 10,
       should_retry: fn
-        {:ok, %{status: status}} when status in [400, 500] -> true
-        {:ok, _} -> false
-        {:error, _} -> true
+        {:ok, %{status: status}}, _env, _context when status in [400, 500] ->
+          true
+
+        {:ok, _reason}, _env, _context ->
+          false
+
+        {:error, _reason}, %Tesla.Env{method: :post}, _context ->
+          false
+
+        {:error, _reason}, %Tesla.Env{method: :put}, %{retries: 2} ->
+          false
+
+        {:error, _reason}, _env, _context ->
+          true
       end
 
     adapter LaggyAdapter
@@ -74,6 +86,14 @@ defmodule Tesla.Middleware.RetryTest do
              ClientWithShouldRetryFunction.get("/retry_status")
   end
 
+  test "use custom retry determination function matching on env" do
+    assert {:error, :econnrefused} = ClientWithShouldRetryFunction.post("/maybe", "payload")
+  end
+
+  test "use custom retry determination function matching on context" do
+    assert {:error, :nxdomain} = ClientWithShouldRetryFunction.put("/maybe", "payload")
+  end
+
   defmodule DefunctClient do
     use Tesla
 
@@ -171,4 +191,30 @@ defmodule Tesla.Middleware.RetryTest do
                    ClientWithJitterFactorGt1.get("/ok")
                  end
   end
+
+  test "ensures should_retry option is a function with arity of 1 or 3" do
+    defmodule ClientWithShouldRetryArity0 do
+      use Tesla
+      plug Tesla.Middleware.Retry, should_retry: fn -> true end
+      adapter LaggyAdapter
+    end
+
+    defmodule ClientWithShouldRetryArity2 do
+      use Tesla
+      plug Tesla.Middleware.Retry, should_retry: fn _res, _env -> true end
+      adapter LaggyAdapter
+    end
+
+    assert_raise ArgumentError,
+                 ~r/expected :should_retry to be a function with arity of 1 or 3, got #Function<\d.\d+\/0/,
+                 fn ->
+                   ClientWithShouldRetryArity0.get("/ok")
+                 end
+
+    assert_raise ArgumentError,
+                 ~r/expected :should_retry to be a function with arity of 1 or 3, got #Function<\d.\d+\/2/,
+                 fn ->
+                   ClientWithShouldRetryArity2.get("/ok")
+                 end
+  end
 end