From 16726e24610bd1c8926c8338824b2a019e435609 Mon Sep 17 00:00:00 2001
From: Dan Sosedoff <dan.sosedoff@gmail.com>
Date: Fri, 2 Dec 2022 14:20:20 -0600
Subject: [PATCH] Add idle timeout into session manager

---
 pkg/api/session_manager.go      | 21 +++++++++++++++++----
 pkg/api/session_manager_test.go | 10 +++-------
 pkg/cli/cli.go                  |  2 ++
 pkg/client/client.go            | 10 +++++++++-
 4 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/pkg/api/session_manager.go b/pkg/api/session_manager.go
index 00478a8..9166f9f 100644
--- a/pkg/api/session_manager.go
+++ b/pkg/api/session_manager.go
@@ -10,9 +10,10 @@ import (
 )
 
 type SessionManager struct {
-	logger   *logrus.Logger
-	sessions map[string]*client.Client
-	mu       sync.Mutex
+	logger      *logrus.Logger
+	sessions    map[string]*client.Client
+	mu          sync.Mutex
+	idleTimeout time.Duration
 }
 
 func NewSessionManager(logger *logrus.Logger) *SessionManager {
@@ -23,6 +24,10 @@ func NewSessionManager(logger *logrus.Logger) *SessionManager {
 	}
 }
 
+func (m *SessionManager) SetIdleTimeout(timeout time.Duration) {
+	m.idleTimeout = timeout
+}
+
 func (m *SessionManager) IDs() []string {
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -79,6 +84,10 @@ func (m *SessionManager) Len() int {
 }
 
 func (m *SessionManager) Cleanup() int {
+	if m.idleTimeout == 0 {
+		return 0
+	}
+
 	removed := 0
 
 	m.logger.Debug("starting idle sessions cleanup")
@@ -97,6 +106,8 @@ func (m *SessionManager) Cleanup() int {
 }
 
 func (m *SessionManager) RunPeriodicCleanup() {
+	m.logger.WithField("timeout", m.idleTimeout).Info("session manager cleanup enabled")
+
 	for range time.Tick(time.Minute) {
 		m.Cleanup()
 	}
@@ -106,9 +117,11 @@ func (m *SessionManager) staleSessions() []string {
 	m.mu.TryLock()
 	defer m.mu.Unlock()
 
+	now := time.Now()
 	ids := []string{}
+
 	for id, conn := range m.sessions {
-		if conn.IsIdle() {
+		if now.Sub(conn.LastQueryTime()) > m.idleTimeout {
 			ids = append(ids, id)
 		}
 	}
diff --git a/pkg/api/session_manager_test.go b/pkg/api/session_manager_test.go
index 60c0bd5..6ac8c12 100644
--- a/pkg/api/session_manager_test.go
+++ b/pkg/api/session_manager_test.go
@@ -3,12 +3,12 @@ package api
 import (
 	"sort"
 	"testing"
+	"time"
 
 	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 
 	"github.com/sosedoff/pgweb/pkg/client"
-	"github.com/sosedoff/pgweb/pkg/command"
 )
 
 func TestSessionManager(t *testing.T) {
@@ -60,20 +60,16 @@ func TestSessionManager(t *testing.T) {
 	})
 
 	t.Run("clean up stale sessions", func(t *testing.T) {
-		defer func() {
-			command.Opts.ConnectionIdleTimeout = 0
-		}()
-
 		manager := NewSessionManager(logrus.New())
 		conn := &client.Client{}
 		manager.Add("foo", conn)
 
-		command.Opts.ConnectionIdleTimeout = 0
 		assert.Equal(t, 1, manager.Len())
 		assert.Equal(t, 0, manager.Cleanup())
 		assert.Equal(t, 1, manager.Len())
 
-		command.Opts.ConnectionIdleTimeout = 1
+		conn.Query("select 1")
+		manager.SetIdleTimeout(time.Minute)
 		assert.Equal(t, 1, manager.Cleanup())
 		assert.Equal(t, 0, manager.Len())
 		assert.True(t, conn.IsClosed())
diff --git a/pkg/cli/cli.go b/pkg/cli/cli.go
index 7bd07bf..bb3e306 100644
--- a/pkg/cli/cli.go
+++ b/pkg/cli/cli.go
@@ -8,6 +8,7 @@ import (
 	"regexp"
 	"strings"
 	"syscall"
+	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/jessevdk/go-flags"
@@ -266,6 +267,7 @@ func Run() {
 		api.DbSessions = api.NewSessionManager(logger)
 
 		if !command.Opts.DisableConnectionIdleTimeout {
+			api.DbSessions.SetIdleTimeout(time.Minute * time.Duration(command.Opts.ConnectionIdleTimeout))
 			go api.DbSessions.RunPeriodicCleanup()
 		}
 	}
diff --git a/pkg/client/client.go b/pkg/client/client.go
index f24272c..3d8bd50 100644
--- a/pkg/client/client.go
+++ b/pkg/client/client.go
@@ -335,6 +335,10 @@ func (client *Client) ServerVersion() string {
 }
 
 func (client *Client) query(query string, args ...interface{}) (*Result, error) {
+	if client.db == nil {
+		return nil, nil
+	}
+
 	// Update the last usage time
 	defer func() {
 		client.lastQueryTime = time.Now().UTC()
@@ -366,7 +370,7 @@ func (client *Client) query(query string, args ...interface{}) (*Result, error)
 		result := Result{
 			Columns: []string{"Rows Affected"},
 			Rows: []Row{
-				Row{affected},
+				{affected},
 			},
 		}
 
@@ -446,6 +450,10 @@ func (c *Client) IsClosed() bool {
 	return c.closed
 }
 
+func (c *Client) LastQueryTime() time.Time {
+	return c.lastQueryTime
+}
+
 func (client *Client) IsIdle() bool {
 	mins := int(time.Since(client.lastQueryTime).Minutes())