mirror of
https://github.com/fankes/beszel.git
synced 2025-10-19 17:59:28 +08:00
use sync.Map for system connections
This commit is contained in:
@@ -35,8 +35,7 @@ import (
|
||||
|
||||
type Hub struct {
|
||||
app *pocketbase.PocketBase
|
||||
connectionLock *sync.Mutex
|
||||
systemConnections map[string]*ssh.Client
|
||||
systemConnections sync.Map
|
||||
sshClientConfig *ssh.ClientConfig
|
||||
pubKey string
|
||||
am *alerts.AlertManager
|
||||
@@ -49,8 +48,6 @@ type Hub struct {
|
||||
func NewHub(app *pocketbase.PocketBase) *Hub {
|
||||
return &Hub{
|
||||
app: app,
|
||||
connectionLock: &sync.Mutex{},
|
||||
systemConnections: make(map[string]*ssh.Client),
|
||||
am: alerts.NewAlertManager(app),
|
||||
um: users.NewUserManager(app),
|
||||
rm: records.NewRecordManager(app),
|
||||
@@ -262,9 +259,9 @@ func (h *Hub) updateSystem(record *models.Record) {
|
||||
var client *ssh.Client
|
||||
var err error
|
||||
|
||||
// check if system connection data exists
|
||||
if _, ok := h.systemConnections[record.Id]; ok {
|
||||
client = h.systemConnections[record.Id]
|
||||
// check if system connection exists
|
||||
if existingClient, ok := h.systemConnections.Load(record.Id); ok {
|
||||
client = existingClient.(*ssh.Client)
|
||||
} else {
|
||||
// create system connection
|
||||
client, err = h.createSystemConnection(record)
|
||||
@@ -275,9 +272,7 @@ func (h *Hub) updateSystem(record *models.Record) {
|
||||
}
|
||||
return
|
||||
}
|
||||
h.connectionLock.Lock()
|
||||
h.systemConnections[record.Id] = client
|
||||
h.connectionLock.Unlock()
|
||||
h.systemConnections.Store(record.Id, client)
|
||||
}
|
||||
// get system stats from agent
|
||||
var systemData system.CombinedData
|
||||
@@ -286,6 +281,7 @@ func (h *Hub) updateSystem(record *models.Record) {
|
||||
// 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.deleteSystemConnection(record)
|
||||
time.Sleep(time.Millisecond * 100)
|
||||
h.updateSystem(record)
|
||||
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) {
|
||||
if _, ok := h.systemConnections[record.Id]; ok {
|
||||
if h.systemConnections[record.Id] != nil {
|
||||
h.systemConnections[record.Id].Close()
|
||||
if client, ok := h.systemConnections.Load(record.Id); ok {
|
||||
if sshClient := client.(*ssh.Client); sshClient != nil {
|
||||
sshClient.Close()
|
||||
}
|
||||
h.connectionLock.Lock()
|
||||
defer h.connectionLock.Unlock()
|
||||
delete(h.systemConnections, record.Id)
|
||||
h.systemConnections.Delete(record.Id)
|
||||
}
|
||||
}
|
||||
|
||||
|
Reference in New Issue
Block a user