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

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())
+}

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
+	resp := common.AgentResponse{Id: req.Id}
-	data, _ := cbor.Marshal(monitor.SyncResponse{})
+	if req.Action == common.GetData {
-	response, _ := cbor.Marshal(common.AgentResponse{Id: req.Id, Data: data})
+		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))

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})

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

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() {
+	go system.syncPendingNetworkMonitors()
-		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)
-		}
-	}()
 
 	return nil
 }