From 354a4acd19d5376477275588d3e170df6a688697 Mon Sep 17 00:00:00 2001
From: bodqhrohro <bodqhrohro@gmail.com>
Date: Wed, 4 Dec 2019 17:55:15 +0200
Subject: [PATCH] Fix session restoring

---
 persistence/sessions.go | 13 +++++++------
 xmpp/component.go       | 19 +++++++------------
 2 files changed, 14 insertions(+), 18 deletions(-)

diff --git a/persistence/sessions.go b/persistence/sessions.go
index ab46823..d15c4f3 100644
--- a/persistence/sessions.go
+++ b/persistence/sessions.go
@@ -26,7 +26,7 @@ type Session struct {
 	Login string `yaml:":login"`
 }
 
-var sessionDB SessionsYamlDB
+var sessionDB *SessionsYamlDB
 
 // SessionMarshaller implementation for YamlDB
 func SessionMarshaller() ([]byte, error) {
@@ -34,10 +34,11 @@ func SessionMarshaller() ([]byte, error) {
 }
 
 // LoadSessions restores TDlib sessions from the previous run
-func LoadSessions(path string) (SessionsYamlDB, error) {
+func LoadSessions(path string) (*SessionsYamlDB, error) {
 	var sessionData SessionsMap
+	var err error
 
-	sessionDB, err := initYamlDB(path, &sessionData)
+	sessionDB, err = initYamlDB(path, &sessionData)
 	if err != nil {
 		return sessionDB, errors.Wrap(err, "Sessions restore error")
 	}
@@ -49,12 +50,12 @@ func emptySessionsMap(dataPtr *SessionsMap) {
 	dataPtr.Sessions = make(map[string]Session)
 }
 
-func initYamlDB(path string, dataPtr *SessionsMap) (SessionsYamlDB, error) {
+func initYamlDB(path string, dataPtr *SessionsMap) (*SessionsYamlDB, error) {
 	file, err := ioutil.ReadFile(path)
 	if err == nil {
 		err = yaml.Unmarshal(file, dataPtr)
 		if err != nil {
-			return SessionsYamlDB{}, errors.Wrap(err, "YamlDB is corrupted")
+			return nil, errors.Wrap(err, "YamlDB is corrupted")
 		}
 
 		if dataPtr.Sessions == nil {
@@ -66,7 +67,7 @@ func initYamlDB(path string, dataPtr *SessionsMap) (SessionsYamlDB, error) {
 		emptySessionsMap(dataPtr)
 	}
 
-	return SessionsYamlDB{
+	return &SessionsYamlDB{
 		YamlDB: yamldb.YamlDB{
 			Path:    path,
 			PathNew: path + ".new",
diff --git a/xmpp/component.go b/xmpp/component.go
index 57d9c7b..bc81126 100644
--- a/xmpp/component.go
+++ b/xmpp/component.go
@@ -13,11 +13,9 @@ import (
 	"gosrc.io/xmpp"
 )
 
-const pollingInterval time.Duration = 1e7
-
 var tgConf config.TelegramConfig
 var sessions map[string]*telegram.Client
-var db persistence.SessionsYamlDB
+var db *persistence.SessionsYamlDB
 
 // NewComponent starts a new component and wraps it in
 // a stream manager that you should start yourself
@@ -53,9 +51,9 @@ func NewComponent(conf config.XMPPConfig, tc config.TelegramConfig) (*xmpp.Strea
 		return nil, nil, err
 	}
 
-	sm := xmpp.NewStreamManager(component, nil)
-
-	go heartbeat(component)
+	sm := xmpp.NewStreamManager(component, func(s xmpp.Sender) {
+		go heartbeat(component)
+	})
 
 	return sm, component, nil
 }
@@ -65,12 +63,9 @@ func heartbeat(component *xmpp.Component) {
 	probeType := gateway.SPType("probe")
 
 	for jid := range sessions {
-		for {
-			err = gateway.SendPresence(component, jid, probeType)
-			if err == nil {
-				break
-			}
-			time.Sleep(pollingInterval)
+		err = gateway.SendPresence(component, jid, probeType)
+		if err != nil {
+			log.Error(err)
 		}
 	}