Save sessions on exit

2 files changed, 28 insertions, 7 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()
+}
diff --git a/xmpp/handlers.go b/xmpp/handlers.go
index 3709443..70f985d 100644
--- a/xmpp/handlers.go
+++ b/xmpp/handlers.go
@@ -3,6 +3,8 @@ package xmpp
 import (
 	"github.com/pkg/errors"
 
+	"dev.narayana.im/narayana/telegabber/persistence"
+
 	log "github.com/sirupsen/logrus"
 	"gosrc.io/xmpp"
 	"gosrc.io/xmpp/stanza"
@@ -88,7 +90,7 @@ func handlePresence(s xmpp.Sender, p stanza.Presence) {
 		return
 	}
 	bareFromJid := fromJid.Bare()
-	session, ok := getTelegramInstance(bareFromJid, nil)
+	session, ok := getTelegramInstance(bareFromJid, &persistence.Session{})
 	if !ok {
 		return
 	}