From 0eca531e64b22d7ad16ce9b38eb59d2d8ed93b07 Mon Sep 17 00:00:00 2001
From: Thales Macedo Garitezi <thalesmg@gmail.com>
Date: Tue, 25 Oct 2022 15:01:05 -0300
Subject: [PATCH] feat: add `retry-after` headers to unavailable response

---
 .../src/emqx_mgmt_api_status.erl              |  6 +-
 .../test/emqx_mgmt_api_status_SUITE.erl       | 60 ++++++++++++-------
 2 files changed, 45 insertions(+), 21 deletions(-)

diff --git a/apps/emqx_management/src/emqx_mgmt_api_status.erl b/apps/emqx_management/src/emqx_mgmt_api_status.erl
index 37a8274f8..b313f4bcc 100644
--- a/apps/emqx_management/src/emqx_mgmt_api_status.erl
+++ b/apps/emqx_management/src/emqx_mgmt_api_status.erl
@@ -50,7 +50,11 @@ running_status() ->
                     running -> 200;
                     not_running -> 503
                 end,
-            {StatusCode, #{<<"content-type">> => <<"text/plain">>}, list_to_binary(Body)};
+            Headers = #{
+                <<"content-type">> => <<"text/plain">>,
+                <<"retry-after">> => <<"15">>
+            },
+            {StatusCode, Headers, list_to_binary(Body)};
         false ->
             {503, #{<<"retry-after">> => <<"15">>}, <<>>}
     end.
diff --git a/apps/emqx_management/test/emqx_mgmt_api_status_SUITE.erl b/apps/emqx_management/test/emqx_mgmt_api_status_SUITE.erl
index 4051bd62c..42d833bda 100644
--- a/apps/emqx_management/test/emqx_mgmt_api_status_SUITE.erl
+++ b/apps/emqx_management/test/emqx_mgmt_api_status_SUITE.erl
@@ -54,28 +54,43 @@ end_per_testcase(_TestCase, _Config) ->
 
 do_request(Opts) ->
     #{
-        path := Path,
+        path := Path0,
         method := Method,
         headers := Headers,
         body := Body0
     } = Opts,
-    URL = ?HOST ++ filename:join(Path),
+    URL = ?HOST ++ filename:join(Path0),
+    {ok, #{host := Host, port := Port, path := Path}} = emqx_http_lib:uri_parse(URL),
+    %% we must not use `httpc' here, because it keeps retrying when it
+    %% receives a 503 with `retry-after' header, and there's no option
+    %% to stop that behavior...
+    {ok, Gun} = gun:open(Host, Port, #{retry => 0}),
+    {ok, http} = gun:await_up(Gun),
     Request =
-        case Body0 of
-            no_body -> {URL, Headers};
-            {Encoding, Body} -> {URL, Headers, Encoding, Body}
+        fun() ->
+            case Body0 of
+                no_body -> gun:Method(Gun, Path, Headers);
+                {_Encoding, Body} -> gun:Method(Gun, Path, Headers, Body)
+            end
         end,
-    ct:pal("Method: ~p, Request: ~p", [Method, Request]),
-    case httpc:request(Method, Request, [], []) of
-        {error, socket_closed_remotely} ->
-            {error, socket_closed_remotely};
-        {ok, {{_, StatusCode, _}, Headers1, Body1}} ->
-            Body2 =
-                case emqx_json:safe_decode(Body1, [return_maps]) of
-                    {ok, Json} -> Json;
-                    {error, _} -> Body1
-                end,
-            {ok, #{status_code => StatusCode, headers => Headers1, body => Body2}}
+    Ref = Request(),
+    receive
+        {gun_response, Gun, Ref, nofin, StatusCode, Headers1} ->
+            Data = data_loop(Gun, Ref, _Acc = <<>>),
+            #{status_code => StatusCode, headers => maps:from_list(Headers1), body => Data}
+    after 5_000 ->
+        error({timeout, Opts, process_info(self(), messages)})
+    end.
+
+data_loop(Gun, Ref, Acc) ->
+    receive
+        {gun_data, Gun, Ref, nofin, Data} ->
+            data_loop(Gun, Ref, <<Acc/binary, Data/binary>>);
+        {gun_data, Gun, Ref, fin, Data} ->
+            gun:shutdown(Gun),
+            <<Acc/binary, Data/binary>>
+    after 5000 ->
+        error(timeout)
     end.
 
 %%---------------------------------------------------------------------------------------
@@ -83,10 +98,10 @@ do_request(Opts) ->
 %%---------------------------------------------------------------------------------------
 
 t_status_ok(_Config) ->
-    {ok, #{
+    #{
         body := Resp,
         status_code := StatusCode
-    }} = do_request(#{
+    } = do_request(#{
         method => get,
         path => ["status"],
         headers => [],
@@ -100,10 +115,11 @@ t_status_ok(_Config) ->
     ok.
 
 t_status_not_ok(_Config) ->
-    {ok, #{
+    #{
         body := Resp,
+        headers := Headers,
         status_code := StatusCode
-    }} = do_request(#{
+    } = do_request(#{
         method => get,
         path => ["status"],
         headers => [],
@@ -114,4 +130,8 @@ t_status_not_ok(_Config) ->
         {match, _},
         re:run(Resp, <<"emqx is not_running$">>)
     ),
+    ?assertMatch(
+        #{<<"retry-after">> := <<"15">>},
+        Headers
+    ),
     ok.