fix: prevent WebSocket reconnect loops on slow agent collections (#2294)

The hub waited only 5s (the request manager default) for stats over
WebSocket and closed the connection on any error. On hosts where
`zpool list` stalls (seen on Proxmox, up to its 10s timeout), collection
exceeded that limit, so the hub sent a close (code 1000) and the agent
reconnected. The refresh ran every other cycle, which caused a
disconnect roughly every 2 minutes.

Hub:
- Wait up to 30s for WebSocket stats responses.
- Keep the connection open when a stats request times out; only close
  it (and fall back to SSH) for other errors.

Agent:
- After the first collection, refresh `zpool list` pool stats and
  `zfs list` dataset usage in the background and serve cached values
  meanwhile, so a hung utility cannot delay the stats response.

agent/storage_pool.go

@@ -54,12 +54,18 @@ type poolBackend struct {
 	kernelStatsFn  func() ([]zfs.PoolKernelStat, error) // procfs pool state/I/O source
 	poolStatusesFn func() ([]zfs.PoolStatus, error)     // scrub/vdev detail source
 
-	poolData      []zfs.PoolStat // cached pool inventory (TTL below)
-	lastPoolStats time.Time
-	kernelSamples map[string]poolKernelSample
+	// Utility-backed caches below are refreshed in the background after the
+	// first collection, so cacheMu guards them against those goroutines.
+	cacheMu        sync.Mutex
+	poolData       []zfs.PoolStat // cached pool inventory (TTL below)
+	lastPoolStats  time.Time
+	poolRefreshing bool
 
 	datasetUsage     map[string]zfsDatasetUsage // mountpoint -> usage
 	lastUsageRefresh time.Time
+	usageRefreshing  bool
+
+	kernelSamples map[string]poolKernelSample
 
 	// Detail data (pools, vdevs, scrub, datasets) is cached and refreshed on
 	// an interval. Accessed from handler goroutines, so it is mutex-protected.
@@ -177,21 +183,42 @@ func (b *poolBackend) updateBackendStats(systemStats *system.Stats) {
 }
 
 // poolStats returns the cached pool inventory, calling its collector at most
-// every poolStatsRefreshInterval. On failure the previous inventory is
-// retained and the refresh is retried on the next cadence.
+// every poolStatsRefreshInterval. Only the first collection blocks; later
+// refreshes run in the background because utilities like `zpool list` can hang
+// for seconds on busy hosts, which would otherwise delay the hub's stats
+// response. On failure the previous inventory is retained and the refresh is
+// retried on the next cadence.
 func (b *poolBackend) poolStats() []zfs.PoolStat {
-	if b.lastPoolStats.IsZero() || time.Since(b.lastPoolStats) >= poolStatsRefreshInterval {
-		pools, err := b.poolStatsFn()
-		if err != nil {
-			slog.Debug("Storage pool stats unavailable", "backend", b.name, "err", err)
-		} else {
-			b.poolData = pools
-		}
-		b.lastPoolStats = time.Now()
+	b.cacheMu.Lock()
+	defer b.cacheMu.Unlock()
+	if b.poolRefreshing || (!b.lastPoolStats.IsZero() && time.Since(b.lastPoolStats) < poolStatsRefreshInterval) {
+		return b.poolData
 	}
+	if b.lastPoolStats.IsZero() {
+		b.storePoolStats(b.poolStatsFn())
+		return b.poolData
+	}
+	b.poolRefreshing = true
+	go func() {
+		pools, err := b.poolStatsFn()
+		b.cacheMu.Lock()
+		defer b.cacheMu.Unlock()
+		b.poolRefreshing = false
+		b.storePoolStats(pools, err)
+	}()
 	return b.poolData
 }
 
+// storePoolStats records a pool inventory result. Callers must hold cacheMu.
+func (b *poolBackend) storePoolStats(pools []zfs.PoolStat, err error) {
+	if err != nil {
+		slog.Debug("Storage pool stats unavailable", "backend", b.name, "err", err)
+	} else {
+		b.poolData = pools
+	}
+	b.lastPoolStats = time.Now()
+}
+
 // kernelStats reads cumulative pool counters and converts them to per-second
 // rates. Counter decreases indicate a pool export/import and reset the
 // baseline instead of producing an underflow spike.
@@ -225,12 +252,33 @@ func (b *poolBackend) kernelStats() (map[string]zfs.PoolKernelStat, map[string]z
 }
 
 // refreshDatasetUsage re-runs `zfs list` when the refresh window has elapsed
-// and rebuilds the mountpoint-keyed usage map.
-func (b *poolBackend) refreshDatasetUsage() {
-	if !b.lastUsageRefresh.IsZero() && time.Since(b.lastUsageRefresh) < datasetUsageRefreshInterval {
-		return
+// and returns the mountpoint-keyed usage map. Like poolStats, only the first
+// collection blocks and later refreshes run in the background.
+func (b *poolBackend) refreshDatasetUsage() map[string]zfsDatasetUsage {
+	b.cacheMu.Lock()
+	defer b.cacheMu.Unlock()
+	if b.usageRefreshing || (!b.lastUsageRefresh.IsZero() && time.Since(b.lastUsageRefresh) < datasetUsageRefreshInterval) {
+		return b.datasetUsage
 	}
-	datasets, err := b.datasets()
+	if b.lastUsageRefresh.IsZero() {
+		b.storeDatasetUsage(b.datasets())
+		return b.datasetUsage
+	}
+	b.usageRefreshing = true
+	go func() {
+		datasets, err := b.datasets()
+		b.cacheMu.Lock()
+		defer b.cacheMu.Unlock()
+		b.usageRefreshing = false
+		b.storeDatasetUsage(datasets, err)
+	}()
+	return b.datasetUsage
+}
+
+// storeDatasetUsage rebuilds the usage map from a dataset listing. The map is
+// replaced rather than mutated so returned references stay safe to read.
+// Callers must hold cacheMu.
+func (b *poolBackend) storeDatasetUsage(datasets []zfs.Dataset, err error) {
 	if err != nil {
 		slog.Debug("Storage pool dataset usage unavailable", "backend", b.name, "err", err)
 	} else {
@@ -251,8 +299,7 @@ func (b *poolBackend) refreshDatasetUsage() {
 func (m *StoragePoolManager) DatasetUsage() map[string]zfsDatasetUsage {
 	for _, backend := range m.backends {
 		if backend.name == "zfs" {
-			backend.refreshDatasetUsage()
-			return backend.datasetUsage
+			return backend.refreshDatasetUsage()
 		}
 	}
 	return nil
@@ -442,7 +489,10 @@ func (m *StoragePoolManager) markDuplicateCharts(stats *system.Stats, filesystem
 		}
 	}
 	for _, backend := range m.backends {
-		for _, pool := range backend.poolData {
+		backend.cacheMu.Lock()
+		pools := backend.poolData
+		backend.cacheMu.Unlock()
+		for _, pool := range pools {
 			sample := stats.ZfsPools[pool.Name]
 			if sample == nil || pool.MountID == "" {
 				continue

agent/storage_pool_test.go

@@ -518,3 +518,31 @@ func TestBtrfsPoolIdentities(t *testing.T) {
 	assert.Equal(t, first, zm.GetDetail(true).Pools[1].Name)
 	assert.Equal(t, "renamed", zm.GetDetail(true).Pools[1].DisplayName)
 }
+
+func TestStaleUtilityCachesRefreshInBackground(t *testing.T) {
+	release := make(chan struct{})
+	b := &poolBackend{name: "zfs"}
+	b.poolStatsFn = func() ([]zfs.PoolStat, error) {
+		<-release
+		return []zfs.PoolStat{{Name: "new"}}, nil
+	}
+	b.datasetsFn = func() ([]zfs.Dataset, error) {
+		<-release
+		return []zfs.Dataset{{Name: "new", Mountpoint: "/new"}}, nil
+	}
+	b.poolData = []zfs.PoolStat{{Name: "old"}}
+	b.lastPoolStats = time.Now().Add(-2 * poolStatsRefreshInterval)
+	b.datasetUsage = map[string]zfsDatasetUsage{"/old": {}}
+	b.lastUsageRefresh = time.Now().Add(-2 * datasetUsageRefreshInterval)
+
+	// A hung utility must not block collection; cached data is served meanwhile.
+	for range 2 {
+		assert.Equal(t, "old", b.poolStats()[0].Name)
+		assert.Contains(t, b.refreshDatasetUsage(), "/old")
+	}
+
+	close(release)
+	require.Eventually(t, func() bool {
+		return b.poolStats()[0].Name == "new" && b.refreshDatasetUsage()["/new"] == zfsDatasetUsage{}
+	}, time.Second, time.Millisecond)
+}

internal/hub/systems/system.go

@@ -697,6 +697,11 @@ func (sys *System) fetchDataFromAgent(options common.DataRequestOptions) (*syste
 			sys.syncPendingNetworkMonitors()
 			return wsData, nil
 		}
+		// A slow collection doesn't mean the connection is broken. Closing it
+		// would force the agent into a reconnect loop, so only report the error.
+		if errors.Is(err, context.DeadlineExceeded) {
+			return nil, err
+		}
 		// close the WebSocket connection if error and try SSH
 		sys.closeWebSocketConnection()
 	}
@@ -709,12 +714,19 @@ func (sys *System) fetchDataFromAgent(options common.DataRequestOptions) (*syste
 	return sshData, nil
 }
 
+// wsDataRequestTimeout bounds how long to wait for stats over WebSocket. Agent
+// collection can legitimately take several seconds (e.g. a slow `zpool list`),
+// so this must be well above the request manager's 5s default.
+var wsDataRequestTimeout = 30 * time.Second
+
 func (sys *System) fetchDataViaWebSocket(options common.DataRequestOptions) (*system.CombinedData, error) {
 	if sys.WsConn == nil || !sys.WsConn.IsConnected() {
 		return nil, errors.New("no websocket connection")
 	}
+	ctx, cancel := context.WithTimeout(context.Background(), wsDataRequestTimeout)
+	defer cancel()
 	wsTransport := transport.NewWebSocketTransport(sys.WsConn)
-	err := wsTransport.Request(context.Background(), common.GetData, options, sys.data)
+	err := wsTransport.Request(ctx, common.GetData, options, sys.data)
 	if err != nil {
 		return nil, err
 	}

internal/hub/systems/ws_data_timeout_test.go

@@ -0,0 +1,82 @@
+//go:build testing
+
+package systems
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/blang/semver"
+	"github.com/fxamacker/cbor/v2"
+	"github.com/henrygd/beszel/internal/common"
+	esystem "github.com/henrygd/beszel/internal/entities/system"
+	"github.com/henrygd/beszel/internal/hub/ws"
+	"github.com/lxzan/gws"
+	"github.com/stretchr/testify/require"
+)
+
+// slowDataClient answers GetData only after release is closed, simulating an
+// agent whose collection outlasts the hub's request timeout.
+type slowDataClient struct {
+	gws.BuiltinEventHandler
+	release chan struct{}
+}
+
+func (c *slowDataClient) OnMessage(conn *gws.Conn, message *gws.Message) {
+	defer message.Close()
+	var req common.HubRequest[cbor.RawMessage]
+	if err := cbor.Unmarshal(message.Bytes(), &req); err != nil || req.Action != common.GetData {
+		return
+	}
+	<-c.release
+	response, _ := cbor.Marshal(common.AgentResponse{Id: req.Id, SystemData: &esystem.CombinedData{}})
+	_ = conn.WriteMessage(gws.OpcodeBinary, response)
+}
+
+func TestFetchDataTimeoutKeepsWebSocketOpen(t *testing.T) {
+	originalTimeout := wsDataRequestTimeout
+	wsDataRequestTimeout = 50 * time.Millisecond
+	t.Cleanup(func() { wsDataRequestTimeout = originalTimeout })
+
+	connections := make(chan *ws.WsConn, 1)
+	upgrader := gws.NewUpgrader(&monitorSyncServer{}, nil)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := upgrader.Upgrade(w, r)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		wsConn := ws.NewWsConnection(conn, semver.MustParse("0.20.0"))
+		conn.Session().Store("wsConn", wsConn)
+		connections <- wsConn
+		conn.ReadLoop()
+	}))
+	t.Cleanup(server.Close)
+
+	client := &slowDataClient{release: make(chan struct{})}
+	conn, _, err := gws.NewClient(client, &gws.ClientOption{Addr: "ws" + strings.TrimPrefix(server.URL, "http")})
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.NetConn().Close() })
+	go conn.ReadLoop()
+
+	var sys *System
+	select {
+	case wsConn := <-connections:
+		sys = &System{WsConn: wsConn}
+	case <-time.After(3 * time.Second):
+		t.Fatal("websocket connection was not established")
+	}
+
+	_, err = sys.fetchDataFromAgent(common.DataRequestOptions{})
+	require.ErrorIs(t, err, context.DeadlineExceeded)
+	require.True(t, sys.WsConn.IsConnected(), "a slow collection must not close the connection")
+
+	// The late response is discarded and the next request still succeeds.
+	close(client.release)
+	_, err = sys.fetchDataFromAgent(common.DataRequestOptions{})
+	require.NoError(t, err)
+}