fix: resync network monitors after SSH reconnect and retry failed syncs

This commit is contained in:
henrygd
2026-09-18 23:49:54 -04:00
parent 2784460621
commit 187dc886a9
5 changed files with 238 additions and 17 deletions
@@ -0,0 +1,159 @@
//go:build testing
package systems
import (
"context"
"crypto/ed25519"
"crypto/rand"
"net"
"sync/atomic"
"testing"
"time"
"github.com/fxamacker/cbor/v2"
"github.com/henrygd/beszel/internal/common"
"github.com/henrygd/beszel/internal/entities/monitor"
esystem "github.com/henrygd/beszel/internal/entities/system"
"github.com/henrygd/beszel/internal/hub/expirymap"
"github.com/pocketbase/pocketbase/core"
"github.com/stretchr/testify/require"
"golang.org/x/crypto/ssh"
)
func TestSSHNetworkMonitorReconnectSync(t *testing.T) {
sys, app := newTestSystemWithHub(t)
sys.manager.zfsFetchMap = expirymap.New[zfsFetchState](time.Hour)
t.Cleanup(sys.manager.zfsFetchMap.StopCleaner)
sys.ctx = context.Background()
sys.Status = up
_, key, err := ed25519.GenerateKey(rand.Reader)
require.NoError(t, err)
signer, err := ssh.NewSignerFromKey(key)
require.NoError(t, err)
config := &ssh.ServerConfig{NoClientAuth: true, ServerVersion: "SSH-2.0-beszel_0.20.0"}
config.AddHostKey(signer)
listener, err := net.Listen("tcp", "127.0.0.1:0")
require.NoError(t, err)
t.Cleanup(func() { _ = listener.Close() })
sys.Host, sys.Port, err = net.SplitHostPort(listener.Addr().String())
require.NoError(t, err)
sys.manager.sshConfig = &ssh.ClientConfig{User: "test", HostKeyCallback: ssh.InsecureIgnoreHostKey(), Timeout: time.Second}
t.Cleanup(sys.closeSSHConnection)
requests := make(chan monitor.SyncRequest, 10)
var failSync atomic.Bool
go func() {
for {
conn, err := listener.Accept()
if err != nil {
return
}
go func() {
server, channels, reqs, err := ssh.NewServerConn(conn, config)
if err != nil {
_ = conn.Close()
return
}
defer server.Close()
go ssh.DiscardRequests(reqs)
for channel := range channels {
ch, reqs, err := channel.Accept()
if err != nil {
return
}
go func() {
defer ch.Close()
for req := range reqs {
if req.Type != "shell" {
_ = req.Reply(false, nil)
continue
}
_ = req.Reply(true, nil)
var request common.HubRequest[cbor.RawMessage]
if cbor.NewDecoder(ch).Decode(&request) != nil {
return
}
response := common.AgentResponse{}
switch request.Action {
case common.GetData:
response.SystemData = &esystem.CombinedData{}
case common.SyncNetworkMonitors:
var syncReq monitor.SyncRequest
if cbor.Unmarshal(request.Data, &syncReq) != nil {
return
}
requests <- syncReq
if failSync.Load() {
response.Error = "test sync failure"
} else {
response.Data, _ = cbor.Marshal(monitor.SyncResponse{})
}
}
_ = cbor.NewEncoder(ch).Encode(response)
_, _ = ch.SendRequest("exit-status", false, ssh.Marshal(struct{ Status uint32 }{0}))
return
}
}()
}
}()
}
}()
collection, err := app.FindCachedCollectionByNameOrId("network_monitors")
require.NoError(t, err)
probe := core.NewRecord(collection)
probe.Load(map[string]any{"system": sys.Id, "target": "localhost", "protocol": "tcp", "port": 80, "interval": 60, "enabled": true})
require.NoError(t, app.SaveNoValidate(probe))
fetch := func() {
t.Helper()
_, err := sys.fetchDataFromAgent(common.DataRequestOptions{})
require.NoError(t, err, "monitor sync failure must not fail stats fetching")
}
receive := func() monitor.SyncRequest {
t.Helper()
select {
case req := <-requests:
require.Equal(t, monitor.SyncActionReplace, req.Action)
return req
case <-time.After(time.Second):
t.Fatal("missing full monitor sync")
return monitor.SyncRequest{}
}
}
fetch()
require.Equal(t, probe.Id, receive().Configs[0].ID)
require.False(t, sys.monitorsNeedSync.Load())
fetch()
require.Empty(t, requests, "steady-state fetch must not resync")
// Simulate loss of the agent process/connection and its in-memory monitors.
require.NoError(t, sys.client.Load().Close())
fetch()
require.Equal(t, probe.Id, receive().Configs[0].ID)
require.False(t, sys.monitorsNeedSync.Load())
// Failed replacements are retried on the next successful stats fetch.
require.NoError(t, sys.client.Load().Close())
failSync.Store(true)
fetch()
receive()
require.True(t, sys.monitorsNeedSync.Load())
failSync.Store(false)
fetch()
receive()
require.False(t, sys.monitorsNeedSync.Load())
probe.Set("enabled", false)
require.NoError(t, app.SaveNoValidate(probe))
require.NoError(t, sys.client.Load().Close())
fetch()
require.Empty(t, receive().Configs, "empty replacement must clear stale monitors")
}
func TestPendingNetworkMonitorSyncQueryFailure(t *testing.T) {
sys, app := newTestSystemWithHub(t)
_, err := app.DB().NewQuery("DROP TABLE network_monitors").Execute()
require.NoError(t, err)
sys.monitorsNeedSync.Store(true)
sys.syncPendingNetworkMonitors()
require.True(t, sys.monitorsNeedSync.Load())
}
@@ -6,6 +6,7 @@ import (
"net/http"
"net/http/httptest"
"strings"
"sync/atomic"
"testing"
"time"
@@ -13,6 +14,7 @@ import (
"github.com/fxamacker/cbor/v2"
"github.com/henrygd/beszel/internal/common"
"github.com/henrygd/beszel/internal/entities/monitor"
esystem "github.com/henrygd/beszel/internal/entities/system"
"github.com/henrygd/beszel/internal/hub/ws"
"github.com/lxzan/gws"
"github.com/pocketbase/pocketbase/core"
@@ -22,17 +24,31 @@ import (
type monitorSyncClient struct {
gws.BuiltinEventHandler
requests chan common.HubRequest[monitor.SyncRequest]
failSync atomic.Bool
}
func (c *monitorSyncClient) OnMessage(conn *gws.Conn, message *gws.Message) {
defer message.Close()
var req common.HubRequest[monitor.SyncRequest]
var req common.HubRequest[cbor.RawMessage]
if err := cbor.Unmarshal(message.Bytes(), &req); err != nil {
return
}
c.requests <- req
data, _ := cbor.Marshal(monitor.SyncResponse{})
response, _ := cbor.Marshal(common.AgentResponse{Id: req.Id, Data: data})
resp := common.AgentResponse{Id: req.Id}
if req.Action == common.GetData {
resp.SystemData = &esystem.CombinedData{}
} else {
var data monitor.SyncRequest
if err := cbor.Unmarshal(req.Data, &data); err != nil {
return
}
c.requests <- common.HubRequest[monitor.SyncRequest]{Id: req.Id, Action: req.Action, Data: data}
if c.failSync.Load() {
resp.Error = "test sync failure"
} else {
resp.Data, _ = cbor.Marshal(monitor.SyncResponse{})
}
}
response, _ := cbor.Marshal(resp)
_ = conn.WriteMessage(gws.OpcodeBinary, response)
}
@@ -57,7 +73,7 @@ func TestNetworkMonitorSyncSkipsOlderAgents(t *testing.T) {
}
func TestNetworkMonitorReconnectSync(t *testing.T) {
for _, change := range []string{"delete", "disable"} {
for _, change := range []string{"delete", "disable", "retry"} {
t.Run(change, func(t *testing.T) {
sys, app := newTestSystemWithHub(t)
record, err := app.FindRecordById("systems", sys.Id)
@@ -120,9 +136,33 @@ func TestNetworkMonitorReconnectSync(t *testing.T) {
}
}
client.failSync.Store(change == "retry")
initial := connect()
require.Len(t, initial.Configs, 1)
require.Equal(t, probe.Id, initial.Configs[0].ID)
if change == "retry" {
system, err := sm.GetSystem(sys.Id)
require.NoError(t, err)
require.Eventually(t, system.monitorsNeedSync.Load, time.Second, time.Millisecond)
// A second failed sync must not fail the stats fetch or clear pending state.
_, err = system.fetchDataFromAgent(common.DataRequestOptions{})
require.NoError(t, err)
require.True(t, system.monitorsNeedSync.Load())
require.Len(t, client.requests, 1)
<-client.requests
client.failSync.Store(false)
_, err = system.fetchDataFromAgent(common.DataRequestOptions{})
require.NoError(t, err)
require.False(t, system.monitorsNeedSync.Load())
require.Len(t, client.requests, 1)
retry := <-client.requests
require.Equal(t, monitor.SyncActionReplace, retry.Data.Action)
require.Equal(t, initial.Configs, retry.Data.Configs)
_, err = system.fetchDataFromAgent(common.DataRequestOptions{})
require.NoError(t, err)
require.Empty(t, client.requests, "successful sync must not repeat on every fetch")
return
}
require.NoError(t, sm.RemoveSystem(sys.Id))
if change == "delete" {
require.NoError(t, app.Delete(probe))
+22
View File
@@ -2,6 +2,7 @@ package systems
import (
"context"
"fmt"
"time"
"github.com/henrygd/beszel"
@@ -9,6 +10,27 @@ import (
"github.com/henrygd/beszel/internal/entities/monitor"
)
// syncPendingNetworkMonitors runs on WebSocket connect and after successful stats
// fetches. Failed syncs retry on the next update without taking the system down.
func (sys *System) syncPendingNetworkMonitors() {
if !sys.monitorsNeedSync.Swap(false) {
return
}
if err := sys.syncAllNetworkMonitors(); err != nil {
sys.monitorsNeedSync.Store(true)
sys.manager.hub.Logger().Warn("failed to sync monitors to agent", "system", sys.Id, "err", err)
}
}
func (sys *System) syncAllNetworkMonitors() error {
configs, err := sys.manager.GetMonitorConfigsForSystem(sys.Id)
if err != nil {
return fmt.Errorf("failed to load monitors: %w", err)
}
// An empty set must also replace probes retained across a disconnect.
return sys.SyncNetworkMonitors(configs)
}
// SyncNetworkMonitors sends monitor configurations to the agent.
func (sys *System) SyncNetworkMonitors(configs []monitor.Config) error {
_, err := sys.syncNetworkMonitors(monitor.SyncRequest{Action: monitor.SyncActionReplace, Configs: configs})
+10 -1
View File
@@ -56,6 +56,9 @@ type System struct {
smartInterval time.Duration // Interval for periodic SMART data updates
zfsFetching atomic.Bool // True if ZFS pools are currently being fetched
zfsInterval time.Duration // Interval for periodic ZFS detail data updates
// A fresh connection needs a full monitor configuration sync.
monitorsNeedSync atomic.Bool
// Serialize persistence from scheduled updates and resumes through commit.
recordsMu sync.Mutex
// Protected by recordsMu; realtime reads don't consume probes.
@@ -630,7 +633,10 @@ func (sys *System) request(ctx context.Context, action common.WebSocketAction, r
err := sys.sshTransport.RequestWithRetry(ctx, action, req, dest, 1)
// Keep legacy SSH client/version fields in sync for other code paths.
if sys.sshTransport != nil {
sys.client.Store(sys.sshTransport.GetClient())
client := sys.sshTransport.GetClient()
if previous := sys.client.Swap(client); client != nil && client != previous {
sys.monitorsNeedSync.Store(true)
}
sys.agentVersion = sys.sshTransport.GetAgentVersion()
}
return err
@@ -688,6 +694,7 @@ func (sys *System) fetchDataFromAgent(options common.DataRequestOptions) (*syste
if sys.WsConn != nil && sys.WsConn.IsConnected() {
wsData, err := sys.fetchDataViaWebSocket(options)
if err == nil {
sys.syncPendingNetworkMonitors()
return wsData, nil
}
// close the WebSocket connection if error and try SSH
@@ -698,6 +705,7 @@ func (sys *System) fetchDataFromAgent(options common.DataRequestOptions) (*syste
if err != nil {
return nil, err
}
sys.syncPendingNetworkMonitors()
return sshData, nil
}
@@ -932,6 +940,7 @@ func (s *System) createSSHClient() error {
return err
}
s.agentVersion, _ = extractAgentVersion(string(client.Conn.ServerVersion()))
s.monitorsNeedSync.Store(true)
s.manager.resetFailedSmartFetchState(s.Id)
s.manager.resetFailedZfsFetchState(s.Id)
return nil
+2 -11
View File
@@ -349,23 +349,14 @@ func (sm *SystemManager) AddWebSocketSystem(systemId string, agentVersion semver
system := sm.NewSystem(systemId)
system.WsConn = wsConn
system.agentVersion = agentVersion
system.monitorsNeedSync.Store(true)
if err := sm.AddRecord(systemRecord, system); err != nil {
return err
}
// Sync network monitors to the newly connected agent
go func() {
configs, err := sm.GetMonitorConfigsForSystem(systemId)
if err != nil {
sm.hub.Logger().Warn("failed to load monitors for agent", "system", systemId, "err", err)
return
}
// An empty set must also replace any probes retained across a disconnect.
if err := system.SyncNetworkMonitors(configs); err != nil {
sm.hub.Logger().Warn("failed to sync monitors to agent", "system", systemId, "err", err)
}
}()
go system.syncPendingNetworkMonitors()
return nil
}