From b61f535134adb35df35c16b6eebe51b3d14aaf48 Mon Sep 17 00:00:00 2001
From: James Fish <james@fishcakez.com>
Date: Sun, 17 Feb 2013 00:17:43 +0000
Subject: Fix to prevent loop handler awakening immediately after response sent

If a loop handler sent a response (e.g. cowboy_req:chunked_reply/2,/3)
and then returns {loop, Req, HandlerState, hibernate} it
would have a {cowboy_req, resp_sent} message in its message queue. This
message would cause the process to immediately awaken, so it is flushed
before hibernation.
---
 src/cowboy_handler.erl | 32 +++++++++++++++++++++++---------
 1 file changed, 23 insertions(+), 9 deletions(-)

diff --git a/src/cowboy_handler.erl b/src/cowboy_handler.erl
index ab09a97..7aaf9ae 100644
--- a/src/cowboy_handler.erl
+++ b/src/cowboy_handler.erl
@@ -70,17 +70,17 @@ handler_init(Req, State, Handler, HandlerOpts) ->
 		{ok, Req2, HandlerState} ->
 			handler_handle(Req2, State, Handler, HandlerState);
 		{loop, Req2, HandlerState} ->
-			handler_before_loop(Req2, State, Handler, HandlerState);
+			handler_after_callback(Req2, State, Handler, HandlerState);
 		{loop, Req2, HandlerState, hibernate} ->
-			handler_before_loop(Req2, State#state{hibernate=true},
+			handler_after_callback(Req2, State#state{hibernate=true},
 				Handler, HandlerState);
 		{loop, Req2, HandlerState, Timeout} ->
 			State2 = handler_loop_timeout(State#state{loop_timeout=Timeout}),
-			handler_before_loop(Req2, State2, Handler, HandlerState);
+			handler_after_callback(Req2, State2, Handler, HandlerState);
 		{loop, Req2, HandlerState, Timeout, hibernate} ->
 			State2 = handler_loop_timeout(State#state{
 				hibernate=true, loop_timeout=Timeout}),
-			handler_before_loop(Req2, State2, Handler, HandlerState);
+			handler_after_callback(Req2, State2, Handler, HandlerState);
 		{shutdown, Req2, HandlerState} ->
 			terminate_request(Req2, State, Handler, HandlerState,
 				{normal, shutdown});
@@ -133,6 +133,23 @@ handler_handle(Req, State, Handler, HandlerState) ->
 		error_terminate(Req, State)
 	end.
 
+%% Update the state if the response was sent in the callback.
+-spec handler_after_callback(Req, #state{}, module(), any())
+	-> {ok, Req, cowboy_middleware:env()}
+	| {error, 500, Req} | {suspend, module(), atom(), [any()]}
+	when Req::cowboy_req:req().
+handler_after_callback(Req, State=#state{resp_sent=false}, Handler,
+		HandlerState) ->
+	receive
+		{cowboy_req, resp_sent} ->
+			handler_before_loop(Req, State#state{resp_sent=true}, Handler,
+				HandlerState)
+	after 0 ->
+		handler_before_loop(Req, State, Handler, HandlerState)
+	end;
+handler_after_callback(Req, State, Handler, HandlerState) ->
+	handler_before_loop(Req, State, Handler, HandlerState).
+
 %% We don't listen for Transport closes because that would force us
 %% to receive data and buffer it indefinitely.
 -spec handler_before_loop(Req, #state{}, module(), any())
@@ -191,9 +208,6 @@ handler_loop(Req, State=#state{loop_buffer_size=NbBytes,
 		{Error, Socket, Reason} ->
 			terminate_request(Req, State, Handler, HandlerState,
 				{error, Reason});
-		{cowboy_req, resp_sent} ->
-			handler_before_loop(Req, State#state{resp_sent=true},
-				Handler, HandlerState);
 		{timeout, TRef, ?MODULE} ->
 			handler_after_loop(Req, State, Handler, HandlerState,
 				{normal, timeout});
@@ -213,9 +227,9 @@ handler_call(Req, State, Handler, HandlerState, Message) ->
 			handler_after_loop(Req2, State, Handler, HandlerState2,
 				{normal, shutdown});
 		{loop, Req2, HandlerState2} ->
-			handler_before_loop(Req2, State, Handler, HandlerState2);
+			handler_after_callback(Req2, State, Handler, HandlerState2);
 		{loop, Req2, HandlerState2, hibernate} ->
-			handler_before_loop(Req2, State#state{hibernate=true},
+			handler_after_callback(Req2, State#state{hibernate=true},
 				Handler, HandlerState2)
 	catch Class:Reason ->
 		error_logger:error_msg(
-- 
cgit v1.2.3