Add feature to allow disabling reconnection. (#203)

* Add feature to allow disabling reconnection.

* Address PR feedback.

* Add newline to end of test file.

Author: kinyoklion

src/shotgun.erl

@@ -41,7 +41,8 @@
     #{async => boolean(),
       async_mode => binary | sse,
       handle_event => fun((fin | nofin, reference(), binary()) -> any()),
-      timeout => timeout()}. %% Default 5000 ms
+      timeout => timeout(), %% Default 5000 ms
+      allow_reconnect => boolean()}. %% Default true
 -type response() ::
     #{status_code => integer(),
       headers => proplists:proplist(),
@@ -241,13 +242,15 @@ request(Pid, get, Uri, Headers0, Body, Options) ->
           async := IsAsync,
           async_mode := AsyncMode,
           headers := Headers,
-          timeout := Timeout} =
+          timeout := Timeout,
+          allow_reconnect := AllowReconnect
+        } =
             process_options(Options, Headers0, get),
 
         Event =
             case IsAsync of
                 true ->
-                    {get_async, {HandleEvent, AsyncMode}, {Uri, Headers, Body}};
+                    {get_async, {HandleEvent, AsyncMode, AllowReconnect}, {Uri, Headers, Body}};
                 false ->
                     {get, {Uri, Headers, Body}}
             end,
@@ -320,7 +323,9 @@ parse_event(EventBin) ->
       pid => pid() | undefined,
       responses => queue:queue() | undefined,
       status_code => integer() | undefined,
-      stream => reference() | undefined}.
+      stream => reference() | undefined,
+      allow_reconnect => boolean()
+    }.
 
 %% @private
 -spec callback_mode() -> state_functions.
@@ -392,12 +397,17 @@ handle_info({gun_up, Pid, _Protocol}, StateName, StateData = #{pid := Pid}) ->
     {next_state, StateName, StateData};
 handle_info({gun_down, Pid, Protocol, Reason, KilledStreams},
             _StateName,
-            StateData = #{pid := Pid}) ->
-    error_logger:warning_msg("~p connection down on ~p: ~p (Killed: ~p)",
-                             [Protocol, Pid, Reason, KilledStreams]),
-    gun:shutdown(Pid),
-    CleanStateData = clean_state_data(StateData),
-    {next_state, down, CleanStateData};
+            StateData = #{pid := Pid, allow_reconnect := AllowReconnect }) ->
+    error_logger:warning_msg("~p connection down on ~p: ~p (Killed: ~p) (AllowReconnect: ~p)",
+                            [Protocol, Pid, Reason, KilledStreams, AllowReconnect]),
+    case AllowReconnect of
+        true ->
+            gun:shutdown(Pid),
+            CleanStateData = clean_state_data(StateData),
+            {next_state, down, CleanStateData};
+        false ->
+            keep_state_and_data
+    end;
 handle_info(Event, StateName, StateData) ->
     Module = ?MODULE,
     % forward messages to state machine
@@ -450,7 +460,7 @@ at_rest(cast, {gun_down, _Args, _From}, StateData = #{pid := Pid}) ->
     CleanStateData = clean_state_data(StateData),
     {next_state, down, CleanStateData};
 at_rest(cast,
-        {get_async, {HandleEvent, AsyncMode}, Args, From},
+        {get_async, {HandleEvent, AsyncMode, AllowReconnect}, Args, From},
         StateData = #{pid := Pid}) ->
     StreamRef = do_http_verb(get, Pid, Args),
     CleanStateData = clean_state_data(StateData),
@@ -460,7 +470,8 @@ at_rest(cast,
                         stream => StreamRef,
                         handle_event => HandleEvent,
                         async => true,
-                        async_mode => AsyncMode},
+                        async_mode => AsyncMode,
+                        allow_reconnect => AllowReconnect},
     {next_state, wait_response, NewStateData};
 at_rest(cast, {HttpVerb, {_, _, Body} = Args, From}, StateData = #{pid := Pid}) ->
     StreamRef = do_http_verb(HttpVerb, Pid, Args),
@@ -641,6 +652,7 @@ clean_state_data() ->
 clean_state_data(StateData) ->
     Responses = maps:get(responses, StateData, queue:new()),
     Requests = maps:get(pending_requests, StateData, queue:new()),
+    AllowReconnect = maps:get(allow_reconnect, StateData, true),
     #{pid => undefined,
       stream => undefined,
       handle_event => undefined,
@@ -652,7 +664,8 @@ clean_state_data(StateData) ->
       async => false,
       async_mode => binary,
       buffer => <<"">>,
-      pending_requests => Requests}.
+      pending_requests => Requests,
+      allow_reconnect => AllowReconnect}.
 
 %% @private
 -spec do_http_verb(http_verb(), pid(), tuple()) -> reference().
@@ -711,6 +724,7 @@ process_options(Options, Headers0, HttpVerb) ->
     Async = maps:get(async, Options, false),
     AsyncMode = maps:get(async_mode, Options, binary),
     Timeout = maps:get(timeout, Options, 5000),
+    AllowReconnect = maps:get(allow_reconnect, Options, true),
     case {Async, HttpVerb} of
         {true, get} ->
             ok;
@@ -723,7 +737,8 @@ process_options(Options, Headers0, HttpVerb) ->
       async => Async,
       async_mode => AsyncMode,
       headers => Headers,
-      timeout => Timeout}.
+      timeout => Timeout,
+      allow_reconnect => AllowReconnect}.
 
 %% @private
 -spec basic_auth_header(headers()) -> proplists:proplist().
@@ -809,8 +824,8 @@ enqueue_work_or_stop(_StateName, Event, From, StateData, Timeout) ->
 
 %% @private
 -spec create_work({atom(), list()}, gen_statem:from()) -> not_work | {ok, work()}.
-create_work({M = get_async, {HandleEvent, AsyncMode}, Args}, From) ->
-    {ok, {M, {HandleEvent, AsyncMode}, Args, From}};
+create_work({M = get_async, {HandleEvent, AsyncMode, AllowReconnect}, Args}, From) ->
+    {ok, {M, {HandleEvent, AsyncMode, AllowReconnect}, Args, From}};
 create_work({M, Args}, From)
     when M == get
          orelse M == post

test/shotgun_async_SUITE.erl

@@ -2,7 +2,15 @@
 
 -export([all/0, init_per_suite/1, end_per_suite/1, init_per_testcase/2,
          end_per_testcase/2]).
--export([get_sse/1, get_binary/1, work_queue/1, get_handle_event/1, async_unsupported/1]).
+-export([
+    get_sse/1,
+    get_binary/1,
+    work_queue/1,
+    get_handle_event/1,
+    async_unsupported/1,
+    async_gun_down_no_reopen/1,
+    async_gun_down_with_reopen/1
+]).
 
 -include_lib("common_test/include/ct.hrl").
 
@@ -163,3 +171,86 @@ async_unsupported(Config) ->
     {error, {async_unsupported, put}} = shotgun:put(Conn, "/", #{}, <<>>, #{async => true}),
 
     {comment, ""}.
+
+-spec async_gun_down_no_reopen(shotgun_test_utils:config()) -> {comment, string()}.
+async_gun_down_no_reopen(Config) ->
+    Conn = ?config(conn, Config),
+    Self = self(),
+
+    ct:comment("SSE: GET should return a ref"),
+    HandleEvent =
+        fun(_, _, EventBin) ->
+           case shotgun:parse_event(EventBin) of
+               #{id := Data} ->
+                   Self ! Data;
+               _ ->
+                   ok
+           end
+        end,
+    Opts =
+        #{async => true,
+          async_mode => sse,
+          allow_reconnect => false,
+          handle_event => HandleEvent},
+    {ok, _Ref} = shotgun:get(Conn, <<"/chunked-sse/3">>, #{}, Opts),
+
+    timer:sleep(500),
+
+    ok = shotgun_test_utils:wait_receive(<<"1">>, 500),
+    {at_rest, _} = sys:get_state(Conn),
+    %% Server is stopped immediately after receiving event.
+    %% We should be at rest waiting for the next event.
+    http_server:stop_listener(),
+    timer:sleep(500),
+    %% Listener needs to be running at test end.
+    http_server:start_listener(),
+
+    ct:comment("Connection should be closed after server shutdown"),
+    false = erlang:is_process_alive(Conn),
+
+    {comment, ""}.
+
+-spec async_gun_down_with_reopen(shotgun_test_utils:config()) -> {comment, string()}.
+async_gun_down_with_reopen(Config) ->
+    Conn = ?config(conn, Config),
+    Self = self(),
+
+    ct:comment("SSE: GET should return a ref"),
+    HandleEvent =
+        fun(_, _, EventBin) ->
+           case shotgun:parse_event(EventBin) of
+               #{id := Data} ->
+                   Self ! Data;
+               _ ->
+                   ok
+           end
+        end,
+    Opts =
+        #{async => true,
+          async_mode => sse,
+          allow_reconnect => true,
+          handle_event => HandleEvent},
+    {ok, _Ref} = shotgun:get(Conn, <<"/chunked-sse/3">>, #{}, Opts),
+
+    timer:sleep(500),
+
+    ok = shotgun_test_utils:wait_receive(<<"1">>, 500),
+    {at_rest, _} = sys:get_state(Conn),
+    %% Server is stopped immediately after receiving event.
+    %% We should be at rest waiting for the next event.
+    http_server:stop_listener(),
+    timer:sleep(500),
+
+    %% Listener needs to be running at test end.
+    http_server:start_listener(),
+
+    ct:comment("Connection should still be alive"),
+    true = erlang:is_process_alive(Conn),
+
+    Port = application:get_env(http_server, http_port, 8888),
+    shotgun:reopen(Conn, "localhost", Port),
+
+    {ok, _Ref2} = shotgun:get(Conn, <<"/chunked-sse/3">>, #{}, Opts),
+    ok = shotgun_test_utils:wait_receive(<<"1">>, 500),
+
+    {comment, ""}.