fankes/beszel
2
0
Fork 0
mirror of https://github.com/fankes/beszel.git synced 2025-10-20 02:09:28 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

use sync.Map for system connections

Browse Source
This commit is contained in:
Henry Dollman
2024-11-15 15:35:26 -05:00
parent 4e64d9efad
commit e6839480d9
1 changed files with 15 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
beszel/internal/hub/hub.go
View File

@@ -35,8 +35,7 @@ import (
type Hub struct { type Hub struct {
app *pocketbase.PocketBase app *pocketbase.PocketBase
connectionLock *sync.Mutex systemConnections sync.Map
systemConnections map[string]*ssh.Client
sshClientConfig *ssh.ClientConfig sshClientConfig *ssh.ClientConfig
pubKey string pubKey string
am *alerts.AlertManager am *alerts.AlertManager
@@ -49,8 +48,6 @@ type Hub struct {
func NewHub(app *pocketbase.PocketBase) *Hub { func NewHub(app *pocketbase.PocketBase) *Hub {
return &Hub{ return &Hub{
app: app, app: app,
connectionLock: &sync.Mutex{},
systemConnections: make(map[string]*ssh.Client),
am: alerts.NewAlertManager(app), am: alerts.NewAlertManager(app),
um: users.NewUserManager(app), um: users.NewUserManager(app),
rm: records.NewRecordManager(app), rm: records.NewRecordManager(app),
@@ -262,9 +259,9 @@ func (h *Hub) updateSystem(record *models.Record) {
var client *ssh.Client var client *ssh.Client
var err error var err error
// check if system connection data exists // check if system connection exists
if _, ok := h.systemConnections[record.Id]; ok { if existingClient, ok := h.systemConnections.Load(record.Id); ok {
client = h.systemConnections[record.Id] client = existingClient.(*ssh.Client)
} else { } else {
// create system connection // create system connection
client, err = h.createSystemConnection(record) client, err = h.createSystemConnection(record)
@@ -275,9 +272,7 @@ func (h *Hub) updateSystem(record *models.Record) {
} }
return return
} }
h.connectionLock.Lock() h.systemConnections.Store(record.Id, client)
h.systemConnections[record.Id] = client
h.connectionLock.Unlock()
} }
// get system stats from agent // get system stats from agent
var systemData system.CombinedData var systemData system.CombinedData
@@ -286,6 +281,7 @@ func (h *Hub) updateSystem(record *models.Record) {
// if previous connection was closed, try again // if previous connection was closed, try again
h.app.Logger().Error("Existing SSH connection closed. Retrying...", "host", record.GetString("host"), "port", record.GetString("port")) h.app.Logger().Error("Existing SSH connection closed. Retrying...", "host", record.GetString("host"), "port", record.GetString("port"))
h.deleteSystemConnection(record) h.deleteSystemConnection(record)
time.Sleep(time.Millisecond * 100)
h.updateSystem(record) h.updateSystem(record)
return return
} }
@@ -359,14 +355,13 @@ func (h *Hub) updateSystemStatus(record *models.Record, status string) {
} }
} }
// delete system connection from map and close connection
func (h *Hub) deleteSystemConnection(record *models.Record) { func (h *Hub) deleteSystemConnection(record *models.Record) {
if _, ok := h.systemConnections[record.Id]; ok { if client, ok := h.systemConnections.Load(record.Id); ok {
if h.systemConnections[record.Id] != nil { if sshClient := client.(*ssh.Client); sshClient != nil {
h.systemConnections[record.Id].Close() sshClient.Close()
} }
h.connectionLock.Lock() h.systemConnections.Delete(record.Id)
defer h.connectionLock.Unlock()
delete(h.systemConnections, record.Id)
} }
} }