diff --git a/internal/hub/systems/network_monitor_ssh_test.go b/internal/hub/systems/network_monitor_ssh_test.go new file mode 100644 index 00000000..b733daf7 --- /dev/null +++ b/internal/hub/systems/network_monitor_ssh_test.go @@ -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()) +} diff --git a/internal/hub/systems/network_monitor_sync_test.go b/internal/hub/systems/network_monitor_sync_test.go index f9c7666e..ea50d097 100644 --- a/internal/hub/systems/network_monitor_sync_test.go +++ b/internal/hub/systems/network_monitor_sync_test.go @@ -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)) diff --git a/internal/hub/systems/network_monitors.go b/internal/hub/systems/network_monitors.go index 0fb44ed4..ec669ce5 100644 --- a/internal/hub/systems/network_monitors.go +++ b/internal/hub/systems/network_monitors.go @@ -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}) diff --git a/internal/hub/systems/system.go b/internal/hub/systems/system.go index 6b0ca713..7465ff6a 100644 --- a/internal/hub/systems/system.go +++ b/internal/hub/systems/system.go @@ -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 diff --git a/internal/hub/systems/system_manager.go b/internal/hub/systems/system_manager.go index db279609..d1ff9633 100644 --- a/internal/hub/systems/system_manager.go +++ b/internal/hub/systems/system_manager.go @@ -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 }