2 files changed, 32 insertions, 2 deletions

diff --git a/plugins/mod_c2s.lua b/plugins/mod_c2s.lua
index 1a9580fb..05c37824 100644
--- a/plugins/mod_c2s.lua
+++ b/plugins/mod_c2s.lua
@@ -16,7 +16,8 @@ local statsmanager = require "core.statsmanager";
 local st = require "util.stanza";
 local sm_new_session, sm_destroy_session = sessionmanager.new_session, sessionmanager.destroy_session;
 local uuid_generate = require "util.uuid".generate;
-local runner = require "util.async".runner;
+local async = require "util.async";
+local runner = async.runner;
 
 local tostring, type = tostring, type;
 
@@ -382,6 +383,7 @@ function listener.ondisconnect(conn, err)
 		session.conn = nil;
 		sessions[conn]  = nil;
 	end
+	module:fire_event("c2s-closed", { session = session; conn = conn });
 end
 
 function listener.onreadtimeout(conn)
@@ -431,11 +433,24 @@ module:hook("server-stopping", function(event)
 end, -80);
 
 module:hook("server-stopping", function(event)
+	local wait, done = async.waiter();
+	module:hook("c2s-closed", function ()
+		if next(sessions) == nil then done(); end
+	end)
+
 	-- Close sessions
 	local reason = event.reason;
 	for _, session in pairs(sessions) do
 		session:close{ condition = "system-shutdown", text = reason };
 	end
+
+	-- Wait for them to close properly if they haven't already
+	if next(sessions) ~= nil then
+		add_task(stream_close_timeout+1, done);
+		module:log("info", "Waiting for sessions to close");
+		wait();
+	end
+
 end, -100);
 
 
diff --git a/plugins/mod_s2s.lua b/plugins/mod_s2s.lua
index 4b721237..836cf347 100644
--- a/plugins/mod_s2s.lua
+++ b/plugins/mod_s2s.lua
@@ -26,7 +26,8 @@ local s2s_new_incoming = require "core.s2smanager".new_incoming;
 local s2s_new_outgoing = require "core.s2smanager".new_outgoing;
 local s2s_destroy_session = require "core.s2smanager".destroy_session;
 local uuid_gen = require "util.uuid".generate;
-local runner = require "util.async".runner;
+local async = require "util.async";
+local runner = async.runner;
 local connect = require "net.connect".connect;
 local service = require "net.resolvers.service";
 local resolver_chain = require "net.resolvers.chain";
@@ -859,6 +860,7 @@ function listener.ondisconnect(conn, err)
 		end
 		s2s_destroy_session(session, err);
 	end
+	module:fire_event("s2s-closed", { session = session; conn = conn });
 end
 
 function listener.onfail(data, err)
@@ -971,11 +973,24 @@ module:hook("server-stopping", function(event)
 		end
 	end
 
+	local wait, done = async.waiter();
+	module:hook("s2s-closed", function ()
+		if next(sessions) == nil then done(); end
+	end, 1)
+
 	-- Close sessions
 	local reason = event.reason;
 	for _, session in pairs(sessions) do
 		session:close{ condition = "system-shutdown", text = reason };
 	end
+
+	-- Wait for them to close properly if they haven't already
+	if next(sessions) ~= nil then
+		module:log("info", "Waiting for sessions to close");
+		add_task(stream_close_timeout + 1, done);
+		wait();
+	end
+
 end, -200);