[disk][windows]: add context handling on PartionsWithContext

disk/disk_windows.go

@@ -81,17 +81,27 @@ func UsageWithContext(ctx context.Context, path string) (*UsageStat, error) {
 	return ret, nil
 }
 
+// PartitionsWithContext returns disk partitions.
+// Since GetVolumeInformation doesn't have a timeout, this method uses context to set deadline by users.
 func PartitionsWithContext(ctx context.Context, all bool) ([]PartitionStat, error) {
 	warnings := Warnings{
 		Verbose: true,
 	}
 	var ret []PartitionStat
+	retChan := make(chan []PartitionStat)
+	errChan := make(chan error)
+	defer close(retChan)
+	defer close(errChan)
+
 	lpBuffer := make([]byte, 254)
+
+	f := func() {
 		diskret, _, err := procGetLogicalDriveStringsW.Call(
 			uintptr(len(lpBuffer)),
 			uintptr(unsafe.Pointer(&lpBuffer[0])))
 		if diskret == 0 {
-		return ret, err
+			errChan <- err
+			return
 		}
 		for _, v := range lpBuffer {
 			if v >= 65 && v <= 90 {
@@ -146,7 +156,18 @@ func PartitionsWithContext(ctx context.Context, all bool) ([]PartitionStat, erro
 				}
 			}
 		}
+		retChan <- ret
+	}
+
+	go f()
+	select {
+	case err := <-errChan:
+		return ret, err
+	case ret := <-retChan:
 		return ret, warnings.Reference()
+	case <-ctx.Done():
+		return ret, ctx.Err()
+	}
 }
 
 func IOCountersWithContext(ctx context.Context, names ...string) (map[string]IOCountersStat, error) {