Merge pull request #936 from ninedraft/fix-context-usage

Use cancelable sleep

cpu/cpu.go

@@ -149,7 +149,9 @@ func PercentWithContext(ctx context.Context, interval time.Duration, percpu bool
 		return nil, err
 	}
 
-	time.Sleep(interval)
+	if err := common.Sleep(ctx, interval); err != nil {
+		return nil, err
+	}
 
 	// And at the end of the interval.
 	cpuTimes2, err := Times(percpu)

internal/common/sleep.go

@@ -0,0 +1,18 @@
+package common
+
+import (
+	"context"
+	"time"
+)
+
+// Sleep awaits for provided interval.
+// Can be interrupted by context cancelation.
+func Sleep(ctx context.Context, interval time.Duration) error {
+	var timer = time.NewTimer(interval)
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-timer.C:
+		return nil
+	}
+}

internal/common/sleep_test.go

@@ -0,0 +1,28 @@
+package common_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/shirou/gopsutil/internal/common"
+)
+
+func TestSleep(test *testing.T) {
+	const dt = 50 * time.Millisecond
+	var t = func(name string, ctx context.Context, expected error) {
+		test.Run(name, func(test *testing.T) {
+			var err = common.Sleep(ctx, dt)
+			if err != expected {
+				test.Errorf("expected %v, got %v", expected, err)
+			}
+		})
+	}
+
+	var ctx = context.Background()
+	var canceled, cancel = context.WithCancel(ctx)
+	cancel()
+
+	t("background context", ctx, nil)
+	t("canceled context", canceled, context.Canceled)
+}

process/process.go

@@ -201,7 +201,9 @@ func (p *Process) PercentWithContext(ctx context.Context, interval time.Duration
 	if interval > 0 {
 		p.lastCPUTimes = cpuTimes
 		p.lastCPUTime = now
-		time.Sleep(interval)
+		if err := common.Sleep(ctx, interval); err != nil {
+			return 0, err
+		}
 		cpuTimes, err = p.Times()
 		now = time.Now()
 		if err != nil {
@@ -316,5 +318,5 @@ func (p *Process) CPUPercentWithContext(ctx context.Context) (float64, error) {
 
 // Groups returns all group IDs(include supplementary groups) of the process as a slice of the int
 func (p *Process) Groups() ([]int32, error) {
-       return p.GroupsWithContext(context.Background())
+	return p.GroupsWithContext(context.Background())
 }