Add gun_up and gun_down messages

The flush(Pid) function was enhanced to also discard Websocket
messages and the new up/down messages.

4 files changed, 59 insertions, 20 deletions

diff --git a/src/gun.erl b/src/gun.erl
index 4bb2806..b3623eb 100644
--- a/src/gun.erl
+++ b/src/gun.erl
@@ -341,6 +341,10 @@ flush(StreamRef) ->
 
 flush_pid(ServerPid) ->
 	receive
+		{gun_up, ServerPid, _} ->
+			flush_pid(ServerPid);
+		{gun_down, ServerPid, _, _, _, _} ->
+			flush_pid(ServerPid);
 		{gun_response, ServerPid, _, _, _, _} ->
 			flush_pid(ServerPid);
 		{gun_data, ServerPid, _, _, _} ->
@@ -351,6 +355,10 @@ flush_pid(ServerPid) ->
 			flush_pid(ServerPid);
 		{gun_error, ServerPid, _} ->
 			flush_pid(ServerPid);
+		{gun_ws_upgrade, ServerPid, _} ->
+			flush_pid(ServerPid);
+		{gun_ws, ServerPid, _} ->
+			flush_pid(ServerPid);
 		{'DOWN', _, process, ServerPid, _} ->
 			flush_pid(ServerPid)
 	after 0 ->
@@ -424,7 +432,7 @@ init(Parent, Owner, Host, Port, Opts) ->
 default_transport(443) -> ssl;
 default_transport(_) -> tcp.
 
-connect(State=#state{owner=Owner, host=Host, port=Port, opts=Opts, transport=Transport=ranch_ssl}, Retries) ->
+connect(State=#state{host=Host, port=Port, opts=Opts, transport=Transport=ranch_ssl}, Retries) ->
 	Protocols = lists:flatten([case P of
 		http -> <<"http/1.1">>;
 		spdy -> [<<"spdy/3.1">>, <<"spdy/3">>]
@@ -438,14 +446,11 @@ connect(State=#state{owner=Owner, host=Host, port=Port, opts=Opts, transport=Tra
 				{ok, <<"spdy/3", _/bits>>} -> {gun_spdy, spdy_opts};
 				_ -> {gun_http, http_opts}
 			end,
-			ProtoOpts = maps:get(ProtoOptsKey, Opts, #{}),
-			ProtoState = Protocol:init(Owner, Socket, Transport, ProtoOpts),
-			before_loop(State#state{socket=Socket, transport=Transport,
-				protocol=Protocol, protocol_state=ProtoState});
+			up(State, Socket, Protocol, ProtoOptsKey);
 		{error, _} ->
 			retry(State, Retries - 1)
 	end;
-connect(State=#state{owner=Owner, host=Host, port=Port, opts=Opts, transport=Transport}, Retries) ->
+connect(State=#state{host=Host, port=Port, opts=Opts, transport=Transport}, Retries) ->
 	TransportOpts = [binary, {active, false}
 		|maps:get(transport_opts, Opts, [])],
 	case Transport:connect(Host, Port, TransportOpts) of
@@ -454,15 +459,20 @@ connect(State=#state{owner=Owner, host=Host, port=Port, opts=Opts, transport=Tra
 				[http] -> {gun_http, http_opts};
 				[spdy] -> {gun_spdy, spdy_opts}
 			end,
-			ProtoOpts = maps:get(ProtoOptsKey, Opts, #{}),
-			ProtoState = Protocol:init(Owner, Socket, Transport, ProtoOpts),
-			before_loop(State#state{socket=Socket, transport=Transport,
-				protocol=Protocol, protocol_state=ProtoState});
+			up(State, Socket, Protocol, ProtoOptsKey);
 		{error, _} ->
 			retry(State, Retries - 1)
 	end.
 
-retry(State=#state{opts=Opts}) ->
+up(State=#state{owner=Owner, opts=Opts, transport=Transport}, Socket, Protocol, ProtoOptsKey) ->
+	ProtoOpts = maps:get(ProtoOptsKey, Opts, #{}),
+	ProtoState = Protocol:init(Owner, Socket, Transport, ProtoOpts),
+	Owner ! {gun_up, self(), Protocol:name()},
+	before_loop(State#state{socket=Socket, protocol=Protocol, protocol_state=ProtoState}).
+
+down(State=#state{owner=Owner, opts=Opts, protocol=Protocol, protocol_state=ProtoState}, Reason) ->
+	{KilledStreams, UnprocessedStreams} = Protocol:down(ProtoState),
+	Owner ! {gun_down, self(), Protocol:name(), Reason, KilledStreams, UnprocessedStreams},
 	retry(State#state{socket=undefined, protocol=undefined, protocol_state=undefined},
 		maps:get(retry, Opts, 5)).
 
@@ -513,7 +523,7 @@ loop(State=#state{parent=Parent, owner=Owner, host=Host, port=Port,
 			case Protocol:handle(Data, ProtoState) of
 				close ->
 					Transport:close(Socket),
-					retry(State);
+					down(State, normal);
 				{upgrade, Protocol2, ProtoState2} ->
 					ws_loop(State#state{protocol=Protocol2, protocol_state=ProtoState2});
 			ProtoState2 ->
@@ -522,11 +532,11 @@ loop(State=#state{parent=Parent, owner=Owner, host=Host, port=Port,
 		{Closed, Socket} ->
 			Protocol:close(ProtoState),
 			Transport:close(Socket),
-			retry(State);
-		{Error, Socket, _} ->
+			down(State, closed);
+		{Error, Socket, Reason} ->
 			Protocol:close(ProtoState),
 			Transport:close(Socket),
-			retry(State);
+			down(State, {error, Reason});
 		{OK, _PreviousSocket, _Data} ->
 			loop(State);
 		{Closed, _PreviousSocket} ->
@@ -591,16 +601,16 @@ ws_loop(State=#state{parent=Parent, owner=Owner, socket=Socket,
 			case Protocol:handle(Data, ProtoState) of
 				close ->
 					Transport:close(Socket),
-					retry(State);
+					down(State, normal);
 				ProtoState2 ->
 					ws_loop(State#state{protocol_state=ProtoState2})
 			end;
 		{Closed, Socket} ->
 			Transport:close(Socket),
-			retry(State);
-		{Error, Socket, _} ->
+			down(State, closed);
+		{Error, Socket, Reason} ->
 			Transport:close(Socket),
-			retry(State);
+			down(State, {error, Reason});
 		%% Ignore any previous HTTP keep-alive.
 		keepalive ->
 			ws_loop(State);
@@ -610,7 +620,7 @@ ws_loop(State=#state{parent=Parent, owner=Owner, socket=Socket,
 			case Protocol:send(Frame, ProtoState) of
 				close ->
 					Transport:close(Socket),
-					retry(State);
+					down(State, normal);
 				ProtoState2 ->
 					ws_loop(State#state{protocol_state=ProtoState2})
 			end;
diff --git a/src/gun_http.erl b/src/gun_http.erl
index 8b176fe..d6a69d1 100644
--- a/src/gun_http.erl
+++ b/src/gun_http.erl
@@ -15,6 +15,7 @@
 -module(gun_http).
 
 -export([check_options/1]).
+-export([name/0]).
 -export([init/4]).
 -export([handle/2]).
 -export([close/1]).
@@ -23,6 +24,7 @@
 -export([request/8]).
 -export([data/4]).
 -export([cancel/2]).
+-export([down/1]).
 -export([ws_upgrade/7]).
 
 -type io() :: head | {body, non_neg_integer()} | body_close | body_chunked.
@@ -52,6 +54,8 @@ do_check_options([{version, V}|Opts]) when V =:= 'HTTP/1.1'; V =:= 'HTTP/1.0' ->
 do_check_options([Opt|_]) ->
 	{error, {options, {http, Opt}}}.
 
+name() -> http.
+
 init(Owner, Socket, Transport, Opts) ->
 	Version = maps:get(version, Opts, 'HTTP/1.1'),
 	#http_state{owner=Owner, socket=Socket, transport=Transport, version=Version}.
@@ -286,6 +290,14 @@ cancel(State, StreamRef) ->
 			error_stream_not_found(State)
 	end.
 
+%% HTTP does not provide any way to figure out what streams are unprocessed.
+down(#http_state{streams=Streams}) ->
+	KilledStreams = [case Ref of
+		{websocket, Ref2, _, _, _} -> Ref2;
+		_ -> Ref
+	end || {Ref, _} <- Streams],
+	{KilledStreams, []}.
+
 error_stream_closed(State=#http_state{owner=Owner}) ->
 	Owner ! {gun_error, self(), {badstate,
 		"The stream has already been closed."}},
diff --git a/src/gun_spdy.erl b/src/gun_spdy.erl
index 6aeb517..c417e78 100644
--- a/src/gun_spdy.erl
+++ b/src/gun_spdy.erl
@@ -15,6 +15,7 @@
 -module(gun_spdy).
 
 -export([check_options/1]).
+-export([name/0]).
 -export([init/4]).
 -export([handle/2]).
 -export([close/1]).
@@ -23,6 +24,7 @@
 -export([request/8]).
 -export([data/4]).
 -export([cancel/2]).
+-export([down/1]).
 
 -record(stream, {
 	id :: non_neg_integer(),
@@ -52,6 +54,8 @@ do_check_options([{keepalive, K}|Opts]) when is_integer(K), K > 0 ->
 do_check_options([Opt|_]) ->
 	{error, {options, {spdy, Opt}}}.
 
+name() -> spdy.
+
 init(Owner, Socket, Transport, _Opts) ->
 	#spdy_state{owner=Owner, socket=Socket, transport=Transport,
 		zdef=cow_spdy:deflate_init(), zinf=cow_spdy:inflate_init()}.
@@ -263,6 +267,11 @@ cancel(State=#spdy_state{socket=Socket, transport=Transport},
 			error_stream_not_found(State)
 	end.
 
+%% @todo Add unprocessed streams when GOAWAY handling is done.
+down(#spdy_state{streams=Streams}) ->
+	KilledStreams = [Ref || #stream{ref=Ref} <- Streams],
+	{KilledStreams, []}.
+
 error_stream_closed(State=#spdy_state{owner=Owner}) ->
 	Owner ! {gun_error, self(), {badstate,
 		"The stream has already been closed."}},
diff --git a/src/gun_ws.erl b/src/gun_ws.erl
index 5379362..246cc57 100644
--- a/src/gun_ws.erl
+++ b/src/gun_ws.erl
@@ -14,9 +14,11 @@
 
 -module(gun_ws).
 
+-export([name/0]).
 -export([init/5]).
 -export([handle/2]).
 -export([send/2]).
+-export([down/1]).
 
 -record(payload, {
 	type = undefined :: cow_ws:frame_type(),
@@ -40,6 +42,8 @@
 	extensions = #{} :: cow_ws:extensions()
 }).
 
+name() -> ws.
+
 %% @todo Protocols
 init(Owner, Socket, Transport, Extensions, _Protocols) ->
 	Owner ! {gun_ws_upgrade, self(), ok},
@@ -123,3 +127,7 @@ send(Frame, State=#ws_state{socket=Socket, transport=Transport, extensions=Exten
 		{close, _, _} -> close;
 		_ -> State
 	end.
+
+%% Websocket has no concept of streams.
+down(_) ->
+	{[], []}.