From f0c0d0ba94c923ca06b49211c70db47616e5a8e1 Mon Sep 17 00:00:00 2001 From: bodqhrohro Date: Tue, 19 Nov 2019 22:25:14 +0200 Subject: [PATCH] Save sessions on exit --- persistence/sessions.go | 12 +++++++++++- telegabber.go | 38 ++++++++++++++++++++++++++++++++++++-- telegram/client.go | 4 ++-- xmpp/component.go | 31 +++++++++++++++++++++++++------ xmpp/handlers.go | 4 +++- yamldb/yamldb.go | 3 +++ 6 files changed, 80 insertions(+), 12 deletions(-) diff --git a/persistence/sessions.go b/persistence/sessions.go index 36fac52..ab46823 100644 --- a/persistence/sessions.go +++ b/persistence/sessions.go @@ -6,6 +6,7 @@ import ( "dev.narayana.im/narayana/telegabber/yamldb" + log "github.com/sirupsen/logrus" "gopkg.in/yaml.v2" ) @@ -44,6 +45,10 @@ func LoadSessions(path string) (SessionsYamlDB, error) { return sessionDB, nil } +func emptySessionsMap(dataPtr *SessionsMap) { + dataPtr.Sessions = make(map[string]Session) +} + func initYamlDB(path string, dataPtr *SessionsMap) (SessionsYamlDB, error) { file, err := ioutil.ReadFile(path) if err == nil { @@ -51,9 +56,14 @@ func initYamlDB(path string, dataPtr *SessionsMap) (SessionsYamlDB, error) { if err != nil { return SessionsYamlDB{}, errors.Wrap(err, "YamlDB is corrupted") } + + if dataPtr.Sessions == nil { + emptySessionsMap(dataPtr) + } + log.Debugf("Unmarshalled YAML: %#v", *dataPtr) } else { // DB file does not exist, create an empty DB - dataPtr.Sessions = make(map[string]Session) + emptySessionsMap(dataPtr) } return SessionsYamlDB{ diff --git a/telegabber.go b/telegabber.go index 8c098e1..a196cc1 100644 --- a/telegabber.go +++ b/telegabber.go @@ -1,10 +1,14 @@ package main import ( + "os" + "os/signal" + "dev.narayana.im/narayana/telegabber/config" "dev.narayana.im/narayana/telegabber/xmpp" log "github.com/sirupsen/logrus" + goxmpp "gosrc.io/xmpp" ) // YAML config, compatible with the format of Zhabogram 2.0.0 @@ -13,7 +17,18 @@ const configPath string = "config.yml" // JSON schema (not for editing by a user) const schemaPath string = "./config_schema.json" +var sm *goxmpp.StreamManager +var component *goxmpp.Component +var err error + +var cleanupDone chan struct{} +var sigintChannel chan os.Signal + func main() { + cleanupDone = make(chan struct{}) + sigintChannel = make(chan os.Signal, 1) + signal.Notify(sigintChannel, os.Interrupt) + config, err := config.ReadConfig(configPath, schemaPath) if err != nil { log.Fatal(err) @@ -21,11 +36,30 @@ func main() { SetLogrusLevel(config.XMPP.Loglevel) - cm, err := xmpp.NewComponent(config.XMPP, config.Telegram) + sm, component, err = xmpp.NewComponent(config.XMPP, config.Telegram) if err != nil { log.Fatal(err) } + go func() { + <-sigintChannel + log.Error("Interrupting...") + exit() + + os.Exit(0) + }() + // reconnect automatically - log.Fatal(cm.Run()) + err = sm.Run() + exit() + + if err != nil { + <-cleanupDone + log.Fatal(err) + } +} + +func exit() { + xmpp.Close(component) + close(cleanupDone) } diff --git a/telegram/client.go b/telegram/client.go index 82f65d0..833d883 100644 --- a/telegram/client.go +++ b/telegram/client.go @@ -36,7 +36,7 @@ type Client struct { client *client.Client jid string parameters *client.TdlibParameters - session *persistence.Session + Session *persistence.Session online bool logVerbosity client.Option } @@ -78,7 +78,7 @@ func NewClient(conf config.TelegramConfig, jid string, session *persistence.Sess return Client{ parameters: ¶meters, jid: jid, - session: session, + Session: session, logVerbosity: logVerbosity, }, nil } 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 } diff --git a/yamldb/yamldb.go b/yamldb/yamldb.go index 1478cc7..76a8292 100644 --- a/yamldb/yamldb.go +++ b/yamldb/yamldb.go @@ -30,13 +30,16 @@ func (db *YamlDB) Transaction(callback func() bool, marshaller func() ([]byte, e if isDataChanged { yamlData, err := marshaller() + log.Debugf("Marshalled YAML: %#v", string(yamlData)) 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")