Skip to content

Commit 129fbdf

Browse files
Set 'HostProcMountinfo' for Disk integration (#37826)
1 parent 4e1c6c5 commit 129fbdf

File tree

6 files changed

+63
-40
lines changed

6 files changed

+63
-40
lines changed

cmd/agent/dist/conf.d/disk.d/conf.yaml.default

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -241,3 +241,12 @@ instances:
241241
# host: nfsserver
242242
# share: /mnt/nfs_share
243243
# type: nfs
244+
245+
## @param proc_mountinfo_path - string - optional
246+
## Path to the file from which to read mount information when enumerating disk partitions.
247+
## By default this points at /proc/self/mounts (the same source used by Python’s disk check).
248+
## If this setting is empty or does not exist, Go’s disk check built-in fallback
249+
## logic will be used in order: /proc/1/mountinfo → /proc/self/mountinfo → /proc/mounts.
250+
#
251+
# proc_mountinfo_path: /proc/self/mounts
252+

pkg/collector/corechecks/system/disk/diskv2/disk.go

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,15 @@
77
package diskv2
88

99
import (
10+
"context"
1011
"errors"
1112
"fmt"
1213
"regexp"
1314
"strings"
1415
"time"
1516

1617
"github.com/benbjohnson/clock"
18+
"github.com/shirou/gopsutil/v4/common"
1719
gopsutil_disk "github.com/shirou/gopsutil/v4/disk"
1820
yaml "gopkg.in/yaml.v2"
1921

@@ -86,6 +88,7 @@ type diskInstanceConfig struct {
8688
DeviceTagRe map[string]string `yaml:"device_tag_re"`
8789
LowercaseDeviceTag bool `yaml:"lowercase_device_tag"`
8890
Timeout uint16 `yaml:"timeout"`
91+
ProcMountInfoPath string `yaml:"proc_mountinfo_path"`
8992
}
9093

9194
func sliceMatchesExpression(slice []regexp.Regexp, expression string) bool {
@@ -111,10 +114,10 @@ func compileRegExp(expr string, ignoreCase bool) (*regexp.Regexp, error) {
111114
// Check represents the Disk check that will be periodically executed via the Run() function
112115
type Check struct {
113116
core.CheckBase
114-
clock clock.Clock
115-
diskPartitions func(bool) ([]gopsutil_disk.PartitionStat, error)
116-
diskUsage func(string) (*gopsutil_disk.UsageStat, error)
117-
diskIOCounters func(...string) (map[string]gopsutil_disk.IOCountersStat, error)
117+
clock clock.Clock
118+
diskPartitionsWithContext func(context.Context, bool) ([]gopsutil_disk.PartitionStat, error)
119+
diskUsage func(string) (*gopsutil_disk.UsageStat, error)
120+
diskIOCounters func(...string) (map[string]gopsutil_disk.IOCountersStat, error)
118121

119122
initConfig diskInitConfig
120123
instanceConfig diskInstanceConfig
@@ -410,7 +413,11 @@ func (c *Check) configureIncludeMountPoint() error {
410413
}
411414

412415
func (c *Check) collectPartitionMetrics(sender sender.Sender) error {
413-
partitions, err := c.diskPartitions(c.instanceConfig.IncludeAllDevices)
416+
ctx := context.Background()
417+
if c.instanceConfig.ProcMountInfoPath != "" {
418+
ctx = context.WithValue(ctx, common.EnvKey, common.EnvMap{common.HostProcMountinfo: c.instanceConfig.ProcMountInfoPath})
419+
}
420+
partitions, err := c.diskPartitionsWithContext(ctx, c.instanceConfig.IncludeAllDevices)
414421
if err != nil {
415422
log.Warnf("Unable to get disk partitions: %s", err)
416423
return err
@@ -671,11 +678,11 @@ func Factory() option.Option[func() check.Check] {
671678

672679
func newCheck() check.Check {
673680
return &Check{
674-
CheckBase: core.NewCheckBase(CheckName),
675-
clock: clock.New(),
676-
diskPartitions: gopsutil_disk.Partitions,
677-
diskUsage: gopsutil_disk.Usage,
678-
diskIOCounters: gopsutil_disk.IOCounters,
681+
CheckBase: core.NewCheckBase(CheckName),
682+
clock: clock.New(),
683+
diskPartitionsWithContext: gopsutil_disk.PartitionsWithContext,
684+
diskUsage: gopsutil_disk.Usage,
685+
diskIOCounters: gopsutil_disk.IOCounters,
679686
initConfig: diskInitConfig{
680687
DeviceGlobalExclude: []string{},
681688
DeviceGlobalBlacklist: []string{},
@@ -714,6 +721,8 @@ func newCheck() check.Check {
714721
DeviceTagRe: make(map[string]string),
715722
LowercaseDeviceTag: false,
716723
Timeout: 5,
724+
// Match psutil exactly setting default value (https://github.com/giampaolo/psutil/blob/3d21a43a47ab6f3c4a08d235d2a9a55d4adae9b1/psutil/_pslinux.py#L1277)
725+
ProcMountInfoPath: "/proc/self/mounts",
717726
},
718727
includedDevices: []regexp.Regexp{},
719728
excludedDevices: []regexp.Regexp{},

pkg/collector/corechecks/system/disk/diskv2/disk_nix_test.go

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
package diskv2_test
99

1010
import (
11+
"context"
1112
"errors"
1213
"testing"
1314

@@ -82,7 +83,7 @@ func TestGivenADiskCheckWithDefaultConfig_WhenCheckRuns_ThenAllUsageMetricsAreRe
8283
func TestGivenADiskCheckWithLowercaseDeviceTagConfigured_WhenCheckRuns_ThenLowercaseDevicesAreReported(t *testing.T) {
8384
setupDefaultMocks()
8485
diskCheck := createDiskCheck(t)
85-
diskCheck = diskv2.WithDiskPartitions(diskv2.WithDiskIOCounters(diskCheck, func(...string) (map[string]gopsutil_disk.IOCountersStat, error) {
86+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskv2.WithDiskIOCounters(diskCheck, func(...string) (map[string]gopsutil_disk.IOCountersStat, error) {
8687
return map[string]gopsutil_disk.IOCountersStat{
8788
"/dev/SDA1": {
8889
Name: "sda1",
@@ -103,7 +104,7 @@ func TestGivenADiskCheckWithLowercaseDeviceTagConfigured_WhenCheckRuns_ThenLower
103104
WriteTime: 150,
104105
},
105106
}, nil
106-
}), func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
107+
}), func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
107108
return []gopsutil_disk.PartitionStat{
108109
{
109110
Device: "/dev/SDA1",
@@ -183,7 +184,7 @@ func TestGivenADiskCheckWithIncludeAllDevicesFalseConfigured_WhenCheckRuns_ThenO
183184
func TestGivenADiskCheckWithDefaultConfig_WhenCheckRunsAndPartitionsSystemReturnsEmptyDevice_ThenNoUsageMetricsAreReportedForThatPartition(t *testing.T) {
184185
setupDefaultMocks()
185186
diskCheck := createDiskCheck(t)
186-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
187+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
187188
return []gopsutil_disk.PartitionStat{
188189
{
189190
Device: "",
@@ -217,7 +218,7 @@ func TestGivenADiskCheckWithDefaultConfig_WhenCheckRunsAndPartitionsSystemReturn
217218
func TestGivenADiskCheckWithAllPartitionsFalseConfigured_WhenCheckRunsAndPartitionsSystemReturnsEmptyDevice_ThenNoUsageMetricsAreReportedForThatPartition(t *testing.T) {
218219
setupDefaultMocks()
219220
diskCheck := createDiskCheck(t)
220-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
221+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
221222
return []gopsutil_disk.PartitionStat{
222223
{
223224
Device: "",
@@ -252,7 +253,7 @@ func TestGivenADiskCheckWithAllPartitionsFalseConfigured_WhenCheckRunsAndPartiti
252253
func TestGivenADiskCheckWithAllPartitionsTrueConfigured_WhenCheckRunsAndPartitionsSystemReturnsEmptyDevice_ThenUsageMetricsAreReportedForThatPartition(t *testing.T) {
253254
setupDefaultMocks()
254255
diskCheck := createDiskCheck(t)
255-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
256+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
256257
return []gopsutil_disk.PartitionStat{
257258
{
258259
Device: "",
@@ -343,7 +344,7 @@ device_whitelist:
343344
func TestGivenADiskCheckWithFileSystemGlobalExcludeNotConfigured_WhenCheckRuns_ThenUsageMetricsAreNotReportedForPartitionsWithIso9660FileSystems(t *testing.T) {
344345
setupDefaultMocks()
345346
diskCheck := createDiskCheck(t)
346-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
347+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
347348
return []gopsutil_disk.PartitionStat{
348349
{
349350
Device: "cdrom",
@@ -376,7 +377,7 @@ func TestGivenADiskCheckWithFileSystemGlobalExcludeNotConfigured_WhenCheckRuns_T
376377
func TestGivenADiskCheckWithFileSystemGlobalExcludeNotConfigured_WhenCheckRuns_ThenUsageMetricsAreNotReportedForPartitionsWithTracefsFileSystems(t *testing.T) {
377378
setupDefaultMocks()
378379
diskCheck := createDiskCheck(t)
379-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
380+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
380381
return []gopsutil_disk.PartitionStat{
381382
{
382383
Device: "trace",
@@ -500,7 +501,7 @@ excluded_filesystems:
500501
func TestGivenADiskCheckWithExcludedFileSystemsConfiguredWithTmpfs_WhenCheckRuns_ThenUsageMetricsAreReportedForPartitionsWithDevTmpfsFileSystem(t *testing.T) {
501502
setupDefaultMocks()
502503
diskCheck := createDiskCheck(t)
503-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
504+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
504505
return []gopsutil_disk.PartitionStat{
505506
{
506507
Device: "devtmpfs",

pkg/collector/corechecks/system/disk/diskv2/disk_test.go

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ package diskv2_test
88
import (
99
"bufio"
1010
"bytes"
11+
"context"
1112
"errors"
1213
"testing"
1314
"time"
@@ -169,11 +170,11 @@ func createDiskCheck(t *testing.T) check.Check {
169170
diskCheckOpt := diskv2.Factory()
170171
diskCheckFunc, _ := diskCheckOpt.Get()
171172
diskCheck := diskCheckFunc()
172-
diskCheck = diskv2.WithDiskPartitions(diskv2.WithDiskUsage(diskv2.WithDiskIOCounters(diskCheck, func(...string) (map[string]gopsutil_disk.IOCountersStat, error) {
173+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskv2.WithDiskUsage(diskv2.WithDiskIOCounters(diskCheck, func(...string) (map[string]gopsutil_disk.IOCountersStat, error) {
173174
return ioCountersData, nil
174175
}), func(mountpoint string) (*gopsutil_disk.UsageStat, error) {
175176
return usageData[mountpoint], nil
176-
}), func(all bool) ([]gopsutil_disk.PartitionStat, error) {
177+
}), func(_ context.Context, all bool) ([]gopsutil_disk.PartitionStat, error) {
177178
if all {
178179
return partitionsTrue, nil
179180
}
@@ -221,7 +222,7 @@ func TestGivenADiskCheckAndStoppedSender(t *testing.T) {
221222
func TestGivenADiskCheckWithDefaultConfig_WhenCheckRunsAndPartitionsSystemCallReturnsError_ThenErrorIsReturnedAndNoUsageMetricsAreReported(t *testing.T) {
222223
setupDefaultMocks()
223224
diskCheck := createDiskCheck(t)
224-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
225+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
225226
return nil, errors.New("error calling disk.DiskPartitions")
226227
})
227228
m := mocksender.NewMockSender(diskCheck.ID())
@@ -725,7 +726,7 @@ excluded_disks:
725726
func TestGivenADiskCheckWithExcludedDisksConfiguredWithDa2_WhenCheckRuns_ThenUsageMetricsAreNotReportedForPartitionsWithSda2Devices(t *testing.T) {
726727
setupDefaultMocks()
727728
diskCheck := createDiskCheck(t)
728-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
729+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
729730
return []gopsutil_disk.PartitionStat{
730731
{
731732
Device: "sda2",
@@ -865,7 +866,7 @@ file_system_include:
865866
func TestGivenADiskCheckWithMountPointGlobalExcludeNotConfigured_WhenCheckRuns_ThenUsageMetricsAreNotReportedForPartitionsWithBinfmt_miscMountPoints(t *testing.T) {
866867
setupDefaultMocks()
867868
diskCheck := createDiskCheck(t)
868-
diskCheck = diskv2.WithDiskPartitions(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
869+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
869870
return &gopsutil_disk.UsageStat{
870871
Path: "/host/proc/sys/fs/binfmt_misc",
871872
Fstype: "ext4",
@@ -878,7 +879,7 @@ func TestGivenADiskCheckWithMountPointGlobalExcludeNotConfigured_WhenCheckRuns_T
878879
InodesFree: 500000,
879880
InodesUsedPercent: 50.0,
880881
}, nil
881-
}), func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
882+
}), func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
882883
return []gopsutil_disk.PartitionStat{
883884
{
884885
Device: "first",
@@ -1024,7 +1025,7 @@ mount_point_include:
10241025
func TestGivenADiskCheckWithDefaultConfig_WhenCheckRunsAndUsageSystemCallReturnsAPartitionWithZeroTotal_ThenNoUsageMetricsAreReportedForThatPartition(t *testing.T) {
10251026
setupDefaultMocks()
10261027
diskCheck := createDiskCheck(t)
1027-
diskCheck = diskv2.WithDiskPartitions(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
1028+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
10281029
return &gopsutil_disk.UsageStat{
10291030
Path: "/dev/shm",
10301031
Fstype: "tmpfs",
@@ -1037,7 +1038,7 @@ func TestGivenADiskCheckWithDefaultConfig_WhenCheckRunsAndUsageSystemCallReturns
10371038
InodesFree: 0,
10381039
InodesUsedPercent: 0,
10391040
}, nil
1040-
}), func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
1041+
}), func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
10411042
return []gopsutil_disk.PartitionStat{
10421043
{
10431044
Device: "shm",
@@ -1059,7 +1060,7 @@ func TestGivenADiskCheckWithDefaultConfig_WhenCheckRunsAndUsageSystemCallReturns
10591060
func TestGivenADiskCheckWithMinDiskSizeConfiguredTo1MiBConfig_WhenCheckRunsAndUsageSystemCallReturnsAPartitionWith1024Total_ThenNoUsageMetricsAreReportedForThatPartition(t *testing.T) {
10601061
setupDefaultMocks()
10611062
diskCheck := createDiskCheck(t)
1062-
diskCheck = diskv2.WithDiskPartitions(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
1063+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
10631064
return &gopsutil_disk.UsageStat{
10641065
Path: "/dev/shm",
10651066
Fstype: "tmpfs",
@@ -1072,7 +1073,7 @@ func TestGivenADiskCheckWithMinDiskSizeConfiguredTo1MiBConfig_WhenCheckRunsAndUs
10721073
InodesFree: 0,
10731074
InodesUsedPercent: 0,
10741075
}, nil
1075-
}), func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
1076+
}), func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
10761077
return []gopsutil_disk.PartitionStat{
10771078
{
10781079
Device: "shm",
@@ -1104,7 +1105,7 @@ func TestGivenADiskCheckWithMinDiskSizeConfiguredTo1MiBConfig_WhenCheckRunsAndUs
11041105
func TestGivenADiskCheckWithMinDiskSizeConfiguredTo1MiBConfig_WhenCheckRunsAndUsageSystemCallReturnsAPartitionWith1048576Total_ThenUsageMetricsAreReportedForThatPartition(t *testing.T) {
11051106
setupDefaultMocks()
11061107
diskCheck := createDiskCheck(t)
1107-
diskCheck = diskv2.WithDiskPartitions(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
1108+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
11081109
return &gopsutil_disk.UsageStat{
11091110
Path: "/dev/shm",
11101111
Fstype: "tmpfs",
@@ -1117,7 +1118,7 @@ func TestGivenADiskCheckWithMinDiskSizeConfiguredTo1MiBConfig_WhenCheckRunsAndUs
11171118
InodesFree: 0,
11181119
InodesUsedPercent: 0,
11191120
}, nil
1120-
}), func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
1121+
}), func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
11211122
return []gopsutil_disk.PartitionStat{
11221123
{
11231124
Device: "shm",
@@ -1307,7 +1308,7 @@ func TestGivenADiskCheckWithDefaultConfig_WhenUsagePartitionTimeout_ThenUsageMet
13071308
afterCalled: afterCalled,
13081309
}
13091310
diskCheck := createDiskCheck(t)
1310-
diskCheck = diskv2.WithClock(diskv2.WithDiskPartitions(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
1311+
diskCheck = diskv2.WithClock(diskv2.WithDiskPartitionsWithContext(diskv2.WithDiskUsage(diskCheck, func(_ string) (*gopsutil_disk.UsageStat, error) {
13111312
// Sleep 10s (longer than default timeout)
13121313
time.Sleep(10 * time.Second)
13131314
return &gopsutil_disk.UsageStat{
@@ -1322,7 +1323,7 @@ func TestGivenADiskCheckWithDefaultConfig_WhenUsagePartitionTimeout_ThenUsageMet
13221323
InodesFree: 0,
13231324
InodesUsedPercent: 0,
13241325
}, nil
1325-
}), func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
1326+
}), func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
13261327
return []gopsutil_disk.PartitionStat{
13271328
{
13281329
Device: "shm",

pkg/collector/corechecks/system/disk/diskv2/disk_windows_test.go

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ package diskv2_test
1010
import (
1111
"bufio"
1212
"bytes"
13+
"context"
1314
"errors"
1415
"testing"
1516

@@ -40,7 +41,7 @@ func createWindowsCheck(t *testing.T) check.Check {
4041
diskCheckOpt := diskv2.Factory()
4142
diskCheckFunc, _ := diskCheckOpt.Get()
4243
diskCheck := diskCheckFunc()
43-
diskCheck = diskv2.WithDiskPartitions(diskv2.WithDiskUsage(diskv2.WithDiskIOCounters(diskCheck, func(...string) (map[string]gopsutil_disk.IOCountersStat, error) {
44+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskv2.WithDiskUsage(diskv2.WithDiskIOCounters(diskCheck, func(...string) (map[string]gopsutil_disk.IOCountersStat, error) {
4445
return map[string]gopsutil_disk.IOCountersStat{
4546
"\\\\?\\Volume{a1b2c3d4-e5f6-7890-abcd-ef1234567890}\\": {
4647
Name: "sda1",
@@ -65,7 +66,7 @@ func createWindowsCheck(t *testing.T) check.Check {
6566
InodesFree: 500000,
6667
InodesUsedPercent: 50.0,
6768
}, nil
68-
}), func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
69+
}), func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
6970
return []gopsutil_disk.PartitionStat{
7071
{
7172
Device: `\\?\Volume{a1b2c3d4-e5f6-7890-abcd-ef1234567890}\`,
@@ -147,7 +148,7 @@ func TestGivenADiskCheckWithIncludeAllDevicesFalseConfigured_WhenCheckRuns_ThenO
147148
func TestGivenADiskCheckWithDefaultConfig_WhenCheckRunsAndPartitionsSystemReturnsEmptyDevice_ThenNoUsageMetricsAreReportedForThatPartition(t *testing.T) {
148149
setupDefaultMocks()
149150
diskCheck := createWindowsCheck(t)
150-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
151+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
151152
return []gopsutil_disk.PartitionStat{
152153
{
153154
Device: "",
@@ -218,7 +219,7 @@ device_whitelist:
218219
func TestGivenADiskCheckWithFileSystemGlobalExcludeNotConfigured_WhenCheckRuns_ThenUsageMetricsAreNotReportedForPartitionsWithIso9660FileSystems(t *testing.T) {
219220
setupDefaultMocks()
220221
diskCheck := createWindowsCheck(t)
221-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
222+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
222223
return []gopsutil_disk.PartitionStat{
223224
{
224225
Device: "cdrom",
@@ -251,7 +252,7 @@ func TestGivenADiskCheckWithFileSystemGlobalExcludeNotConfigured_WhenCheckRuns_T
251252
func TestGivenADiskCheckWithFileSystemGlobalExcludeNotConfigured_WhenCheckRuns_ThenUsageMetricsAreNotReportedForPartitionsWithTracefsFileSystems(t *testing.T) {
252253
setupDefaultMocks()
253254
diskCheck := createWindowsCheck(t)
254-
diskCheck = diskv2.WithDiskPartitions(diskCheck, func(_ bool) ([]gopsutil_disk.PartitionStat, error) {
255+
diskCheck = diskv2.WithDiskPartitionsWithContext(diskCheck, func(_ context.Context, _ bool) ([]gopsutil_disk.PartitionStat, error) {
255256
return []gopsutil_disk.PartitionStat{
256257
{
257258
Device: "trace",

pkg/collector/corechecks/system/disk/diskv2/with_test.go

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77
package diskv2
88

99
import (
10+
"context"
11+
1012
"github.com/DataDog/datadog-agent/pkg/collector/check"
1113
"github.com/benbjohnson/clock"
1214
gopsutil_disk "github.com/shirou/gopsutil/v4/disk"
@@ -20,9 +22,9 @@ func WithClock(c check.Check, clock clock.Clock) check.Check {
2022
return c
2123
}
2224

23-
// WithDiskPartitions sets a diskPartitions call on the Check and returns the updated Check.
24-
func WithDiskPartitions(c check.Check, f func(bool) ([]gopsutil_disk.PartitionStat, error)) check.Check {
25-
c.(*Check).diskPartitions = f
25+
// WithDiskPartitionsWithContext sets a diskPartitionsWithContext call on the Check and returns the updated Check.
26+
func WithDiskPartitionsWithContext(c check.Check, f func(context.Context, bool) ([]gopsutil_disk.PartitionStat, error)) check.Check {
27+
c.(*Check).diskPartitionsWithContext = f
2628
return c
2729
}
2830

0 commit comments

Comments
 (0)