Fix retrying on disconnect with retry=1

2 files changed, 50 insertions, 34 deletions

diff --git a/src/gun.erl b/src/gun.erl
index a72cc85..507eef7 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -891,6 +891,8 @@ init({Owner, Host, Port, Opts}) ->
 default_transport(443) -> tls;
 default_transport(_) -> tcp.
 
+not_connected(_, {retries, 0, normal}, State) ->
+	{stop, normal, State};
 not_connected(_, {retries, 0, Reason}, State) ->
 	{stop, {shutdown, Reason}, State};
 not_connected(_, {retries, Retries0, _}, State=#state{opts=Opts}) ->
@@ -907,9 +909,14 @@ not_connected(Type, Event, State) ->
 	handle_common(Type, Event, ?FUNCTION_NAME, State).
 
 default_retry_fun(Retries, Opts) ->
+	%% We retry immediately after a disconnect.
+	Timeout = case maps:get(retry, Opts, 5) of
+		Retries -> 0;
+		_ -> maps:get(retry_timeout, Opts, 5000)
+	end,
 	#{
 		retries => Retries - 1,
-		timeout => maps:get(retry_timeout, Opts, 5000)
+		timeout => Timeout
 	}.
 
 domain_lookup(_, {retries, Retries, _}, State=#state{host=Host, port=Port, opts=Opts,
@@ -1427,17 +1434,10 @@ disconnect(State0=#state{owner=Owner, status=Status, opts=Opts,
 			KilledStreams = Protocol:down(ProtoState),
 			Owner ! {gun_down, self(), Protocol:name(), Reason, KilledStreams},
 			Retry = maps:get(retry, Opts, 5),
-			case Retry of
-				0 when Reason =:= normal ->
-					{stop, normal, State};
-				0 ->
-					{stop, {shutdown, Reason}, State};
-				_ ->
-					{next_state, not_connected,
-						keepalive_cancel(State#state{socket=undefined,
-							protocol=undefined, protocol_state=undefined}),
-						{next_event, internal, {retries, Retry - 1, Reason}}}
-			end
+			{next_state, not_connected,
+				keepalive_cancel(State#state{socket=undefined,
+					protocol=undefined, protocol_state=undefined}),
+				{next_event, internal, {retries, Retry, Reason}}}
 	end.
 
 disconnect_flush(State=#state{socket=Socket, messages={OK, Closed, Error}}) ->
diff --git a/test/gun_SUITE.erl b/test/gun_SUITE.erl
index cd6b419..2fbbb05 100644
--- a/test/gun_SUITE.erl
+++ b/test/gun_SUITE.erl
@@ -309,6 +309,29 @@ retry_0(_) ->
 		error(timeout)
 	end.
 
+retry_0_disconnect(_) ->
+	doc("Ensure Gun gives up immediately with retry=0 after a successful connection."),
+	{ok, ListenSocket} = gen_tcp:listen(0, [binary, {active, false}]),
+	{ok, {_, Port}} = inet:sockname(ListenSocket),
+	{ok, Pid} = gun:open("localhost", Port, #{retry => 0, retry_timeout => 500}),
+	Ref = monitor(process, Pid),
+	%% On Windows when the connection is refused the OS will retry
+	%% 3 times before giving up, with a 500ms delay between tries.
+	%% This adds approximately 1 second to connection failures.
+	After = case os:type() of
+		{win32, _} -> 1200;
+		_ -> 200
+	end,
+	%% We accept the connection and then close it to trigger a disconnect.
+	{ok, ClientSocket} = gen_tcp:accept(ListenSocket, 5000),
+	gen_tcp:close(ClientSocket),
+	receive
+		{'DOWN', Ref, process, Pid, {shutdown, _}} ->
+			ok
+	after After ->
+		error(timeout)
+	end.
+
 retry_1(_) ->
 	doc("Ensure Gun gives up with retry=1."),
 	{ok, Pid} = gun:open("localhost", 12345, #{retry => 1, retry_timeout => 500}),
@@ -324,6 +347,18 @@ retry_1(_) ->
 		error(timeout)
 	end.
 
+retry_1_disconnect(_) ->
+	doc("Ensure Gun gives up with retry=1 after a successful connection."),
+	{ok, ListenSocket} = gen_tcp:listen(0, [binary, {active, false}]),
+	{ok, {_, Port}} = inet:sockname(ListenSocket),
+	{ok, Pid} = gun:open("localhost", Port, #{retry => 1, retry_timeout => 500}),
+	%% We accept the connection and then close it to trigger a disconnect.
+	{ok, ClientSocket} = gen_tcp:accept(ListenSocket, 5000),
+	gen_tcp:close(ClientSocket),
+	%% We confirm that Gun reconnects immediately.
+	{ok, _} = gen_tcp:accept(ListenSocket, 200),
+	gun:close(Pid).
+
 retry_fun(_) ->
 	doc("Ensure the retry_fun is used when provided."),
 	{ok, Pid} = gun:open("localhost", 12345, #{
@@ -343,29 +378,10 @@ retry_fun(_) ->
 		error(shutdown_too_late)
 	end.
 
-retry_immediately(_) ->
-	doc("Ensure Gun retries immediately."),
-	%% We have to make a first successful connection in order to test this.
-	{ok, _, OriginPort} = init_origin(tcp, http,
-		fun(_, ClientSocket, ClientTransport) ->
-			ClientTransport:close(ClientSocket)
-		end),
-	{ok, Pid} = gun:open("localhost", OriginPort, #{retry => 1, retry_timeout => 500}),
-	Ref = monitor(process, Pid),
-	After = case os:type() of
-		{win32, _} -> 1200;
-		_ -> 200
-	end,
-	receive
-		{'DOWN', Ref, process, Pid, {shutdown, _}} ->
-			ok
-	after After ->
-		error(timeout)
-	end.
-
 retry_timeout(_) ->
-	doc("Ensure the retry_timeout value is enforced."),
-	{ok, Pid} = gun:open("localhost", 12345, #{retry => 1, retry_timeout => 1000}),
+	doc("Ensure the retry_timeout value is enforced. The first retry is immediate "
+		"and therefore does not use the timeout."),
+	{ok, Pid} = gun:open("localhost", 12345, #{retry => 2, retry_timeout => 1000}),
 	Ref = monitor(process, Pid),
 	After = case os:type() of
 		{win32, _} -> 2800;