Merge branch '107-restart-user-session' into 'master'

feat: dump user's session data in json file during shutdown to restore active sessions

Closes #107

See merge request postgres-ai/joe!158

Author: agneum

.golangci.yml

@@ -1,5 +1,5 @@
 run:
-  timeout: 1m
+  timeout: 90s
   issues-exit-code: 1
   tests: true
   skip-dirs:
@@ -48,6 +48,8 @@ linters-settings:
       - hugeParam
     enabled-tags:
       - performance
+  godot:
+    scope: declarations
 
 linters:
   enable:
@@ -76,6 +78,7 @@ linters:
     - unused
     - unparam
     - wsl
+    - godot
   enable-all: false
   disable:
     - gosec

cmd/joe/main.go

@@ -25,10 +25,14 @@ import (
 	"gitlab.com/postgres-ai/joe/pkg/bot"
 	"gitlab.com/postgres-ai/joe/pkg/config"
 	"gitlab.com/postgres-ai/joe/pkg/services/platform"
+	"gitlab.com/postgres-ai/joe/pkg/services/storage"
 )
 
 const (
 	shutdownTimeout = 60 * time.Second
+
+	configFilePath   = "config/config.yml"
+	sessionsFilePath = "config/sessions.json"
 )
 
 // ldflag variables.
@@ -37,7 +41,7 @@ var buildTime, version string
 func main() {
 	version := formatBotVersion()
 
-	botCfg, err := loadConfig("config/config.yml")
+	botCfg, err := loadConfig(configFilePath)
 	if err != nil {
 		log.Fatal("failed to load config: %v", err)
 	}
@@ -56,7 +60,14 @@ func main() {
 	ctx, cancel := context.WithCancel(context.Background())
 	shutdownCh := setShutdownListener()
 
-	joeBot := bot.NewApp(botCfg, platformClient, features.NewPack())
+	sessionsStorage := storage.NewJSONSessionData(sessionsFilePath)
+	if err := sessionsStorage.Load(); err != nil {
+		log.Fatal("unable to load sessions data: ", err)
+	}
+
+	joeBot := bot.NewApp(botCfg, platformClient, features.NewPack(), sessionsStorage)
+
+	go setSighupListener(ctx, joeBot)
 
 	go func() {
 		if err := joeBot.RunServer(ctx); err != nil && err != http.ErrServerClosed {
@@ -65,6 +76,7 @@ func main() {
 	}()
 
 	<-shutdownCh
+	log.Dbg("shutdown request received")
 	cancel()
 
 	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout)
@@ -111,3 +123,20 @@ func setShutdownListener() chan os.Signal {
 
 	return c
 }
+
+// setSighupListener allows to dump active sessions.
+func setSighupListener(ctx context.Context, app *bot.App) {
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, syscall.SIGHUP)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-c:
+			if err := app.SaveSessions(); err != nil {
+				log.Err("failed to save user session data: ", err)
+			}
+		}
+	}
+}

pkg/bot/bot.go

@@ -27,6 +27,7 @@ import (
 	"gitlab.com/postgres-ai/joe/pkg/connection/slackrtm"
 	"gitlab.com/postgres-ai/joe/pkg/connection/webui"
 	"gitlab.com/postgres-ai/joe/pkg/services/dblab"
+	"gitlab.com/postgres-ai/joe/pkg/services/storage"
 	"gitlab.com/postgres-ai/joe/pkg/util"
 )
 
@@ -43,6 +44,8 @@ type App struct {
 
 	dblabMu        *sync.RWMutex
 	dblabInstances map[string]*dblab.Instance
+
+	sessionStorage storage.PersistentSessionStorage
 }
 
 // HealthResponse represents a response for heath-check requests.
@@ -53,13 +56,15 @@ type HealthResponse struct {
 }
 
 // Creates a new application.
-func NewApp(cfg *config.Config, platformClient *platform.Client, enterprise *features.Pack) *App {
+func NewApp(cfg *config.Config, platformClient *platform.Client,
+	enterprise *features.Pack, sessions storage.PersistentSessionStorage) *App {
 	bot := App{
 		Config:         cfg,
 		dblabMu:        &sync.RWMutex{},
 		dblabInstances: make(map[string]*dblab.Instance, len(cfg.ChannelMapping.DBLabInstances)),
 		featurePack:    enterprise,
 		platformClient: platformClient,
+		sessionStorage: sessions,
 	}
 
 	return &bot
@@ -95,12 +100,25 @@ func (a *App) Shutdown(ctx context.Context) error {
 	}
 
 	if len(a.assistants) > 0 {
+		if err := a.SaveSessions(); err != nil {
+			log.Err("unable to dump sessionStorage data: ", err)
+		}
+
 		a.deregisterAssistants(ctx)
 	}
 
 	return nil
 }
 
+// SaveSessions dumps session data to the disk.
+func (a *App) SaveSessions() error {
+	for _, assistantSvc := range a.assistants {
+		assistantSvc.DumpSessions()
+	}
+
+	return a.sessionStorage.Save()
+}
+
 func (a *App) deregisterAssistants(ctx context.Context) {
 	wg := sync.WaitGroup{}
 
@@ -185,13 +203,13 @@ func (a *App) getAssistant(communicationType string, workspaceCfg config.Workspa
 
 	switch communicationType {
 	case slack.CommunicationType:
-		return slack.NewAssistant(&workspaceCfg.Credentials, a.Config, handlerPrefix, a.featurePack, a.platformClient), nil
+		return slack.NewAssistant(&workspaceCfg.Credentials, a.Config, handlerPrefix, a.featurePack, a.platformClient, a.sessionStorage), nil
 
 	case slackrtm.CommunicationType:
-		return slackrtm.NewAssistant(&workspaceCfg.Credentials, a.Config, a.featurePack, a.platformClient), nil
+		return slackrtm.NewAssistant(&workspaceCfg.Credentials, a.Config, a.featurePack, a.platformClient, a.sessionStorage), nil
 
 	case webui.CommunicationType:
-		return webui.NewAssistant(&workspaceCfg.Credentials, a.Config, handlerPrefix, a.featurePack, a.platformClient), nil
+		return webui.NewAssistant(&workspaceCfg.Credentials, a.Config, handlerPrefix, a.featurePack, a.platformClient, a.sessionStorage), nil
 
 	default:
 		return nil, errors.New("unknown workspace type given")
@@ -222,6 +240,10 @@ func (a *App) setupChannels(ctx context.Context, assistant connection.Assistant,
 			return errors.Wrapf(err, "failed to register the %q assistant", channel.ChannelID)
 		}
 
+		if err := assistant.RestoreSessions(ctx); err != nil {
+			return errors.Wrapf(err, "failed to restore active sessions for the %q assistant", channel.ChannelID)
+		}
+
 		_ = util.RunInterval(InactiveCloneCheckInterval, func() {
 			assistant.CheckIdleSessions(ctx)
 		})

pkg/connection/assistant.go

@@ -9,6 +9,7 @@ import (
 
 	"gitlab.com/postgres-ai/joe/pkg/models"
 	"gitlab.com/postgres-ai/joe/pkg/services/dblab"
+	"gitlab.com/postgres-ai/joe/pkg/services/usermanager"
 )
 
 // Assistant defines the interface of a Query Optimization assistant.
@@ -22,11 +23,17 @@ type Assistant interface {
 	// Deregister defines the method to deregister the assistant.
 	Deregister(ctx context.Context) error
 
+	// RestoreSessions checks sessions after restart and establishes DB connection
+	RestoreSessions(context.Context) error
+
 	// CheckIdleSessions defines the method for checking user idle sessions and notification about them.
 	CheckIdleSessions(context.Context)
 
 	// AddChannel adds a new Database Lab instance to communication via the assistant.
 	AddChannel(channelID, project string, dbLabInstance *dblab.Instance)
+
+	// DumpSessions iterates over channels and collects user's sessions to storage
+	DumpSessions()
 }
 
 // MessageProcessor defines the interface of a message processor.
@@ -37,6 +44,12 @@ type MessageProcessor interface {
 	// ProcessAppMentionEvent defines the method for replying to an application mention event.
 	ProcessAppMentionEvent(incomingMessage models.IncomingMessage)
 
+	// RestoreSessions checks sessions after restart and establishes DB connection
+	RestoreSessions(ctx context.Context) error
+
+	// Users returns user's session data
+	Users() usermanager.UserList
+
 	// CheckIdleSessions defines the method of check idleness sessions.
 	CheckIdleSessions(ctx context.Context)
 }

pkg/connection/slack/assistant.go

@@ -30,6 +30,7 @@ import (
 	"gitlab.com/postgres-ai/joe/pkg/services/dblab"
 	"gitlab.com/postgres-ai/joe/pkg/services/msgproc"
 	"gitlab.com/postgres-ai/joe/pkg/services/platform"
+	"gitlab.com/postgres-ai/joe/pkg/services/storage"
 	"gitlab.com/postgres-ai/joe/pkg/services/usermanager"
 )
 
@@ -50,19 +51,19 @@ type Assistant struct {
 	appCfg         *config.Config
 	featurePack    *features.Pack
 	messenger      *Messenger
-	userManager    *usermanager.UserManager
+	userInformer   usermanager.UserInformer
 	platformClient *platform.Client
+	sessionStorage storage.SessionStorage
 }
 
 // NewAssistant returns a new assistant service.
 func NewAssistant(cfg *config.Credentials, appCfg *config.Config, handlerPrefix string, pack *features.Pack,
-	platformClient *platform.Client) *Assistant {
+	platformClient *platform.Client, sessionStorage storage.SessionStorage) *Assistant {
 	prefix := fmt.Sprintf("/%s", strings.Trim(handlerPrefix, "/"))
 
 	chatAPI := slack.New(cfg.AccessToken)
 	messenger := NewMessenger(chatAPI, &MessengerConfig{AccessToken: cfg.AccessToken})
 	userInformer := NewUserInformer(chatAPI)
-	userManager := usermanager.NewUserManager(userInformer, appCfg.Enterprise.Quota)
 
 	assistant := &Assistant{
 		credentialsCfg: cfg,
@@ -71,8 +72,9 @@ func NewAssistant(cfg *config.Credentials, appCfg *config.Config, handlerPrefix
 		prefix:         prefix,
 		featurePack:    pack,
 		messenger:      messenger,
-		userManager:    userManager,
+		userInformer:   userInformer,
 		platformClient: platformClient,
+		sessionStorage: sessionStorage,
 	}
 
 	return assistant
@@ -117,12 +119,12 @@ func (a *Assistant) Deregister(_ context.Context) error {
 
 // AddChannel sets a message processor for a specific channel.
 func (a *Assistant) AddChannel(channelID, project string, dbLabInstance *dblab.Instance) {
-	messageProcessor := a.buildMessageProcessor(project, dbLabInstance)
+	messageProcessor := a.buildMessageProcessor(channelID, project, dbLabInstance)
 
 	a.addProcessingService(channelID, messageProcessor)
 }
 
-func (a *Assistant) buildMessageProcessor(project string, dbLabInstance *dblab.Instance) *msgproc.ProcessingService {
+func (a *Assistant) buildMessageProcessor(channelID, project string, dbLabInstance *dblab.Instance) *msgproc.ProcessingService {
 	processingCfg := msgproc.ProcessingConfig{
 		App:      a.appCfg.App,
 		Platform: a.appCfg.Platform,
@@ -132,7 +134,10 @@ func (a *Assistant) buildMessageProcessor(project string, dbLabInstance *dblab.I
 		Project:  project,
 	}
 
-	return msgproc.NewProcessingService(a.messenger, MessageValidator{}, dbLabInstance.Client(), a.userManager, a.platformClient,
+	users := a.sessionStorage.GetUsers(CommunicationType, channelID)
+	um := usermanager.NewUserManager(a.userInformer, a.appCfg.Enterprise.Quota, users)
+
+	return msgproc.NewProcessingService(a.messenger, MessageValidator{}, dbLabInstance.Client(), um, a.platformClient,
 		processingCfg, a.featurePack)
 }
 
@@ -156,6 +161,22 @@ func (a *Assistant) getProcessingService(channelID string) (connection.MessagePr
 	return messageProcessor, nil
 }
 
+// RestoreSessions checks sessions after restart and establishes DB connection.
+func (a *Assistant) RestoreSessions(ctx context.Context) error {
+	log.Dbg("Restore sessions", a.prefix)
+
+	a.procMu.RLock()
+	defer a.procMu.RUnlock()
+
+	for _, proc := range a.msgProcessors {
+		if err := proc.RestoreSessions(ctx); err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
 // CheckIdleSessions check the running user sessions for idleness.
 func (a *Assistant) CheckIdleSessions(ctx context.Context) {
 	log.Dbg("Check idle sessions", a.prefix)
@@ -358,3 +379,15 @@ func (a *Assistant) verifyRequest(r *http.Request) error {
 
 	return nil
 }
+
+// DumpSessions collects user's data from every message processor to sessionStorage.
+func (a *Assistant) DumpSessions() {
+	log.Dbg("dump sessions", a.prefix)
+
+	a.procMu.RLock()
+	defer a.procMu.RUnlock()
+
+	for channelID, proc := range a.msgProcessors {
+		a.sessionStorage.SetUsers(CommunicationType, channelID, proc.Users())
+	}
+}