Spawn TDlib instances for recovered sessions

5 files changed, 69 insertions, 33 deletions

diff --git a/persistence/sessions.go b/persistence/sessions.go
index 50e3f3e..36fac52 100644
--- a/persistence/sessions.go
+++ b/persistence/sessions.go
@@ -27,8 +27,8 @@ type Session struct {
 
 var sessionDB SessionsYamlDB
 
-// Marshaller implementation for YamlDB
-func Marshaller() ([]byte, error) {
+// SessionMarshaller implementation for YamlDB
+func SessionMarshaller() ([]byte, error) {
 	return yaml.Marshal(sessionDB.Data)
 }
 
@@ -41,9 +41,6 @@ func LoadSessions(path string) (SessionsYamlDB, error) {
 		return sessionDB, errors.Wrap(err, "Sessions restore error")
 	}
 
-	sessionDB.Transaction(func() {
-	}, Marshaller)
-
 	return sessionDB, nil
 }
 
diff --git a/telegram/client.go b/telegram/client.go
index 688912b..82f65d0 100644
--- a/telegram/client.go
+++ b/telegram/client.go
@@ -7,6 +7,7 @@ import (
 	"strconv"
 
 	"dev.narayana.im/narayana/telegabber/config"
+	"dev.narayana.im/narayana/telegabber/persistence"
 
 	"github.com/zelenin/go-tdlib/client"
 )
@@ -35,12 +36,13 @@ type Client struct {
 	client       *client.Client
 	jid          string
 	parameters   *client.TdlibParameters
+	session      *persistence.Session
 	online       bool
 	logVerbosity client.Option
 }
 
 // NewClient instantiates a Telegram App
-func NewClient(conf config.TelegramConfig, jid string) (Client, error) {
+func NewClient(conf config.TelegramConfig, jid string, session *persistence.Session) (Client, error) {
 	logVerbosity := client.WithLogVerbosity(&client.SetLogVerbosityLevelRequest{
 		NewVerbosityLevel: stringToLogConstant(conf.Loglevel),
 	})
@@ -76,6 +78,7 @@ func NewClient(conf config.TelegramConfig, jid string) (Client, error) {
 	return Client{
 		parameters:   &parameters,
 		jid:          jid,
+		session:      session,
 		logVerbosity: logVerbosity,
 	}, nil
 }
diff --git a/xmpp/component.go b/xmpp/component.go
index fcef91b..a341a2e 100644
--- a/xmpp/component.go
+++ b/xmpp/component.go
@@ -1,10 +1,13 @@
 package xmpp
 
 import (
+	"github.com/pkg/errors"
+
 	"dev.narayana.im/narayana/telegabber/config"
 	"dev.narayana.im/narayana/telegabber/persistence"
 	"dev.narayana.im/narayana/telegabber/telegram"
 
+	log "github.com/sirupsen/logrus"
 	"gosrc.io/xmpp"
 )
 
@@ -24,9 +27,8 @@ func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.Strea
 	}
 
 	tgConf = tc
-	sessions = make(map[string]telegram.Client)
 
-	db, err = persistence.LoadSessions(conf.Db)
+	err = loadSessions(conf.Db)
 	if err != nil {
 		return nil, err
 	}
@@ -50,5 +52,46 @@ func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.Strea
 
 	cm := xmpp.NewStreamManager(component, nil)
 
+	go maintenance()
+
 	return cm, nil
 }
+
+func maintenance() {
+	// TODO
+}
+
+func loadSessions(dbPath string) error {
+	var err error
+
+	sessions = make(map[string]telegram.Client)
+
+	db, err = persistence.LoadSessions(dbPath)
+	if err != nil {
+		return err
+	}
+
+	db.Transaction(func() bool {
+		for jid, session := range db.Data.Sessions {
+			getTelegramInstance(jid, &session)
+		}
+
+		return false
+	}, persistence.SessionMarshaller)
+
+	return nil
+}
+
+func getTelegramInstance(jid string, savedSession *persistence.Session) (telegram.Client, bool) {
+	session, ok := sessions[jid]
+	if !ok {
+		session, err := telegram.NewClient(tgConf, jid, savedSession)
+		if err != nil {
+			log.Error(errors.Wrap(err, "TDlib initialization failure"))
+			return session, false
+		}
+		sessions[jid] = session
+	}
+
+	return session, true
+}
diff --git a/xmpp/handlers.go b/xmpp/handlers.go
index ad6e3ba..c19a4ef 100644
--- a/xmpp/handlers.go
+++ b/xmpp/handlers.go
@@ -3,8 +3,6 @@ package xmpp
 import (
 	"github.com/pkg/errors"
 
-	"dev.narayana.im/narayana/telegabber/telegram"
-
 	log "github.com/sirupsen/logrus"
 	"gosrc.io/xmpp"
 	"gosrc.io/xmpp/stanza"
@@ -90,14 +88,9 @@ func handlePresence(s xmpp.Sender, p stanza.Presence) {
 		return
 	}
 	bareFromJid := fromJid.Bare()
-	session, ok := sessions[bareFromJid]
+	session, ok := getTelegramInstance(bareFromJid, nil)
 	if !ok {
-		client, err := telegram.NewClient(tgConf, bareFromJid)
-		if err != nil {
-			log.Error(errors.Wrap(err, "TDlib initialization failure"))
-			return
-		}
-		sessions[bareFromJid] = client
+		return
 	}
 
 	switch p.Type {
diff --git a/yamldb/yamldb.go b/yamldb/yamldb.go
index eec23be..1478cc7 100644
--- a/yamldb/yamldb.go
+++ b/yamldb/yamldb.go
@@ -18,9 +18,7 @@ type YamlDB struct {
 
 // Transaction executes the given callback and safely saves
 // the data after they are modified within the callback
-func (db *YamlDB) Transaction(callback func(), marshaller func() ([]byte, error)) error {
-	var err error
-
+func (db *YamlDB) Transaction(callback func() bool, marshaller func() ([]byte, error)) error {
 	log.Debug("Enter transaction")
 	db.lock.Lock()
 	defer func() {
@@ -28,19 +26,21 @@ func (db *YamlDB) Transaction(callback func(), marshaller func() ([]byte, error)
 		log.Debug("Exit transaction")
 	}()
 
-	callback()
-
-	yamlData, err := marshaller()
-	if err != nil {
-		return errors.Wrap(err, "Data marshalling error")
-	}
-	err = ioutil.WriteFile(db.PathNew, yamlData, 0644)
-	if err != nil {
-		return errors.Wrap(err, "YamlDB write failure")
-	}
-	err = os.Rename(db.PathNew, db.Path)
-	if err != nil {
-		return errors.Wrap(err, "Couldn't rewrite an old YamlDB file")
+	isDataChanged := callback()
+
+	if isDataChanged {
+		yamlData, err := marshaller()
+		if err != nil {
+			return errors.Wrap(err, "Data marshalling error")
+		}
+		err = ioutil.WriteFile(db.PathNew, yamlData, 0644)
+		if err != nil {
+			return errors.Wrap(err, "YamlDB write failure")
+		}
+		err = os.Rename(db.PathNew, db.Path)
+		if err != nil {
+			return errors.Wrap(err, "Couldn't rewrite an old YamlDB file")
+		}
 	}
 
 	return nil