feat: add support for should_retry with arity of 2 to Retry middleware

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,13 @@ defmodule Tesla.Middleware.Retry do
         {:ok, _} -> false
         {:error, _} -> true
       end
+      # or
+      plug Tesla.Middleware.Retry, should_retry: fn
+        {:ok, %{status: status}}, _env when status in [400, 500] -> true
+        {:ok, _reason}, _env -> false
+        {:error, _reason}, %Tesla.Env{method: :post} -> false
+        {:error, _reason}, _env -> true
+      end
   end
   ```
 
@@ -42,7 +49,8 @@ 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 arity of 1 or 2 to determine if request should be retried
+      arity of 1 passes the result, while arity of 2 passes env as second argument (defaults to a match on `{:error, _reason}`)
   - `:jitter_factor` - additive noise proportionality constant
       (float between 0 and 1, defaults to 0.2)
   """
@@ -65,7 +73,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 +92,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 == 2 and context.should_retry.(res, env) ->
+        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 +143,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, 2) ->
+        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 +166,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 2, got #{inspect(value)}"
+    )
+  end
 end

test/tesla/middleware/retry_test.exs

@@ -39,9 +39,10 @@ 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 when status in [400, 500] -> true
+        {:ok, _reason}, _env -> false
+        {:error, _reason}, %Tesla.Env{method: :post} -> false
+        {:error, _reason}, _env -> true
       end
 
     adapter LaggyAdapter
@@ -74,6 +75,10 @@ defmodule Tesla.Middleware.RetryTest do
              ClientWithShouldRetryFunction.get("/retry_status")
   end
 
+  test "use custom retry determination function matching on method" do
+    assert {:error, :econnrefused} = ClientWithShouldRetryFunction.post("/maybe", "payload")
+  end
+
   defmodule DefunctClient do
     use Tesla
 
@@ -171,4 +176,30 @@ defmodule Tesla.Middleware.RetryTest do
                    ClientWithJitterFactorGt1.get("/ok")
                  end
   end
+
+  test "ensures should_retry option is a function with arity of 1 or 2" do
+    defmodule ClientWithShouldRetryArity0 do
+      use Tesla
+      plug Tesla.Middleware.Retry, should_retry: fn -> true end
+      adapter LaggyAdapter
+    end
+
+    defmodule ClientWithShouldRetryArity3 do
+      use Tesla
+      plug Tesla.Middleware.Retry, should_retry: fn _res, _env, _other -> true end
+      adapter LaggyAdapter
+    end
+
+    assert_raise ArgumentError,
+                 ~r/expected :should_retry to be a function with arity of 1 or 2, 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 2, got #Function<\d.\d+\/3/,
+                 fn ->
+                   ClientWithShouldRetryArity3.get("/ok")
+                 end
+  end
 end