Merge pull request #488 from sadag/master

windows: fix WithContext functions

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

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

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
 		}

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...)

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