Save sessions on exit

1 files changed, 25 insertions, 6 deletions

diff --git a/xmpp/component.go b/xmpp/component.go
index d82f4bb..fa1a650 100644
--- a/xmpp/component.go
+++ b/xmpp/component.go
@@ -25,19 +25,19 @@ var db persistence.SessionsYamlDB
 
 // NewComponent starts a new component and wraps it in
 // a stream manager that you should start yourself
-func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.StreamManager, error) {
+func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.StreamManager, *xmpp.Component, error) {
 	var err error
 
 	jid, err = xmpp.NewJid(conf.Jid)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	tgConf = tc
 
 	err = loadSessions(conf.Db)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	options := xmpp.ComponentOptions{
@@ -54,14 +54,14 @@ func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.Strea
 
 	component, err := xmpp.NewComponent(options, router)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	cm := xmpp.NewStreamManager(component, nil)
+	sm := xmpp.NewStreamManager(component, nil)
 
 	go heartbeat(component)
 
-	return cm, nil
+	return sm, component, nil
 }
 
 func logPresence(err error, presence *stanza.Presence) {
@@ -228,3 +228,22 @@ func sendPresence(component *xmpp.Component, to string, args ...args.V) error {
 
 	return nil
 }
+
+// Close gracefully terminates the component and saves active sessions
+func Close(component *xmpp.Component) {
+	log.Error("Disconnecting...")
+
+	for _, session := range sessions {
+		session.Disconnect()
+	}
+
+	db.Transaction(func() bool {
+		for jid, session := range sessions {
+			db.Data.Sessions[jid] = *session.Session
+		}
+
+		return true
+	}, persistence.SessionMarshaller)
+
+	component.Disconnect()
+}