From eb5dd230cf01995dd16518e2af18b0486ebe1d11 Mon Sep 17 00:00:00 2001 From: Sven van Ginkel Date: Sat, 15 Aug 2026 19:40:30 +0200 Subject: [PATCH] fix(hub): remove stale smart_devices records when a drive is no longer reported (#2178) * Fix duplicate /dev/sdg-style entries * only prune devices after complete refreshes --------- Co-authored-by: henrygd --- agent/handlers.go | 12 +- agent/handlers_test.go | 13 +++ agent/response.go | 3 + agent/smart.go | 7 +- internal/common/common-ws.go | 3 +- internal/entities/smart/smart.go | 9 +- internal/hub/systems/system.go | 11 +- internal/hub/systems/system_smart.go | 59 +++++++--- internal/hub/systems/system_smart_test.go | 135 ++++++++++++++++++++++ internal/hub/transport/transport.go | 22 ++-- internal/hub/transport/transport_test.go | 33 ++++++ 11 files changed, 273 insertions(+), 34 deletions(-) create mode 100644 internal/hub/transport/transport_test.go diff --git a/agent/handlers.go b/agent/handlers.go index 2db07d31..adf9c918 100644 --- a/agent/handlers.go +++ b/agent/handlers.go @@ -166,14 +166,16 @@ type GetSmartDataHandler struct{} func (h *GetSmartDataHandler) Handle(hctx *HandlerContext) error { if hctx.Agent.smartManager == nil { - // return empty map to indicate no data - return hctx.SendResponse(map[string]smart.SmartData{}, hctx.RequestID) + return hctx.SendResponse(smart.SmartDataResponse{Data: map[string]smart.SmartData{}}, hctx.RequestID) } - if err := hctx.Agent.smartManager.Refresh(false); err != nil { + complete, err := hctx.Agent.smartManager.Refresh(false) + if err != nil { slog.Debug("smart refresh failed", "err", err) } - data := hctx.Agent.smartManager.GetCurrentData() - return hctx.SendResponse(data, hctx.RequestID) + return hctx.SendResponse(smart.SmartDataResponse{ + Data: hctx.Agent.smartManager.GetCurrentData(), + Complete: complete, + }, hctx.RequestID) } //////////////////////////////////////////////////////////////////////////// diff --git a/agent/handlers_test.go b/agent/handlers_test.go index 5df132b4..df7fa869 100644 --- a/agent/handlers_test.go +++ b/agent/handlers_test.go @@ -7,6 +7,7 @@ import ( "github.com/fxamacker/cbor/v2" "github.com/henrygd/beszel/internal/common" + "github.com/henrygd/beszel/internal/entities/smart" "github.com/stretchr/testify/assert" ) @@ -17,6 +18,18 @@ type MockHandler struct { handleFunc func(ctx *HandlerContext) error } +func TestNewAgentResponseSmartData(t *testing.T) { + response := newAgentResponse(smart.SmartDataResponse{ + Data: map[string]smart.SmartData{ + "AAA": {SerialNumber: "AAA"}, + }, + Complete: true, + }, nil) + + assert.Equal(t, "AAA", response.SmartData["AAA"].SerialNumber) + assert.True(t, response.SmartComplete) +} + func (m *MockHandler) Handle(ctx *HandlerContext) error { if m.handleFunc != nil { return m.handleFunc(ctx) diff --git a/agent/response.go b/agent/response.go index 5bede047..5d588ad7 100644 --- a/agent/response.go +++ b/agent/response.go @@ -21,6 +21,9 @@ func newAgentResponse(data any, requestID *uint32) common.AgentResponse { response.String = &v case map[string]smart.SmartData: response.SmartData = v + case smart.SmartDataResponse: + response.SmartData = v.Data + response.SmartComplete = v.Complete case systemd.ServiceDetails: response.ServiceInfo = v default: diff --git a/agent/smart.go b/agent/smart.go index 0c052f21..41abb14f 100644 --- a/agent/smart.go +++ b/agent/smart.go @@ -70,8 +70,9 @@ type deviceKey struct { var errNoValidSmartData = fmt.Errorf("no valid SMART data found") // Error for missing data -// Refresh updates SMART data for all known devices -func (sm *SmartManager) Refresh(forceScan bool) error { +// Refresh updates SMART data for all known devices and reports whether every +// discovered device was collected successfully. +func (sm *SmartManager) Refresh(forceScan bool) (bool, error) { sm.refreshMutex.Lock() defer sm.refreshMutex.Unlock() @@ -92,7 +93,7 @@ func (sm *SmartManager) Refresh(forceScan bool) error { } } - return sm.resolveRefreshError(scanErr, collectErr) + return scanErr == nil && collectErr == nil, sm.resolveRefreshError(scanErr, collectErr) } // devicesSnapshot returns a copy of the current device slice to avoid iterating diff --git a/internal/common/common-ws.go b/internal/common/common-ws.go index 4f2f5d22..dee77918 100644 --- a/internal/common/common-ws.go +++ b/internal/common/common-ws.go @@ -42,7 +42,8 @@ type AgentResponse struct { SmartData map[string]smart.SmartData `cbor:"5,keyasint,omitempty,omitzero"` // Legacy (<= 0.17) ServiceInfo systemd.ServiceDetails `cbor:"6,keyasint,omitempty,omitzero"` // Legacy (<= 0.17) // Data is the generic response payload for new endpoints (0.18+) - Data cbor.RawMessage `cbor:"7,keyasint,omitempty,omitzero"` + Data cbor.RawMessage `cbor:"7,keyasint,omitempty,omitzero"` + SmartComplete bool `cbor:"8,keyasint,omitempty,omitzero"` } type FingerprintRequest struct { diff --git a/internal/entities/smart/smart.go b/internal/entities/smart/smart.go index 24c7eac2..6d832e70 100644 --- a/internal/entities/smart/smart.go +++ b/internal/entities/smart/smart.go @@ -494,7 +494,7 @@ type SmartInfoForNvme struct { FirmwareVersion string `json:"firmware_version"` // NVMePCIVendor NVMePCIVendor `json:"nvme_pci_vendor"` // NVMeIEEEOUIIdentifier uint32 `json:"nvme_ieee_oui_identifier"` - NVMeTotalCapacity uint64 `json:"nvme_total_capacity"` + NVMeTotalCapacity uint64 `json:"nvme_total_capacity"` // NVMeUnallocatedCapacity uint64 `json:"nvme_unallocated_capacity"` // NVMeControllerID uint16 `json:"nvme_controller_id"` // NVMeVersion VersionStringInfo `json:"nvme_version"` @@ -531,6 +531,13 @@ type SmartData struct { Attributes []*SmartAttribute `json:"a,omitempty" cbor:"9,keyasint,omitempty"` } +// SmartDataResponse contains the collected data and whether every discovered +// device was collected. Older agents omit Complete, so hubs must not prune from it. +type SmartDataResponse struct { + Data map[string]SmartData `json:"data" cbor:"0,keyasint"` + Complete bool `json:"complete" cbor:"1,keyasint,omitempty"` // Whether every discovered device was collected +} + type SmartAttribute struct { ID uint16 `json:"id,omitempty" cbor:"0,keyasint,omitempty"` Name string `json:"n" cbor:"1,keyasint"` diff --git a/internal/hub/systems/system.go b/internal/hub/systems/system.go index c92c1884..0fc99700 100644 --- a/internal/hub/systems/system.go +++ b/internal/hub/systems/system.go @@ -532,11 +532,16 @@ func (sys *System) FetchSystemdInfoFromAgent(serviceName string) (systemd.Servic return result, err } -// FetchSmartDataFromAgent fetches SMART data from the agent -func (sys *System) FetchSmartDataFromAgent() (map[string]smart.SmartData, error) { +// FetchSmartDataFromAgent fetches SMART data from the agent. +func (sys *System) FetchSmartDataFromAgent() (smart.SmartDataResponse, error) { ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) defer cancel() - var result map[string]smart.SmartData + if sys.agentVersion.LT(beszel.MinVersionAgentResponse) { + var data map[string]smart.SmartData + err := sys.request(ctx, common.GetSmartData, nil, &data) + return smart.SmartDataResponse{Data: data}, err + } + var result smart.SmartDataResponse err := sys.request(ctx, common.GetSmartData, nil, &result) return result, err } diff --git a/internal/hub/systems/system_smart.go b/internal/hub/systems/system_smart.go index c40f643d..68b750fe 100644 --- a/internal/hub/systems/system_smart.go +++ b/internal/hub/systems/system_smart.go @@ -7,6 +7,7 @@ import ( "time" "github.com/henrygd/beszel/internal/entities/smart" + "github.com/pocketbase/dbx" "github.com/pocketbase/pocketbase/core" ) @@ -17,13 +18,13 @@ type smartFetchState struct { // FetchAndSaveSmartDevices fetches SMART data from the agent and saves it to the database func (sys *System) FetchAndSaveSmartDevices() error { - smartData, err := sys.FetchSmartDataFromAgent() + response, err := sys.FetchSmartDataFromAgent() if err != nil { sys.recordSmartFetchResult(err, 0) return err } - err = sys.saveSmartDevices(smartData) - sys.recordSmartFetchResult(err, len(smartData)) + err = sys.saveSmartDevices(response.Data, response.Complete) + sys.recordSmartFetchResult(err, len(response.Data)) return err } @@ -61,8 +62,9 @@ func (sys *System) smartFetchInterval() time.Duration { return time.Hour } -// saveSmartDevices saves SMART device data to the smart_devices collection -func (sys *System) saveSmartDevices(smartData map[string]smart.SmartData) error { +// saveSmartDevices saves SMART device data and, after a complete refresh, +// removes rows for devices that are no longer reported. +func (sys *System) saveSmartDevices(smartData map[string]smart.SmartData, complete bool) error { if len(smartData) == 0 { return nil } @@ -73,20 +75,49 @@ func (sys *System) saveSmartDevices(smartData map[string]smart.SmartData) error return err } - for deviceKey, device := range smartData { - if err := sys.upsertSmartDeviceRecord(collection, deviceKey, device); err != nil { - return err - } + currentIDs := make(map[string]struct{}, len(smartData)) + for deviceKey := range smartData { + currentIDs[makeStableHashId(sys.Id, deviceKey)] = struct{}{} } - return nil + err = hub.RunInTransaction(func(txApp core.App) error { + if complete { + existing, err := txApp.FindRecordsByFilter( + collection, + "system = {:system}", + "", + 0, + 0, + dbx.Params{"system": sys.Id}, + ) + if err != nil { + return err + } + for _, record := range existing { + if _, ok := currentIDs[record.Id]; ok { + continue + } + if err := txApp.Delete(record); err != nil { + return err + } + } + } + + for deviceKey, device := range smartData { + if err := sys.upsertSmartDeviceRecord(txApp, collection, deviceKey, device); err != nil { + return err + } + } + + return nil + }) + return err } -func (sys *System) upsertSmartDeviceRecord(collection *core.Collection, deviceKey string, device smart.SmartData) error { - hub := sys.manager.hub +func (sys *System) upsertSmartDeviceRecord(app core.App, collection *core.Collection, deviceKey string, device smart.SmartData) error { recordID := makeStableHashId(sys.Id, deviceKey) - record, err := hub.FindRecordById(collection, recordID) + record, err := app.FindRecordById(collection, recordID) if err != nil { if !errors.Is(err, sql.ErrNoRows) { return err @@ -114,7 +145,7 @@ func (sys *System) upsertSmartDeviceRecord(collection *core.Collection, deviceKe record.Set("cycles", powerCycles) record.Set("attributes", device.Attributes) - return hub.SaveNoValidate(record) + return app.SaveNoValidate(record) } // extractPowerMetrics extracts power on hours and power cycles from SMART attributes diff --git a/internal/hub/systems/system_smart_test.go b/internal/hub/systems/system_smart_test.go index e186b1ff..7b47d770 100644 --- a/internal/hub/systems/system_smart_test.go +++ b/internal/hub/systems/system_smart_test.go @@ -7,10 +7,53 @@ import ( "testing" "time" + "github.com/henrygd/beszel/internal/entities/smart" + esystem "github.com/henrygd/beszel/internal/entities/system" "github.com/henrygd/beszel/internal/hub/expirymap" + _ "github.com/henrygd/beszel/internal/migrations" + + "github.com/pocketbase/pocketbase/core" + pbtests "github.com/pocketbase/pocketbase/tests" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "golang.org/x/crypto/ssh" ) +// stubHub implements hubLike using a plain pocketbase test app, so +// smart-device DB tests can run in-package without an import cycle to +// internal/hub (which imports this package). +type stubHub struct{ core.App } + +func (stubHub) GetSSHKey(dataDir string) (ssh.Signer, error) { return nil, nil } +func (stubHub) HandleSystemAlerts(systemRecord *core.Record, data *esystem.CombinedData) error { + return nil +} +func (stubHub) HandleStatusAlerts(status string, systemRecord *core.Record) error { return nil } +func (stubHub) CancelPendingStatusAlerts(systemID string) {} + +// newTestSystemWithHub creates a System backed by a real (temp) database, along +// with a matching "systems" record, for tests that need to exercise DB reads/writes. +func newTestSystemWithHub(t *testing.T) (*System, *pbtests.TestApp) { + t.Helper() + testApp, err := pbtests.NewTestApp(t.TempDir()) + require.NoError(t, err) + t.Cleanup(testApp.Cleanup) + + sm := &SystemManager{hub: stubHub{testApp}, smartFetchMap: expirymap.New[smartFetchState](time.Hour)} + t.Cleanup(sm.smartFetchMap.StopCleaner) + + col, err := testApp.FindCachedCollectionByNameOrId("systems") + require.NoError(t, err) + systemRecord := core.NewRecord(col) + systemRecord.Set("name", "test-system") + systemRecord.Set("host", "127.0.0.1") + systemRecord.Set("port", "45876") + require.NoError(t, testApp.SaveNoValidate(systemRecord)) + + sys := &System{Id: systemRecord.Id, manager: sm} + return sys, testApp +} + func TestRecordSmartFetchResult(t *testing.T) { sm := &SystemManager{smartFetchMap: expirymap.New[smartFetchState](time.Hour)} t.Cleanup(sm.smartFetchMap.StopCleaner) @@ -92,3 +135,95 @@ func TestResetFailedSmartFetchState(t *testing.T) { _, ok = sm.smartFetchMap.GetOk("system-1") assert.True(t, ok, "expected successful smart fetch state to be preserved") } + +// countSmartDeviceRecords returns the number of smart_devices rows for the given system. +func countSmartDeviceRecords(t *testing.T, app core.App, systemID string) []*core.Record { + t.Helper() + records, err := app.FindAllRecords("smart_devices", nil) + require.NoError(t, err) + var forSystem []*core.Record + for _, r := range records { + if r.GetString("system") == systemID { + forSystem = append(forSystem, r) + } + } + return forSystem +} + +func TestSaveSmartDevices_RemovesStaleDevices(t *testing.T) { + sys, testApp := newTestSystemWithHub(t) + + // first fetch reports two devices: sda (serial AAA) and sdb (serial BBB) + err := sys.saveSmartDevices(map[string]smart.SmartData{ + "AAA": {SerialNumber: "AAA", DiskName: "sda", ModelName: "Disk A"}, + "BBB": {SerialNumber: "BBB", DiskName: "sdb", ModelName: "Disk B"}, + }, true) + require.NoError(t, err) + + records := countSmartDeviceRecords(t, testApp, sys.Id) + require.Len(t, records, 2, "expected both devices to be saved") + + var recordA *core.Record + for _, r := range records { + if r.GetString("serial") == "AAA" { + recordA = r + } + } + require.NotNil(t, recordA, "expected to find device AAA") + originalID := recordA.Id + + deleteEvents := 0 + testApp.OnRecordAfterDeleteSuccess("smart_devices").BindFunc(func(e *core.RecordEvent) error { + deleteEvents++ + return e.Next() + }) + + // A complete refresh confirms that BBB is gone, so remove it through + // PocketBase and notify realtime subscribers. + err = sys.saveSmartDevices(map[string]smart.SmartData{ + "AAA": {SerialNumber: "AAA", DiskName: "sda", ModelName: "Disk A", Temperature: 42}, + }, true) + require.NoError(t, err) + + records = countSmartDeviceRecords(t, testApp, sys.Id) + require.Len(t, records, 1, "expected stale device BBB to be removed") + assert.Equal(t, "AAA", records[0].GetString("serial")) + assert.Equal(t, originalID, records[0].Id, "expected existing device to be updated in place, not recreated") + assert.EqualValues(t, 42, records[0].GetInt("temp")) + assert.Equal(t, 1, deleteEvents, "expected PocketBase delete hooks to run") +} + +func TestSaveSmartDevices_IncompleteDataDoesNotRemoveDevices(t *testing.T) { + sys, testApp := newTestSystemWithHub(t) + + require.NoError(t, sys.saveSmartDevices(map[string]smart.SmartData{ + "AAA": {SerialNumber: "AAA", DiskName: "sda"}, + "BBB": {SerialNumber: "BBB", DiskName: "sdb"}, + }, true)) + + // AAA was collected but BBB failed. The response is useful for updating AAA, + // but it is not authoritative enough to remove BBB. + require.NoError(t, sys.saveSmartDevices(map[string]smart.SmartData{ + "AAA": {SerialNumber: "AAA", DiskName: "sda", Temperature: 42}, + }, false)) + + assert.Len(t, countSmartDeviceRecords(t, testApp, sys.Id), 2) + recordA, err := testApp.FindRecordById("smart_devices", makeStableHashId(sys.Id, "AAA")) + require.NoError(t, err) + assert.EqualValues(t, 42, recordA.GetInt("temp")) +} + +func TestSaveSmartDevices_EmptyDataIsNoop(t *testing.T) { + sys, testApp := newTestSystemWithHub(t) + + err := sys.saveSmartDevices(map[string]smart.SmartData{ + "AAA": {SerialNumber: "AAA", DiskName: "sda"}, + }, true) + require.NoError(t, err) + + err = sys.saveSmartDevices(map[string]smart.SmartData{}, true) + require.NoError(t, err) + + records := countSmartDeviceRecords(t, testApp, sys.Id) + assert.Len(t, records, 1, "empty fetch result should not delete existing devices") +} diff --git a/internal/hub/transport/transport.go b/internal/hub/transport/transport.go index 3b92e1f1..8f99940f 100644 --- a/internal/hub/transport/transport.go +++ b/internal/hub/transport/transport.go @@ -88,15 +88,23 @@ func unmarshalLegacyResponse(resp common.AgentResponse, action common.WebSocketA *d = *resp.String return nil case common.GetSmartData: - d, ok := dest.(*map[string]smart.SmartData) - if !ok { + switch d := dest.(type) { + case *map[string]smart.SmartData: + if resp.SmartData == nil { + return errors.New("no SMART data in response") + } + *d = resp.SmartData + return nil + case *smart.SmartDataResponse: + if resp.SmartData == nil { + return errors.New("no SMART data in response") + } + d.Data = resp.SmartData + d.Complete = resp.SmartComplete + return nil + default: return fmt.Errorf("unexpected dest type for GetSmartData: %T", dest) } - if resp.SmartData == nil { - return errors.New("no SMART data in response") - } - *d = resp.SmartData - return nil case common.GetSystemdInfo: d, ok := dest.(*systemd.ServiceDetails) if !ok { diff --git a/internal/hub/transport/transport_test.go b/internal/hub/transport/transport_test.go new file mode 100644 index 00000000..9ad417b5 --- /dev/null +++ b/internal/hub/transport/transport_test.go @@ -0,0 +1,33 @@ +package transport + +import ( + "testing" + + "github.com/henrygd/beszel/internal/common" + "github.com/henrygd/beszel/internal/entities/smart" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestUnmarshalSmartDataResponse(t *testing.T) { + for _, test := range []struct { + name string + complete bool + }{ + {name: "complete response", complete: true}, + {name: "older agent defaults to incomplete", complete: false}, + } { + t.Run(test.name, func(t *testing.T) { + response := common.AgentResponse{ + SmartData: map[string]smart.SmartData{ + "AAA": {SerialNumber: "AAA"}, + }, + SmartComplete: test.complete, + } + var result smart.SmartDataResponse + require.NoError(t, UnmarshalResponse(response, common.GetSmartData, &result)) + assert.Equal(t, test.complete, result.Complete) + assert.Equal(t, "AAA", result.Data["AAA"].SerialNumber) + }) + } +}