mod_c2s,mod_s2s: Wait for sessions to close before proceeding with shutdown steps

Ensures unavailable presence and other outgoing stanzas are sent.

Waiting for c2s sessions to close first before proceeding to disable and
close s2s ensures that unavailable presence can go out, even if it
requires dialback to complete first.

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);