From 9644b7c7864d76f1c74cee34f12e99a882ff4b41 Mon Sep 17 00:00:00 2001
From: Gustavo Caso <gustavo.caso@datadoghq.com>
Date: Fri, 11 Apr 2025 13:23:20 +0200
Subject: [PATCH] Use the context provide to UsageWithContext public function

---
 disk/disk.go           | 25 +++++++++++++++++++++++++
 disk/disk_aix_cgo.go   |  2 +-
 disk/disk_aix_nocgo.go |  2 +-
 disk/disk_fallback.go  |  2 +-
 disk/disk_netbsd.go    |  2 +-
 disk/disk_openbsd.go   |  2 +-
 disk/disk_solaris.go   |  2 +-
 disk/disk_unix.go      |  2 +-
 disk/disk_windows.go   |  2 +-
 9 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/disk/disk.go b/disk/disk.go
index 310ea04..4861a6a 100644
--- a/disk/disk.go
+++ b/disk/disk.go
@@ -71,6 +71,31 @@ func Usage(path string) (*UsageStat, error) {
 	return UsageWithContext(context.Background(), path)
 }
 
+func UsageWithContext(ctx context.Context, path string) (*UsageStat, error) {
+	type result struct {
+		usage *UsageStat
+		err   error
+	}
+
+	resultCh := make(chan result, 1)
+
+	go func() {
+		usage, err := getUsage(ctx, path)
+		resultCh <- result{usage: usage, err: err}
+	}()
+
+	select {
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case res := <-resultCh:
+		if res.err != nil {
+			return nil, res.err
+		}
+
+		return res.usage, nil
+	}
+}
+
 // Partitions returns disk partitions. If all is false, returns
 // physical devices only (e.g. hard disks, cd-rom drives, USB keys)
 // and ignore all others (e.g. memory partitions such as /dev/shm)
diff --git a/disk/disk_aix_cgo.go b/disk/disk_aix_cgo.go
index a0d0829..2ea5cda 100644
--- a/disk/disk_aix_cgo.go
+++ b/disk/disk_aix_cgo.go
@@ -44,7 +44,7 @@ func PartitionsWithContext(ctx context.Context, all bool) ([]PartitionStat, erro
 	return ret, err
 }
 
-func UsageWithContext(ctx context.Context, path string) (*UsageStat, error) {
+func getUsage(_ context.Context, path string) (*UsageStat, error) {
 	f, err := perfstat.FileSystemStat()
 	if err != nil {
 		return nil, err
diff --git a/disk/disk_aix_nocgo.go b/disk/disk_aix_nocgo.go
index a0d621e..c208a80 100644
--- a/disk/disk_aix_nocgo.go
+++ b/disk/disk_aix_nocgo.go
@@ -84,7 +84,7 @@ func getFsType(stat unix.Statfs_t) string {
 	return FSType[int(stat.Vfstype)]
 }
 
-func UsageWithContext(ctx context.Context, path string) (*UsageStat, error) {
+func getUsage(ctx context.Context, path string) (*UsageStat, error) {
 	out, err := invoke.CommandWithContext(ctx, "df", "-v")
 	if err != nil {
 		return nil, err
diff --git a/disk/disk_fallback.go b/disk/disk_fallback.go
index 3b58ceb..8eecd5a 100644
--- a/disk/disk_fallback.go
+++ b/disk/disk_fallback.go
@@ -17,7 +17,7 @@ func PartitionsWithContext(_ context.Context, _ bool) ([]PartitionStat, error) {
 	return []PartitionStat{}, common.ErrNotImplementedError
 }
 
-func UsageWithContext(_ context.Context, _ string) (*UsageStat, error) {
+func getUsage(_ context.Context, _ string) (*UsageStat, error) {
 	return nil, common.ErrNotImplementedError
 }
 
diff --git a/disk/disk_netbsd.go b/disk/disk_netbsd.go
index 9d9e56f..1bc1db7 100644
--- a/disk/disk_netbsd.go
+++ b/disk/disk_netbsd.go
@@ -102,7 +102,7 @@ func IOCountersWithContext(_ context.Context, _ ...string) (map[string]IOCounter
 	return ret, common.ErrNotImplementedError
 }
 
-func UsageWithContext(_ context.Context, path string) (*UsageStat, error) {
+func getUsage(_ context.Context, path string) (*UsageStat, error) {
 	stat := Statvfs{}
 	flag := uint64(1) // ST_WAIT/MNT_WAIT, see sys/fstypes.h
 
diff --git a/disk/disk_openbsd.go b/disk/disk_openbsd.go
index 2b0ce10..dd71863 100644
--- a/disk/disk_openbsd.go
+++ b/disk/disk_openbsd.go
@@ -122,7 +122,7 @@ func parseDiskstats(buf []byte) (Diskstats, error) {
 	return ds, nil
 }
 
-func UsageWithContext(_ context.Context, path string) (*UsageStat, error) {
+func getUsage(_ context.Context, path string) (*UsageStat, error) {
 	stat := unix.Statfs_t{}
 	err := unix.Statfs(path, &stat)
 	if err != nil {
diff --git a/disk/disk_solaris.go b/disk/disk_solaris.go
index 16239ef..d8d8d59 100644
--- a/disk/disk_solaris.go
+++ b/disk/disk_solaris.go
@@ -199,7 +199,7 @@ func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOC
 	return ret, nil
 }
 
-func UsageWithContext(_ context.Context, path string) (*UsageStat, error) {
+func getUsage(_ context.Context, path string) (*UsageStat, error) {
 	statvfs := unix.Statvfs_t{}
 	if err := unix.Statvfs(path, &statvfs); err != nil {
 		return nil, fmt.Errorf("unable to call statvfs(2) on %q: %w", path, err)
diff --git a/disk/disk_unix.go b/disk/disk_unix.go
index 482372d..67d380e 100644
--- a/disk/disk_unix.go
+++ b/disk/disk_unix.go
@@ -10,7 +10,7 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-func UsageWithContext(_ context.Context, path string) (*UsageStat, error) {
+func getUsage(_ context.Context, path string) (*UsageStat, error) {
 	stat := unix.Statfs_t{}
 	err := unix.Statfs(path, &stat)
 	if err != nil {
diff --git a/disk/disk_windows.go b/disk/disk_windows.go
index b9343cb..a53c465 100644
--- a/disk/disk_windows.go
+++ b/disk/disk_windows.go
@@ -55,7 +55,7 @@ func init() {
 	}
 }
 
-func UsageWithContext(_ context.Context, path string) (*UsageStat, error) {
+func getUsage(_ context.Context, path string) (*UsageStat, error) {
 	lpFreeBytesAvailable := int64(0)
 	lpTotalNumberOfBytes := int64(0)
 	lpTotalNumberOfFreeBytes := int64(0)