agent: improve disk discovery / IO mapping and add tests (#1811)

agent/disk.go

@@ -14,6 +14,25 @@ import (
 	"github.com/shirou/gopsutil/v4/disk"
 )
 
+// fsRegistrationContext holds the shared lookup state needed to resolve a
+// filesystem into the tracked fsStats key and metadata.
+type fsRegistrationContext struct {
+	filesystem     string // value of optional FILESYSTEM env var
+	isWindows      bool
+	efPath         string // path to extra filesystems (default "/extra-filesystems")
+	diskIoCounters map[string]disk.IOCountersStat
+}
+
+// diskDiscovery groups the transient state for a single initializeDiskInfo run so
+// helper methods can share the same partitions, mount paths, and lookup functions
+type diskDiscovery struct {
+	agent          *Agent
+	rootMountPoint string
+	partitions     []disk.PartitionStat
+	usageFn        func(string) (*disk.UsageStat, error)
+	ctx            fsRegistrationContext
+}
+
 // parseFilesystemEntry parses a filesystem entry in the format "device__customname"
 // Returns the device/filesystem part and the custom name part
 func parseFilesystemEntry(entry string) (device, customName string) {
@@ -27,19 +46,230 @@ func parseFilesystemEntry(entry string) (device, customName string) {
 	return device, customName
 }
 
+// extraFilesystemPartitionInfo derives the I/O device and optional display name
+// for a mounted /extra-filesystems partition. Prefer the partition device reported
+// by the system and only use the folder name for custom naming metadata.
+func extraFilesystemPartitionInfo(p disk.PartitionStat) (device, customName string) {
+	device = strings.TrimSpace(p.Device)
+	folderDevice, customName := parseFilesystemEntry(filepath.Base(p.Mountpoint))
+	if device == "" {
+		device = folderDevice
+	}
+	return device, customName
+}
+
 func isDockerSpecialMountpoint(mountpoint string) bool {
 	switch mountpoint {
 	case "/etc/hosts", "/etc/resolv.conf", "/etc/hostname":
 		return true
-	default:
+	}
+	return false
+}
+
+// registerFilesystemStats resolves the tracked key and stats payload for a
+// filesystem before it is inserted into fsStats.
+func registerFilesystemStats(existing map[string]*system.FsStats, device, mountpoint string, root bool, customName string, ctx fsRegistrationContext) (string, *system.FsStats, bool) {
+	key := device
+	if !ctx.isWindows {
+		key = filepath.Base(device)
+	}
+
+	if root {
+		// Try to map root device to a diskIoCounters entry. First checks for an
+		// exact key match, then uses findIoDevice for normalized / prefix-based
+		// matching (e.g. nda0p2 -> nda0), and finally falls back to FILESYSTEM.
+		if _, ioMatch := ctx.diskIoCounters[key]; !ioMatch {
+			if matchedKey, match := findIoDevice(key, ctx.diskIoCounters); match {
+				key = matchedKey
+			} else if ctx.filesystem != "" {
+				if matchedKey, match := findIoDevice(ctx.filesystem, ctx.diskIoCounters); match {
+					key = matchedKey
+				}
+			}
+			if _, ioMatch = ctx.diskIoCounters[key]; !ioMatch {
+				slog.Warn("Root I/O unmapped; set FILESYSTEM", "device", device, "mountpoint", mountpoint)
+			}
+		}
+	} else {
+		// Check if non-root has diskstats and prefer the folder device for
+		// /extra-filesystems mounts when the discovered partition device is a
+		// mapper path (e.g. luks UUID) that obscures the underlying block device.
+		if _, ioMatch := ctx.diskIoCounters[key]; !ioMatch {
+			if strings.HasPrefix(mountpoint, ctx.efPath) {
+				folderDevice, _ := parseFilesystemEntry(filepath.Base(mountpoint))
+				if folderDevice != "" {
+					if matchedKey, match := findIoDevice(folderDevice, ctx.diskIoCounters); match {
+						key = matchedKey
+					}
+				}
+			}
+			if _, ioMatch = ctx.diskIoCounters[key]; !ioMatch {
+				if matchedKey, match := findIoDevice(key, ctx.diskIoCounters); match {
+					key = matchedKey
+				}
+			}
+		}
+	}
+
+	if _, exists := existing[key]; exists {
+		return "", nil, false
+	}
+
+	fsStats := &system.FsStats{Root: root, Mountpoint: mountpoint}
+	if customName != "" {
+		fsStats.Name = customName
+	}
+	return key, fsStats, true
+}
+
+// addFsStat inserts a discovered filesystem if it resolves to a new tracking
+// key. The key selection itself lives in buildFsStatRegistration so that logic
+// can stay directly unit-tested.
+func (d *diskDiscovery) addFsStat(device, mountpoint string, root bool, customName string) {
+	key, fsStats, ok := registerFilesystemStats(d.agent.fsStats, device, mountpoint, root, customName, d.ctx)
+	if !ok {
+		return
+	}
+	d.agent.fsStats[key] = fsStats
+	name := key
+	if customName != "" {
+		name = customName
+	}
+	slog.Info("Detected disk", "name", name, "device", device, "mount", mountpoint, "io", key, "root", root)
+}
+
+// addConfiguredRootFs resolves FILESYSTEM against partitions first, then falls
+// back to direct diskstats matching for setups like ZFS where partitions do not
+// expose the physical device name.
+func (d *diskDiscovery) addConfiguredRootFs() bool {
+	if d.ctx.filesystem == "" {
 		return false
 	}
+
+	for _, p := range d.partitions {
+		if filesystemMatchesPartitionSetting(d.ctx.filesystem, p) {
+			d.addFsStat(p.Device, p.Mountpoint, true, "")
+			return true
+		}
+	}
+
+	// FILESYSTEM may name a physical disk absent from partitions (e.g. ZFS lists
+	// dataset paths like zroot/ROOT/default, not block devices).
+	if ioKey, match := findIoDevice(d.ctx.filesystem, d.ctx.diskIoCounters); match {
+		d.agent.fsStats[ioKey] = &system.FsStats{Root: true, Mountpoint: d.rootMountPoint}
+		return true
+	}
+
+	slog.Warn("Partition details not found", "filesystem", d.ctx.filesystem)
+	return false
+}
+
+func isRootFallbackPartition(p disk.PartitionStat, rootMountPoint string) bool {
+	return p.Mountpoint == rootMountPoint ||
+		(isDockerSpecialMountpoint(p.Mountpoint) && strings.HasPrefix(p.Device, "/dev"))
+}
+
+// addPartitionRootFs handles the non-configured root fallback path when a
+// partition looks like the active root mount but still needs translating to an
+// I/O device key.
+func (d *diskDiscovery) addPartitionRootFs(device, mountpoint string) bool {
+	fs, match := findIoDevice(filepath.Base(device), d.ctx.diskIoCounters)
+	if !match {
+		return false
+	}
+	// The resolved I/O device is already known here, so use it directly to avoid
+	// a second fallback search inside buildFsStatRegistration.
+	d.addFsStat(fs, mountpoint, true, "")
+	return true
+}
+
+// addLastResortRootFs is only used when neither FILESYSTEM nor partition-based
+// heuristics can identify root, so it picks the busiest I/O device as a final
+// fallback and preserves the root mountpoint for usage collection.
+func (d *diskDiscovery) addLastResortRootFs() {
+	rootKey := mostActiveIoDevice(d.ctx.diskIoCounters)
+	if rootKey != "" {
+		slog.Warn("Using most active device for root I/O; set FILESYSTEM to override", "device", rootKey)
+	} else {
+		rootKey = filepath.Base(d.rootMountPoint)
+		if _, exists := d.agent.fsStats[rootKey]; exists {
+			rootKey = "root"
+		}
+		slog.Warn("Root I/O device not detected; set FILESYSTEM to override")
+	}
+	d.agent.fsStats[rootKey] = &system.FsStats{Root: true, Mountpoint: d.rootMountPoint}
+}
+
+// findPartitionByFilesystemSetting matches an EXTRA_FILESYSTEMS entry against a
+// discovered partition either by mountpoint or by device suffix.
+func findPartitionByFilesystemSetting(filesystem string, partitions []disk.PartitionStat) (disk.PartitionStat, bool) {
+	for _, p := range partitions {
+		if strings.HasSuffix(p.Device, filesystem) || p.Mountpoint == filesystem {
+			return p, true
+		}
+	}
+	return disk.PartitionStat{}, false
+}
+
+// addConfiguredExtraFsEntry resolves one EXTRA_FILESYSTEMS entry, preferring a
+// discovered partition and falling back to any path that disk.Usage accepts.
+func (d *diskDiscovery) addConfiguredExtraFsEntry(filesystem, customName string) {
+	if p, found := findPartitionByFilesystemSetting(filesystem, d.partitions); found {
+		d.addFsStat(p.Device, p.Mountpoint, false, customName)
+		return
+	}
+
+	if _, err := d.usageFn(filesystem); err == nil {
+		d.addFsStat(filepath.Base(filesystem), filesystem, false, customName)
+		return
+	} else {
+		slog.Error("Invalid filesystem", "name", filesystem, "err", err)
+	}
+}
+
+// addConfiguredExtraFilesystems parses and registers the comma-separated
+// EXTRA_FILESYSTEMS env var entries.
+func (d *diskDiscovery) addConfiguredExtraFilesystems(extraFilesystems string) {
+	for fsEntry := range strings.SplitSeq(extraFilesystems, ",") {
+		filesystem, customName := parseFilesystemEntry(fsEntry)
+		d.addConfiguredExtraFsEntry(filesystem, customName)
+	}
+}
+
+// addPartitionExtraFs registers partitions mounted under /extra-filesystems so
+// their display names can come from the folder name while their I/O keys still
+// prefer the underlying partition device.
+func (d *diskDiscovery) addPartitionExtraFs(p disk.PartitionStat) {
+	if !strings.HasPrefix(p.Mountpoint, d.ctx.efPath) {
+		return
+	}
+	device, customName := extraFilesystemPartitionInfo(p)
+	d.addFsStat(device, p.Mountpoint, false, customName)
+}
+
+// addExtraFilesystemFolders handles bare directories under /extra-filesystems
+// that may not appear in partition discovery, while skipping mountpoints that
+// were already registered from higher-fidelity sources.
+func (d *diskDiscovery) addExtraFilesystemFolders(folderNames []string) {
+	existingMountpoints := make(map[string]bool, len(d.agent.fsStats))
+	for _, stats := range d.agent.fsStats {
+		existingMountpoints[stats.Mountpoint] = true
+	}
+
+	for _, folderName := range folderNames {
+		mountpoint := filepath.Join(d.ctx.efPath, folderName)
+		slog.Debug("/extra-filesystems", "mountpoint", mountpoint)
+		if existingMountpoints[mountpoint] {
+			continue
+		}
+		device, customName := parseFilesystemEntry(folderName)
+		d.addFsStat(device, mountpoint, false, customName)
+	}
 }
 
 // Sets up the filesystems to monitor for disk usage and I/O.
 func (a *Agent) initializeDiskInfo() {
 	filesystem, _ := utils.GetEnv("FILESYSTEM")
-	efPath := "/extra-filesystems"
 	hasRoot := false
 	isWindows := runtime.GOOS == "windows"
 
@@ -56,167 +286,57 @@ func (a *Agent) initializeDiskInfo() {
 		}
 	}
 
-	// ioContext := context.WithValue(a.sensorsContext,
-	// 	common.EnvKey, common.EnvMap{common.HostProcEnvKey: "/tmp/testproc"},
-	// )
-	// diskIoCounters, err := disk.IOCountersWithContext(ioContext)
-
 	diskIoCounters, err := disk.IOCounters()
 	if err != nil {
 		slog.Error("Error getting diskstats", "err", err)
 	}
 	slog.Debug("Disk I/O", "diskstats", diskIoCounters)
-
-	// Helper function to add a filesystem to fsStats if it doesn't exist
-	addFsStat := func(device, mountpoint string, root bool, customName ...string) {
-		var key string
-		if isWindows {
-			key = device
-		} else {
-			key = filepath.Base(device)
-		}
-		var ioMatch bool
-		if _, exists := a.fsStats[key]; !exists {
-			if root {
-				slog.Info("Detected root device", "name", key)
-				// Try to map root device to a diskIoCounters entry. First
-				// checks for an exact key match, then uses findIoDevice for
-				// normalized / prefix-based matching (e.g. nda0p2 → nda0),
-				// and finally falls back to the FILESYSTEM env var.
-				if _, ioMatch = diskIoCounters[key]; !ioMatch {
-					if matchedKey, match := findIoDevice(key, diskIoCounters); match {
-						key = matchedKey
-						ioMatch = true
-					} else if filesystem != "" {
-						if matchedKey, match := findIoDevice(filesystem, diskIoCounters); match {
-							key = matchedKey
-							ioMatch = true
-						}
-					}
-					if !ioMatch {
-						slog.Warn("Root I/O unmapped; set FILESYSTEM", "device", device, "mountpoint", mountpoint)
-					}
-				}
-			} else {
-				// Check if non-root has diskstats and fall back to folder name if not
-				// Scenario: device is encrypted and named luks-2bcb02be-999d-4417-8d18-5c61e660fb6e - not in /proc/diskstats.
-				// However, the device can be specified by mounting folder from luks device at /extra-filesystems/sda1
-				if _, ioMatch = diskIoCounters[key]; !ioMatch {
-					efBase := filepath.Base(mountpoint)
-					if _, ioMatch = diskIoCounters[efBase]; ioMatch {
-						key = efBase
-					}
-				}
-			}
-			fsStats := &system.FsStats{Root: root, Mountpoint: mountpoint}
-			if len(customName) > 0 && customName[0] != "" {
-				fsStats.Name = customName[0]
-			}
-			a.fsStats[key] = fsStats
-		}
+	ctx := fsRegistrationContext{
+		filesystem:     filesystem,
+		isWindows:      isWindows,
+		diskIoCounters: diskIoCounters,
+		efPath:         "/extra-filesystems",
 	}
 
 	// Get the appropriate root mount point for this system
-	rootMountPoint := a.getRootMountPoint()
-
-	// Use FILESYSTEM env var to find root filesystem
-	if filesystem != "" {
-		for _, p := range partitions {
-			if filesystemMatchesPartitionSetting(filesystem, p) {
-				addFsStat(p.Device, p.Mountpoint, true)
-				hasRoot = true
-				break
-			}
-		}
-		if !hasRoot {
-			// FILESYSTEM may name a physical disk absent from partitions (e.g.
-			// ZFS lists dataset paths like zroot/ROOT/default, not block devices).
-			// Try matching directly against diskIoCounters.
-			if ioKey, match := findIoDevice(filesystem, diskIoCounters); match {
-				a.fsStats[ioKey] = &system.FsStats{Root: true, Mountpoint: rootMountPoint}
-				hasRoot = true
-			} else {
-				slog.Warn("Partition details not found", "filesystem", filesystem)
-			}
-		}
+	discovery := diskDiscovery{
+		agent:          a,
+		rootMountPoint: a.getRootMountPoint(),
+		partitions:     partitions,
+		usageFn:        disk.Usage,
+		ctx:            ctx,
 	}
 
+	hasRoot = discovery.addConfiguredRootFs()
+
 	// Add EXTRA_FILESYSTEMS env var values to fsStats
 	if extraFilesystems, exists := utils.GetEnv("EXTRA_FILESYSTEMS"); exists {
-		for fsEntry := range strings.SplitSeq(extraFilesystems, ",") {
-			// Parse custom name from format: device__customname
-			fs, customName := parseFilesystemEntry(fsEntry)
-
-			found := false
-			for _, p := range partitions {
-				if strings.HasSuffix(p.Device, fs) || p.Mountpoint == fs {
-					addFsStat(p.Device, p.Mountpoint, false, customName)
-					found = true
-					break
-				}
-			}
-			// if not in partitions, test if we can get disk usage
-			if !found {
-				if _, err := disk.Usage(fs); err == nil {
-					addFsStat(filepath.Base(fs), fs, false, customName)
-				} else {
-					slog.Error("Invalid filesystem", "name", fs, "err", err)
-				}
-			}
-		}
+		discovery.addConfiguredExtraFilesystems(extraFilesystems)
 	}
 
 	// Process partitions for various mount points
 	for _, p := range partitions {
-		// fmt.Println(p.Device, p.Mountpoint)
-		// Binary root fallback or docker root fallback
-		if !hasRoot && (p.Mountpoint == rootMountPoint || (isDockerSpecialMountpoint(p.Mountpoint) && strings.HasPrefix(p.Device, "/dev"))) {
-			fs, match := findIoDevice(filepath.Base(p.Device), diskIoCounters)
-			if match {
-				addFsStat(fs, p.Mountpoint, true)
-				hasRoot = true
-			}
-		}
-
-		// Check if device is in /extra-filesystems
-		if strings.HasPrefix(p.Mountpoint, efPath) {
-			device, customName := parseFilesystemEntry(p.Mountpoint)
-			addFsStat(device, p.Mountpoint, false, customName)
+		if !hasRoot && isRootFallbackPartition(p, discovery.rootMountPoint) {
+			hasRoot = discovery.addPartitionRootFs(p.Device, p.Mountpoint)
 		}
+		discovery.addPartitionExtraFs(p)
 	}
 
 	// Check all folders in /extra-filesystems and add them if not already present
-	if folders, err := os.ReadDir(efPath); err == nil {
-		existingMountpoints := make(map[string]bool)
-		for _, stats := range a.fsStats {
-			existingMountpoints[stats.Mountpoint] = true
-		}
+	if folders, err := os.ReadDir(discovery.ctx.efPath); err == nil {
+		folderNames := make([]string, 0, len(folders))
 		for _, folder := range folders {
 			if folder.IsDir() {
-				mountpoint := filepath.Join(efPath, folder.Name())
-				slog.Debug("/extra-filesystems", "mountpoint", mountpoint)
-				if !existingMountpoints[mountpoint] {
-					device, customName := parseFilesystemEntry(folder.Name())
-					addFsStat(device, mountpoint, false, customName)
-				}
+				folderNames = append(folderNames, folder.Name())
 			}
 		}
+		discovery.addExtraFilesystemFolders(folderNames)
 	}
 
 	// If no root filesystem set, try the most active I/O device as a last
 	// resort (e.g. ZFS where dataset names are unrelated to disk names).
 	if !hasRoot {
-		rootKey := mostActiveIoDevice(diskIoCounters)
-		if rootKey != "" {
-			slog.Warn("Using most active device for root I/O; set FILESYSTEM to override", "device", rootKey)
-		} else {
-			rootKey = filepath.Base(rootMountPoint)
-			if _, exists := a.fsStats[rootKey]; exists {
-				rootKey = "root"
-			}
-			slog.Warn("Root I/O device not detected; set FILESYSTEM to override")
-		}
-		a.fsStats[rootKey] = &system.FsStats{Root: true, Mountpoint: rootMountPoint}
+		discovery.addLastResortRootFs()
 	}
 
 	a.pruneDuplicateRootExtraFilesystems()
@@ -381,6 +501,8 @@ func normalizeDeviceName(value string) string {
 
 // Sets start values for disk I/O stats.
 func (a *Agent) initializeDiskIoStats(diskIoCounters map[string]disk.IOCountersStat) {
+	a.fsNames = a.fsNames[:0]
+	now := time.Now()
 	for device, stats := range a.fsStats {
 		// skip if not in diskIoCounters
 		d, exists := diskIoCounters[device]
@@ -389,7 +511,7 @@ func (a *Agent) initializeDiskIoStats(diskIoCounters map[string]disk.IOCountersS
 			continue
 		}
 		// populate initial values
-		stats.Time = time.Now()
+		stats.Time = now
 		stats.TotalRead = d.ReadBytes
 		stats.TotalWrite = d.WriteBytes
 		// add to list of valid io device names

agent/disk_test.go

@@ -93,6 +93,443 @@ func TestParseFilesystemEntry(t *testing.T) {
 	}
 }
 
+func TestExtraFilesystemPartitionInfo(t *testing.T) {
+	t.Run("uses partition device for label-only mountpoint", func(t *testing.T) {
+		device, customName := extraFilesystemPartitionInfo(disk.PartitionStat{
+			Device:     "/dev/sdc",
+			Mountpoint: "/extra-filesystems/Share",
+		})
+
+		assert.Equal(t, "/dev/sdc", device)
+		assert.Equal(t, "", customName)
+	})
+
+	t.Run("uses custom name from mountpoint suffix", func(t *testing.T) {
+		device, customName := extraFilesystemPartitionInfo(disk.PartitionStat{
+			Device:     "/dev/sdc",
+			Mountpoint: "/extra-filesystems/sdc__Share",
+		})
+
+		assert.Equal(t, "/dev/sdc", device)
+		assert.Equal(t, "Share", customName)
+	})
+
+	t.Run("falls back to folder device when partition device is unavailable", func(t *testing.T) {
+		device, customName := extraFilesystemPartitionInfo(disk.PartitionStat{
+			Mountpoint: "/extra-filesystems/sdc__Share",
+		})
+
+		assert.Equal(t, "sdc", device)
+		assert.Equal(t, "Share", customName)
+	})
+
+	t.Run("supports custom name without folder device prefix", func(t *testing.T) {
+		device, customName := extraFilesystemPartitionInfo(disk.PartitionStat{
+			Device:     "/dev/sdc",
+			Mountpoint: "/extra-filesystems/__Share",
+		})
+
+		assert.Equal(t, "/dev/sdc", device)
+		assert.Equal(t, "Share", customName)
+	})
+}
+
+func TestBuildFsStatRegistration(t *testing.T) {
+	t.Run("uses basename for non-windows exact io match", func(t *testing.T) {
+		key, stats, ok := registerFilesystemStats(
+			map[string]*system.FsStats{},
+			"/dev/sda1",
+			"/mnt/data",
+			false,
+			"archive",
+			fsRegistrationContext{
+				isWindows: false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"sda1": {Name: "sda1"},
+				},
+			},
+		)
+
+		assert.True(t, ok)
+		assert.Equal(t, "sda1", key)
+		assert.Equal(t, "/mnt/data", stats.Mountpoint)
+		assert.Equal(t, "archive", stats.Name)
+		assert.False(t, stats.Root)
+	})
+
+	t.Run("maps root partition to io device by prefix", func(t *testing.T) {
+		key, stats, ok := registerFilesystemStats(
+			map[string]*system.FsStats{},
+			"/dev/ada0p2",
+			"/",
+			true,
+			"",
+			fsRegistrationContext{
+				isWindows: false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"ada0": {Name: "ada0", ReadBytes: 1000, WriteBytes: 1000},
+				},
+			},
+		)
+
+		assert.True(t, ok)
+		assert.Equal(t, "ada0", key)
+		assert.True(t, stats.Root)
+		assert.Equal(t, "/", stats.Mountpoint)
+	})
+
+	t.Run("uses filesystem setting as root fallback", func(t *testing.T) {
+		key, _, ok := registerFilesystemStats(
+			map[string]*system.FsStats{},
+			"overlay",
+			"/",
+			true,
+			"",
+			fsRegistrationContext{
+				filesystem: "nvme0n1p2",
+				isWindows:  false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"nvme0n1": {Name: "nvme0n1", ReadBytes: 1000, WriteBytes: 1000},
+				},
+			},
+		)
+
+		assert.True(t, ok)
+		assert.Equal(t, "nvme0n1", key)
+	})
+
+	t.Run("prefers parsed extra-filesystems device over mapper device", func(t *testing.T) {
+		key, stats, ok := registerFilesystemStats(
+			map[string]*system.FsStats{},
+			"/dev/mapper/luks-2bcb02be-999d-4417-8d18-5c61e660fb6e",
+			"/extra-filesystems/nvme0n1p2__Archive",
+			false,
+			"Archive",
+			fsRegistrationContext{
+				isWindows: false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"dm-1":      {Name: "dm-1", Label: "luks-2bcb02be-999d-4417-8d18-5c61e660fb6e"},
+					"nvme0n1p2": {Name: "nvme0n1p2"},
+				},
+			},
+		)
+
+		assert.True(t, ok)
+		assert.Equal(t, "nvme0n1p2", key)
+		assert.Equal(t, "Archive", stats.Name)
+	})
+
+	t.Run("falls back to mapper io device when folder device cannot be resolved", func(t *testing.T) {
+		key, stats, ok := registerFilesystemStats(
+			map[string]*system.FsStats{},
+			"/dev/mapper/luks-2bcb02be-999d-4417-8d18-5c61e660fb6e",
+			"/extra-filesystems/Archive",
+			false,
+			"Archive",
+			fsRegistrationContext{
+				isWindows: false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"dm-1": {Name: "dm-1", Label: "luks-2bcb02be-999d-4417-8d18-5c61e660fb6e"},
+				},
+			},
+		)
+
+		assert.True(t, ok)
+		assert.Equal(t, "dm-1", key)
+		assert.Equal(t, "Archive", stats.Name)
+	})
+
+	t.Run("uses full device name on windows", func(t *testing.T) {
+		key, _, ok := registerFilesystemStats(
+			map[string]*system.FsStats{},
+			`C:`,
+			`C:\\`,
+			false,
+			"",
+			fsRegistrationContext{
+				isWindows: true,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					`C:`: {Name: `C:`},
+				},
+			},
+		)
+
+		assert.True(t, ok)
+		assert.Equal(t, `C:`, key)
+	})
+
+	t.Run("skips existing key", func(t *testing.T) {
+		key, stats, ok := registerFilesystemStats(
+			map[string]*system.FsStats{"sda1": {Mountpoint: "/existing"}},
+			"/dev/sda1",
+			"/mnt/data",
+			false,
+			"",
+			fsRegistrationContext{
+				isWindows: false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"sda1": {Name: "sda1"},
+				},
+			},
+		)
+
+		assert.False(t, ok)
+		assert.Empty(t, key)
+		assert.Nil(t, stats)
+	})
+}
+
+func TestAddConfiguredRootFs(t *testing.T) {
+	t.Run("adds root from matching partition", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{
+			agent:          agent,
+			rootMountPoint: "/",
+			partitions:     []disk.PartitionStat{{Device: "/dev/ada0p2", Mountpoint: "/"}},
+			ctx: fsRegistrationContext{
+				filesystem: "/dev/ada0p2",
+				isWindows:  false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"ada0": {Name: "ada0", ReadBytes: 1000, WriteBytes: 1000},
+				},
+			},
+		}
+
+		ok := discovery.addConfiguredRootFs()
+
+		assert.True(t, ok)
+		stats, exists := agent.fsStats["ada0"]
+		assert.True(t, exists)
+		assert.True(t, stats.Root)
+		assert.Equal(t, "/", stats.Mountpoint)
+	})
+
+	t.Run("adds root from io device when partition is missing", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{
+			agent:          agent,
+			rootMountPoint: "/sysroot",
+			ctx: fsRegistrationContext{
+				filesystem: "zroot",
+				isWindows:  false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"nda0": {Name: "nda0", Label: "zroot", ReadBytes: 1000, WriteBytes: 1000},
+				},
+			},
+		}
+
+		ok := discovery.addConfiguredRootFs()
+
+		assert.True(t, ok)
+		stats, exists := agent.fsStats["nda0"]
+		assert.True(t, exists)
+		assert.True(t, stats.Root)
+		assert.Equal(t, "/sysroot", stats.Mountpoint)
+	})
+
+	t.Run("returns false when filesystem cannot be resolved", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{
+			agent:          agent,
+			rootMountPoint: "/",
+			ctx: fsRegistrationContext{
+				filesystem:     "missing-disk",
+				isWindows:      false,
+				diskIoCounters: map[string]disk.IOCountersStat{},
+			},
+		}
+
+		ok := discovery.addConfiguredRootFs()
+
+		assert.False(t, ok)
+		assert.Empty(t, agent.fsStats)
+	})
+}
+
+func TestAddPartitionRootFs(t *testing.T) {
+	t.Run("adds root from fallback partition candidate", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{
+			agent: agent,
+			ctx: fsRegistrationContext{
+				isWindows: false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"nvme0n1": {Name: "nvme0n1", ReadBytes: 1000, WriteBytes: 1000},
+				},
+			},
+		}
+
+		ok := discovery.addPartitionRootFs("/dev/nvme0n1p2", "/")
+
+		assert.True(t, ok)
+		stats, exists := agent.fsStats["nvme0n1"]
+		assert.True(t, exists)
+		assert.True(t, stats.Root)
+		assert.Equal(t, "/", stats.Mountpoint)
+	})
+
+	t.Run("returns false when no io device matches", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{agent: agent, ctx: fsRegistrationContext{diskIoCounters: map[string]disk.IOCountersStat{}}}
+
+		ok := discovery.addPartitionRootFs("/dev/mapper/root", "/")
+
+		assert.False(t, ok)
+		assert.Empty(t, agent.fsStats)
+	})
+}
+
+func TestAddLastResortRootFs(t *testing.T) {
+	t.Run("uses most active io device when available", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{agent: agent, rootMountPoint: "/", ctx: fsRegistrationContext{diskIoCounters: map[string]disk.IOCountersStat{
+			"sda": {Name: "sda", ReadBytes: 5000, WriteBytes: 5000},
+			"sdb": {Name: "sdb", ReadBytes: 1000, WriteBytes: 1000},
+		}}}
+
+		discovery.addLastResortRootFs()
+
+		stats, exists := agent.fsStats["sda"]
+		assert.True(t, exists)
+		assert.True(t, stats.Root)
+	})
+
+	t.Run("falls back to root key when mountpoint basename collides", func(t *testing.T) {
+		agent := &Agent{fsStats: map[string]*system.FsStats{
+			"sysroot": {Mountpoint: "/extra-filesystems/sysroot"},
+		}}
+		discovery := diskDiscovery{agent: agent, rootMountPoint: "/sysroot", ctx: fsRegistrationContext{diskIoCounters: map[string]disk.IOCountersStat{}}}
+
+		discovery.addLastResortRootFs()
+
+		stats, exists := agent.fsStats["root"]
+		assert.True(t, exists)
+		assert.True(t, stats.Root)
+		assert.Equal(t, "/sysroot", stats.Mountpoint)
+	})
+}
+
+func TestAddConfiguredExtraFsEntry(t *testing.T) {
+	t.Run("uses matching partition when present", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{
+			agent:      agent,
+			partitions: []disk.PartitionStat{{Device: "/dev/sdb1", Mountpoint: "/mnt/backup"}},
+			usageFn: func(string) (*disk.UsageStat, error) {
+				t.Fatal("usage fallback should not be called when partition matches")
+				return nil, nil
+			},
+			ctx: fsRegistrationContext{
+				isWindows: false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"sdb1": {Name: "sdb1"},
+				},
+			},
+		}
+
+		discovery.addConfiguredExtraFsEntry("sdb1", "backup")
+
+		stats, exists := agent.fsStats["sdb1"]
+		assert.True(t, exists)
+		assert.Equal(t, "/mnt/backup", stats.Mountpoint)
+		assert.Equal(t, "backup", stats.Name)
+	})
+
+	t.Run("falls back to usage-validated path", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{
+			agent: agent,
+			usageFn: func(path string) (*disk.UsageStat, error) {
+				assert.Equal(t, "/srv/archive", path)
+				return &disk.UsageStat{}, nil
+			},
+			ctx: fsRegistrationContext{
+				isWindows: false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"archive": {Name: "archive"},
+				},
+			},
+		}
+
+		discovery.addConfiguredExtraFsEntry("/srv/archive", "archive")
+
+		stats, exists := agent.fsStats["archive"]
+		assert.True(t, exists)
+		assert.Equal(t, "/srv/archive", stats.Mountpoint)
+		assert.Equal(t, "archive", stats.Name)
+	})
+
+	t.Run("ignores invalid filesystem entry", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{
+			agent: agent,
+			usageFn: func(string) (*disk.UsageStat, error) {
+				return nil, os.ErrNotExist
+			},
+		}
+
+		discovery.addConfiguredExtraFsEntry("/missing/archive", "")
+
+		assert.Empty(t, agent.fsStats)
+	})
+}
+
+func TestAddConfiguredExtraFilesystems(t *testing.T) {
+	t.Run("parses and registers multiple configured filesystems", func(t *testing.T) {
+		agent := &Agent{fsStats: make(map[string]*system.FsStats)}
+		discovery := diskDiscovery{
+			agent:      agent,
+			partitions: []disk.PartitionStat{{Device: "/dev/sda1", Mountpoint: "/mnt/fast"}},
+			usageFn: func(path string) (*disk.UsageStat, error) {
+				if path == "/srv/archive" {
+					return &disk.UsageStat{}, nil
+				}
+				return nil, os.ErrNotExist
+			},
+			ctx: fsRegistrationContext{
+				isWindows: false,
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"sda1":    {Name: "sda1"},
+					"archive": {Name: "archive"},
+				},
+			},
+		}
+
+		discovery.addConfiguredExtraFilesystems("sda1__fast,/srv/archive__cold")
+
+		assert.Contains(t, agent.fsStats, "sda1")
+		assert.Equal(t, "fast", agent.fsStats["sda1"].Name)
+		assert.Contains(t, agent.fsStats, "archive")
+		assert.Equal(t, "cold", agent.fsStats["archive"].Name)
+	})
+}
+
+func TestAddExtraFilesystemFolders(t *testing.T) {
+	t.Run("adds missing folders and skips existing mountpoints", func(t *testing.T) {
+		agent := &Agent{fsStats: map[string]*system.FsStats{
+			"existing": {Mountpoint: "/extra-filesystems/existing"},
+		}}
+		discovery := diskDiscovery{
+			agent: agent,
+			ctx: fsRegistrationContext{
+				isWindows: false,
+				efPath:    "/extra-filesystems",
+				diskIoCounters: map[string]disk.IOCountersStat{
+					"newdisk": {Name: "newdisk"},
+				},
+			},
+		}
+
+		discovery.addExtraFilesystemFolders([]string{"existing", "newdisk__Archive"})
+
+		assert.Len(t, agent.fsStats, 2)
+		stats, exists := agent.fsStats["newdisk"]
+		assert.True(t, exists)
+		assert.Equal(t, "/extra-filesystems/newdisk__Archive", stats.Mountpoint)
+		assert.Equal(t, "Archive", stats.Name)
+	})
+}
+
 func TestFindIoDevice(t *testing.T) {
 	t.Run("matches by device name", func(t *testing.T) {
 		ioCounters := map[string]disk.IOCountersStat{
@@ -310,7 +747,7 @@ func TestInitializeDiskInfoWithCustomNames(t *testing.T) {
 			// Test the parsing logic by calling the relevant part
 			// We'll create a simplified version to test just the parsing
 			extraFilesystems := tc.envValue
-			for _, fsEntry := range strings.Split(extraFilesystems, ",") {
+			for fsEntry := range strings.SplitSeq(extraFilesystems, ",") {
 				// Parse the entry
 				fsEntry = strings.TrimSpace(fsEntry)
 				var fs, customName string
@@ -506,3 +943,33 @@ func TestHasSameDiskUsage(t *testing.T) {
 		assert.False(t, hasSameDiskUsage(&disk.UsageStat{Total: 0, Used: 0}, &disk.UsageStat{Total: 1, Used: 1}))
 	})
 }
+
+func TestInitializeDiskIoStatsResetsTrackedDevices(t *testing.T) {
+	agent := &Agent{
+		fsStats: map[string]*system.FsStats{
+			"sda": {},
+			"sdb": {},
+		},
+		fsNames: []string{"stale", "sda"},
+	}
+
+	agent.initializeDiskIoStats(map[string]disk.IOCountersStat{
+		"sda": {Name: "sda", ReadBytes: 10, WriteBytes: 20},
+		"sdb": {Name: "sdb", ReadBytes: 30, WriteBytes: 40},
+	})
+
+	assert.ElementsMatch(t, []string{"sda", "sdb"}, agent.fsNames)
+	assert.Len(t, agent.fsNames, 2)
+	assert.Equal(t, uint64(10), agent.fsStats["sda"].TotalRead)
+	assert.Equal(t, uint64(20), agent.fsStats["sda"].TotalWrite)
+	assert.False(t, agent.fsStats["sda"].Time.IsZero())
+	assert.False(t, agent.fsStats["sdb"].Time.IsZero())
+
+	agent.initializeDiskIoStats(map[string]disk.IOCountersStat{
+		"sdb": {Name: "sdb", ReadBytes: 50, WriteBytes: 60},
+	})
+
+	assert.Equal(t, []string{"sdb"}, agent.fsNames)
+	assert.Equal(t, uint64(50), agent.fsStats["sdb"].TotalRead)
+	assert.Equal(t, uint64(60), agent.fsStats["sdb"].TotalWrite)
+}