From 34df4904f66be9b3b4a62c879a4ec4547072bc80 Mon Sep 17 00:00:00 2001
From: ninedraft <clamores.pro@gmail.com>
Date: Thu, 3 Sep 2020 23:01:53 +0300
Subject: [PATCH] use cancelable sleep in cpu.PercentWithContext and
 process.Process.PercentWithContext

---
 cpu/cpu.go                    |  4 +++-
 internal/common/sleep.go      | 18 ++++++++++++++++++
 internal/common/sleep_test.go | 29 +++++++++++++++++++++++++++++
 process/process.go            |  6 ++++--
 4 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 internal/common/sleep.go
 create mode 100644 internal/common/sleep_test.go

diff --git a/cpu/cpu.go b/cpu/cpu.go
index d3ea1f2..24a8116 100644
--- a/cpu/cpu.go
+++ b/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)
diff --git a/internal/common/sleep.go b/internal/common/sleep.go
new file mode 100644
index 0000000..ee27e54
--- /dev/null
+++ b/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
+	}
+}
diff --git a/internal/common/sleep_test.go b/internal/common/sleep_test.go
new file mode 100644
index 0000000..bba4661
--- /dev/null
+++ b/internal/common/sleep_test.go
@@ -0,0 +1,29 @@
+package common_test
+
+import (
+	"context"
+	"errors"
+	"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 !errors.Is(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)
+}
diff --git a/process/process.go b/process/process.go
index 482537d..eea7f40 100644
--- a/process/process.go
+++ b/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())
 }