diff --git a/cpu/cpu_windows.go b/cpu/cpu_windows.go index a44a0893..8aa691cf 100644 --- a/cpu/cpu_windows.go +++ b/cpu/cpu_windows.go @@ -90,8 +90,6 @@ func InfoWithContext(ctx context.Context) ([]InfoStat, error) { var ret []InfoStat var dst []Win32_Processor q := wmi.CreateQuery(&dst, "") - ctx, cancel := context.WithTimeout(context.Background(), common.Timeout) - defer cancel() if err := common.WMIQueryWithContext(ctx, q, &dst); err != nil { return ret, err } @@ -129,8 +127,6 @@ func PerfInfoWithContext(ctx context.Context) ([]Win32_PerfFormattedData_Counter var ret []Win32_PerfFormattedData_Counters_ProcessorInformation q := wmi.CreateQuery(&ret, "") - ctx, cancel := context.WithTimeout(context.Background(), common.Timeout) - defer cancel() err := common.WMIQueryWithContext(ctx, q, &ret) if err != nil { return []Win32_PerfFormattedData_Counters_ProcessorInformation{}, err @@ -148,8 +144,6 @@ func ProcInfo() ([]Win32_PerfFormattedData_PerfOS_System, error) { func ProcInfoWithContext(ctx context.Context) ([]Win32_PerfFormattedData_PerfOS_System, error) { var ret []Win32_PerfFormattedData_PerfOS_System q := wmi.CreateQuery(&ret, "") - ctx, cancel := context.WithTimeout(context.Background(), common.Timeout) - defer cancel() err := common.WMIQueryWithContext(ctx, q, &ret) if err != nil { return []Win32_PerfFormattedData_PerfOS_System{}, err diff --git a/disk/disk_windows.go b/disk/disk_windows.go index 7389b5a7..326bc1f4 100644 --- a/disk/disk_windows.go +++ b/disk/disk_windows.go @@ -144,8 +144,6 @@ func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOC ret := make(map[string]IOCountersStat, 0) var dst []Win32_PerfFormattedData - ctx, cancel := context.WithTimeout(context.Background(), common.Timeout) - defer cancel() err := common.WMIQueryWithContext(ctx, "SELECT * FROM Win32_PerfFormattedData_PerfDisk_LogicalDisk", &dst) if err != nil { return ret, err diff --git a/host/host_windows.go b/host/host_windows.go index 643bd66a..7897319a 100644 --- a/host/host_windows.go +++ b/host/host_windows.go @@ -48,7 +48,7 @@ func InfoWithContext(ctx context.Context) (*InfoStat, error) { } { - platform, family, version, err := PlatformInformation() + platform, family, version, err := PlatformInformationWithContext(ctx) if err == nil { ret.Platform = platform ret.PlatformFamily = family @@ -118,8 +118,6 @@ func GetOSInfo() (Win32_OperatingSystem, error) { func GetOSInfoWithContext(ctx context.Context) (Win32_OperatingSystem, error) { var dst []Win32_OperatingSystem q := wmi.CreateQuery(&dst, "") - ctx, cancel := context.WithTimeout(context.Background(), common.Timeout) - defer cancel() err := common.WMIQueryWithContext(ctx, q, &dst) if err != nil { return Win32_OperatingSystem{}, err @@ -136,7 +134,7 @@ func Uptime() (uint64, error) { func UptimeWithContext(ctx context.Context) (uint64, error) { if osInfo == nil { - _, err := GetOSInfo() + _, err := GetOSInfoWithContext(ctx) if err != nil { return 0, err } @@ -177,7 +175,7 @@ func PlatformInformation() (platform string, family string, version string, err func PlatformInformationWithContext(ctx context.Context) (platform string, family string, version string, err error) { if osInfo == nil { - _, err = GetOSInfo() + _, err = GetOSInfoWithContext(ctx) if err != nil { return } diff --git a/internal/common/common_windows.go b/internal/common/common_windows.go index 1dffe615..b02c5cf2 100644 --- a/internal/common/common_windows.go +++ b/internal/common/common_windows.go @@ -115,6 +115,12 @@ func CreateCounter(query windows.Handle, pname, cname string) (*CounterInfo, err // WMIQueryWithContext - wraps wmi.Query with a timed-out context to avoid hanging func WMIQueryWithContext(ctx context.Context, query string, dst interface{}, connectServerArgs ...interface{}) error { + if _, ok := ctx.Deadline(); !ok { + ctxTimeout, cancel := context.WithTimeout(ctx, Timeout) + defer cancel() + ctx = ctxTimeout + } + errChan := make(chan error, 1) go func() { errChan <- wmi.Query(query, dst, connectServerArgs...) diff --git a/process/process_windows.go b/process/process_windows.go index d35fcc26..f4c8fe2c 100644 --- a/process/process_windows.go +++ b/process/process_windows.go @@ -144,8 +144,6 @@ func GetWin32ProcWithContext(ctx context.Context, pid int32) ([]Win32_Process, e var dst []Win32_Process query := fmt.Sprintf("WHERE ProcessId = %d", pid) q := wmi.CreateQuery(&dst, query) - ctx, cancel := context.WithTimeout(context.Background(), common.Timeout) - defer cancel() err := common.WMIQueryWithContext(ctx, q, &dst) if err != nil { return []Win32_Process{}, fmt.Errorf("could not get win32Proc: %s", err) @@ -457,8 +455,6 @@ func (p *Process) Children() ([]*Process, error) { func (p *Process) ChildrenWithContext(ctx context.Context) ([]*Process, error) { var dst []Win32_Process query := wmi.CreateQuery(&dst, fmt.Sprintf("Where ParentProcessId = %d", p.Pid)) - ctx, cancel := context.WithTimeout(context.Background(), common.Timeout) - defer cancel() err := common.WMIQueryWithContext(ctx, query, &dst) if err != nil { return nil, err