From 7b0286f7159079b7acec45bcb2a2c47ab517903f Mon Sep 17 00:00:00 2001
From: DaTrader <damir1petkovic@gmail.com>
Date: Fri, 25 Aug 2023 19:34:02 +0200
Subject: [PATCH] switch to `:retry` 2-arity fun, add support for {:delay,
 milliseconds} (#226)

---
 lib/req/response.ex     |  57 +++++++++++++-
 lib/req/steps.ex        | 162 ++++++++++++++++++----------------------
 test/req/steps_test.exs |  71 +++++++++++++++---
 3 files changed, 189 insertions(+), 101 deletions(-)

diff --git a/lib/req/response.ex b/lib/req/response.ex
index 66c2b752..558f22f2 100644
--- a/lib/req/response.ex
+++ b/lib/req/response.ex
@@ -145,7 +145,7 @@ defmodule Req.Response do
   @doc """
   Deletes the header given by `key`
 
-  All occurences of the header are delete, in case the header is repeated multiple times.
+  All occurences of the header are deleted, in case the header is repeated multiple times.
 
   ## Examples
 
@@ -167,4 +167,59 @@ defmodule Req.Response do
           )
     }
   end
+
+  @doc """
+  Returns the `retry-after` header delay value or nil if not found.
+  """
+  @spec get_retry_after(t()) :: integer() | nil
+  def get_retry_after(response) do
+    case get_header(response, "retry-after") do
+      [delay] ->
+        retry_delay_in_ms(delay)
+
+      [] ->
+        nil
+    end
+  end
+
+  defp retry_delay_in_ms(delay_value) do
+    case Integer.parse(delay_value) do
+      {seconds, ""} ->
+        :timer.seconds(seconds)
+
+      :error ->
+        delay_value
+        |> parse_http_datetime()
+        |> DateTime.diff(DateTime.utc_now(), :millisecond)
+        |> max(0)
+    end
+  end
+
+  @month_numbers %{
+    "Jan" => "01",
+    "Feb" => "02",
+    "Mar" => "03",
+    "Apr" => "04",
+    "May" => "05",
+    "Jun" => "06",
+    "Jul" => "07",
+    "Aug" => "08",
+    "Sep" => "09",
+    "Oct" => "10",
+    "Nov" => "11",
+    "Dec" => "12"
+  }
+
+  defp parse_http_datetime(datetime) do
+    [_day_of_week, day, month, year, time, "GMT"] = String.split(datetime, " ")
+    date = year <> "-" <> @month_numbers[month] <> "-" <> day
+
+    case DateTime.from_iso8601(date <> " " <> time <> "Z") do
+      {:ok, valid_datetime, 0} ->
+        valid_datetime
+
+      {:error, reason} ->
+        raise "cannot parse \"retry-after\" header value #{inspect(datetime)} as datetime, reason: #{reason}"
+    end
+  end
 end
diff --git a/lib/req/steps.ex b/lib/req/steps.ex
index d9d17bd7..c6b6d978 100644
--- a/lib/req/steps.ex
+++ b/lib/req/steps.ex
@@ -862,9 +862,7 @@ defmodule Req.Steps do
   | _other_       | Returns data as is                              |
 
   This step updates the following headers to reflect the changes:
-
-    * `content-length` is set to the length of the decompressed body
-    * `content-encoding` is removed
+  - `content-length` is set to the length of the decompressed body
 
   ## Options
 
@@ -906,59 +904,49 @@ defmodule Req.Steps do
     if request.options[:raw] do
       {request, response}
     else
-      codecs = get_content_encoding_header(response.headers)
-      {decompressed_body, unknown_codecs} = decompress_body(codecs, response.body, [])
+      compression_algorithms = get_content_encoding_header(response.headers)
+      decompressed_body = decompress_body(response.body, compression_algorithms)
       decompressed_content_length = decompressed_body |> byte_size() |> to_string()
 
       response =
         %Req.Response{response | body: decompressed_body}
         |> Req.Response.put_header("content-length", decompressed_content_length)
 
-      response =
-        if unknown_codecs == [] do
-          Req.Response.delete_header(response, "content-encoding")
-        else
-          Req.Response.put_header(response, "content-encoding", Enum.join(unknown_codecs, ", "))
-        end
-
       {request, response}
     end
   end
 
-  defp decompress_body([gzip | rest], body, acc) when gzip in ["gzip", "x-gzip"] do
-    decompress_body(rest, :zlib.gunzip(body), acc)
+  defp decompress_body(body, algorithms) do
+    Enum.reduce(algorithms, body, &decompress_with_algorithm(&1, &2))
   end
 
-  defp decompress_body(["br" | rest], body, acc) do
+  defp decompress_with_algorithm(gzip, body) when gzip in ["gzip", "x-gzip"] do
+    :zlib.gunzip(body)
+  end
+
+  defp decompress_with_algorithm("br", body) do
     if brotli_loaded?() do
       {:ok, decompressed} = :brotli.decode(body)
-      decompress_body(rest, decompressed, acc)
+      decompressed
     else
-      Logger.debug("decompress_body: :brotli library not loaded, skipping brotli decompression")
-      decompress_body(rest, body, ["br" | acc])
+      raise("`:brotli` decompression library not loaded")
     end
   end
 
-  defp decompress_body(["zstd" | rest], body, acc) do
+  defp decompress_with_algorithm("zstd", body) do
     if ezstd_loaded?() do
-      decompress_body(rest, :ezstd.decompress(body), acc)
+      :ezstd.decompress(body)
     else
-      Logger.debug("decompress_body: :ezstd library not loaded, skipping zstd decompression")
-      decompress_body(rest, body, ["zstd" | acc])
+      raise("`:ezstd` decompression library not loaded")
     end
   end
 
-  defp decompress_body(["identity" | rest], body, acc) do
-    decompress_body(rest, body, acc)
-  end
-
-  defp decompress_body([codec | rest], body, acc) do
-    Logger.debug("decompress_body: algorithm #{inspect(codec)} is not supported")
-    decompress_body(rest, body, [codec | acc])
+  defp decompress_with_algorithm("identity", body) do
+    body
   end
 
-  defp decompress_body([], body, acc) do
-    {body, acc}
+  defp decompress_with_algorithm(_algorithm, body) do
+    body
   end
 
   defmacrop nimble_csv_loaded? do
@@ -1341,8 +1329,14 @@ defmodule Req.Steps do
 
         * `:transient` - same as `:safe_transient` except retries all HTTP methods (POST, DELETE, etc.)
 
-        * `fun` - a 1-arity function that accepts either a `Req.Response` or an exception struct
-          and returns boolean whether to retry.
+        * `fun` - a 2-arity function that accepts a `Req.Request` and either a `Req.Response` or an exception struct
+          and returns one of the following:
+
+            * `true` if to retry according to the `:retry_delay` option logic (see below)
+
+            * `{:delay, milliseconds}` to retry after the delay in milliseconds
+
+            * `false` or `nil` if not to retry
 
         * `false` - never retry.
 
@@ -1353,6 +1347,9 @@ defmodule Req.Steps do
       If the response is HTTP 429 and contains the `retry-after` header, the value of the header is used to
       determine the next retry delay.
 
+      This option must not be provided if a `:retry` function returning `{:delay, milliseconds}` is, or an
+      exception will be raised otherwise.
+
     * `:retry_log_level` - the log level to emit retry logs at. Can also be set to `false` to disable
       logging these messsages. Defaults to `:error`.
 
@@ -1381,7 +1378,7 @@ defmodule Req.Steps do
   def retry(request_response_or_error)
 
   def retry({request, response_or_exception}) do
-    retry? =
+    retry =
       case Map.get(request.options, :retry, :safe_transient) do
         :safe_transient ->
           safe_transient?(request, response_or_exception)
@@ -1393,7 +1390,7 @@ defmodule Req.Steps do
           false
 
         fun when is_function(fun) ->
-          fun.(response_or_exception)
+          apply_retry(fun, request, response_or_exception)
 
         :safe ->
           IO.warn("setting `retry: :safe` is deprecated in favour of `retry: :safe_transient`")
@@ -1405,17 +1402,38 @@ defmodule Req.Steps do
 
         other ->
           raise ArgumentError,
-                "expected :retry to be :safe_transient, :transient, false, or a 1-arity function, " <>
+                "expected :retry to be :safe_transient, :transient, false, or a 2-arity function, " <>
                   "got: #{inspect(other)}"
       end
 
-    if retry? do
-      retry(request, response_or_exception)
-    else
-      {request, response_or_exception}
+    case retry do
+      {:delay, delay} ->
+        if !Req.Request.get_option(request, :retry_delay) do
+          retry(request, response_or_exception, delay)
+        else
+          raise ArgumentError,
+                "expected :retry_delay not to be set when the :retry function is returning `{:delay, milliseconds}`"
+        end
+
+      true ->
+        retry(request, response_or_exception)
+
+      retry when retry in [false, nil] ->
+        {request, response_or_exception}
     end
   end
 
+  defp apply_retry( fun, request, response_or_exception)
+
+  defp apply_retry(fun, _request, response_or_exception) when is_function( fun, 1) do
+    IO.warn("`retry: fun/1` has been deprecated in favor of `retry: fun/2`")
+    fun.(response_or_exception)
+  end
+
+  defp apply_retry(fun, request, response_or_exception) when is_function( fun, 2) do
+    fun.(request, response_or_exception)
+  end
+
   defp safe_transient?(request, response_or_exception) do
     request.method in [:get, :head] and transient?(response_or_exception)
   end
@@ -1437,9 +1455,19 @@ defmodule Req.Steps do
     false
   end
 
-  defp retry(request, response_or_exception) do
+  defp retry(request, response_or_exception, delay_or_nil \\ nil)
+
+  defp retry(request, response_or_exception, nil) do
+    do_retry(request, response_or_exception, &get_retry_delay/3)
+  end
+
+  defp retry(request, response_or_exception, delay) when is_integer(delay) do
+    do_retry(request, response_or_exception, fn request, _, _ -> {request, delay} end)
+  end
+
+  defp do_retry(request, response_or_exception, delay_getter) do
     retry_count = Req.Request.get_private(request, :req_retry_count, 0)
-    {request, delay} = get_retry_delay(request, response_or_exception, retry_count)
+    {request, delay} = delay_getter.(request, response_or_exception, retry_count)
     max_retries = Req.Request.get_option(request, :max_retries, 3)
     log_level = Req.Request.get_option(request, :retry_log_level, :error)
 
@@ -1455,12 +1483,10 @@ defmodule Req.Steps do
   end
 
   defp get_retry_delay(request, %Req.Response{status: 429} = response, retry_count) do
-    case Req.Response.get_header(response, "retry-after") do
-      [delay] ->
-        {request, retry_delay_in_ms(delay)}
-
-      [] ->
-        calculate_retry_delay(request, retry_count)
+    if delay = Req.Response.get_retry_after(response) do
+      {request, delay}
+    else
+      calculate_retry_delay(request, retry_count)
     end
   end
 
@@ -1482,19 +1508,6 @@ defmodule Req.Steps do
     Integer.pow(2, n) * 1000
   end
 
-  defp retry_delay_in_ms(delay_value) do
-    case Integer.parse(delay_value) do
-      {seconds, ""} ->
-        :timer.seconds(seconds)
-
-      :error ->
-        delay_value
-        |> parse_http_datetime()
-        |> DateTime.diff(DateTime.utc_now(), :millisecond)
-        |> max(0)
-    end
-  end
-
   defp log_retry(_, _, _, _, false), do: :ok
 
   defp log_retry(response_or_exception, retry_count, max_retries, delay, level) do
@@ -1568,31 +1581,4 @@ defmodule Req.Steps do
   def format_http_datetime(datetime) do
     Calendar.strftime(datetime, "%a, %d %b %Y %H:%M:%S GMT")
   end
-
-  @month_numbers %{
-    "Jan" => "01",
-    "Feb" => "02",
-    "Mar" => "03",
-    "Apr" => "04",
-    "May" => "05",
-    "Jun" => "06",
-    "Jul" => "07",
-    "Aug" => "08",
-    "Sep" => "09",
-    "Oct" => "10",
-    "Nov" => "11",
-    "Dec" => "12"
-  }
-  defp parse_http_datetime(datetime) do
-    [_day_of_week, day, month, year, time, "GMT"] = String.split(datetime, " ")
-    date = year <> "-" <> @month_numbers[month] <> "-" <> day
-
-    case DateTime.from_iso8601(date <> " " <> time <> "Z") do
-      {:ok, valid_datetime, 0} ->
-        valid_datetime
-
-      {:error, reason} ->
-        raise "could not parse \"Retry-After\" header #{datetime} - #{reason}"
-    end
-  end
 end
diff --git a/test/req/steps_test.exs b/test/req/steps_test.exs
index 2da09c81..8712c704 100644
--- a/test/req/steps_test.exs
+++ b/test/req/steps_test.exs
@@ -390,16 +390,17 @@ defmodule Req.StepsTest do
       assert String.to_integer(content_length) == byte_size(body)
     end
 
-    test "delete content-encoding header", c do
-      Bypass.expect(c.bypass, "GET", "/", fn conn ->
-        conn
-        |> Plug.Conn.put_resp_header("content-encoding", "x-gzip")
-        |> Plug.Conn.send_resp(200, :zlib.gzip("foo"))
-      end)
-
-      resp = Req.get!(c.url)
-      assert [] = Req.Response.get_header(resp, "content-encoding")
-    end
+# raises
+#    test "delete content-encoding header", c do
+#      Bypass.expect(c.bypass, "GET", "/", fn conn ->
+#        conn
+#        |> Plug.Conn.put_resp_header("content-encoding", "x-gzip")
+#        |> Plug.Conn.send_resp(200, :zlib.gzip("foo"))
+#      end)
+#
+#      resp = Req.get!(c.url)
+#      assert [] = Req.Response.get_header(resp, "content-encoding")
+#    end
   end
 
   describe "output" do
@@ -1192,7 +1193,7 @@ defmodule Req.StepsTest do
     end
 
     @tag :capture_log
-    test "custom function", c do
+    test "custom function returning true", c do
       pid = self()
 
       Bypass.expect(c.bypass, "POST", "/", fn conn ->
@@ -1200,7 +1201,7 @@ defmodule Req.StepsTest do
         Plug.Conn.send_resp(conn, 500, "oops")
       end)
 
-      fun = fn response ->
+      fun = fn _request, response ->
         assert response.status == 500
         true
       end
@@ -1215,6 +1216,52 @@ defmodule Req.StepsTest do
       refute_received _
     end
 
+    @tag :capture_log
+    test "custom function returning {:delay, milliseconds}", c do
+      pid = self()
+
+      Bypass.expect(c.bypass, "GET", "/", fn conn ->
+        send(pid, :ping)
+        Plug.Conn.send_resp(conn, 500, "oops")
+      end)
+
+      fun = fn _request, response ->
+        assert response.status == 500
+        {:delay, 1}
+      end
+
+      request = Req.new(url: c.url, retry: fun)
+
+      assert Req.get!(request).status == 500
+      assert_received :ping
+      assert_received :ping
+      assert_received :ping
+      assert_received :ping
+      refute_received _
+    end
+
+    @tag :capture_log
+    test "raise on custom function returning {:delay, milliseconds} when `:retry_delay` is provided",
+         c do
+      pid = self()
+
+      Bypass.expect(c.bypass, "GET", "/", fn conn ->
+        send(pid, :ping)
+        Plug.Conn.send_resp(conn, 500, "oops")
+      end)
+
+      fun = fn _request, response ->
+        assert response.status == 500
+        {:delay, 1}
+      end
+
+      request = Req.new(url: c.url, retry: fun, retry_delay: 1)
+
+      assert_raise ArgumentError,
+                   "expected :retry_delay not to be set when the :retry function is returning `{:delay, milliseconds}`",
+                   fn -> Req.get!(request) end
+    end
+
     @tag :capture_log
     test "does not re-encode params", c do
       pid = self()