From 8e2446b0ac56f0da5de6dc92dac4e925fb778bb4 Mon Sep 17 00:00:00 2001
From: shirou <shirou.faw@gmail.com>
Date: Sat, 31 Oct 2020 23:39:29 +0900
Subject: [PATCH] [v3][process] apply #962

---
 v3/process/process.go         |  7 ++++++-
 v3/process/process_windows.go | 34 ++++++++++++++++++++++++++++------
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/v3/process/process.go b/v3/process/process.go
index 03cdfb4..28deb2f 100644
--- a/v3/process/process.go
+++ b/v3/process/process.go
@@ -6,6 +6,7 @@ import (
 	"errors"
 	"runtime"
 	"sort"
+	"sync"
 	"syscall"
 	"time"
 
@@ -26,6 +27,7 @@ type Process struct {
 	name           string
 	status         string
 	parent         int32
+	parentMutex    *sync.RWMutex // for windows ppid cache
 	numCtxSwitches *NumCtxSwitchesStat
 	uids           []int32
 	gids           []int32
@@ -178,7 +180,10 @@ func NewProcess(pid int32) (*Process, error) {
 }
 
 func NewProcessWithContext(ctx context.Context, pid int32) (*Process, error) {
-	p := &Process{Pid: pid}
+	p := &Process{
+		Pid:         pid,
+		parentMutex: new(sync.RWMutex),
+	}
 
 	exists, err := PidExistsWithContext(ctx, pid)
 	if err != nil {
diff --git a/v3/process/process_windows.go b/v3/process/process_windows.go
index 45e7df8..5cad3ba 100644
--- a/v3/process/process_windows.go
+++ b/v3/process/process_windows.go
@@ -222,8 +222,9 @@ func PidExistsWithContext(ctx context.Context, pid int32) (bool, error) {
 
 func (p *Process) PpidWithContext(ctx context.Context) (int32, error) {
 	// if cached already, return from cache
-	if p.parent != 0 {
-		return p.parent, nil
+	cachedPpid := p.getPpid()
+	if cachedPpid != 0 {
+		return cachedPpid, nil
 	}
 
 	ppid, _, _, err := getFromSnapProcess(p.Pid)
@@ -231,8 +232,8 @@ func (p *Process) PpidWithContext(ctx context.Context) (int32, error) {
 		return 0, err
 	}
 
-	// if no errors, cache it
-	p.parent = ppid
+	// no errors and not cached already, so cache it
+	p.setPpid(ppid)
 
 	return ppid, nil
 }
@@ -243,8 +244,11 @@ func (p *Process) NameWithContext(ctx context.Context) (string, error) {
 		return "", fmt.Errorf("could not get Name: %s", err)
 	}
 
-	// if no errors, cache ppid
+	// if no errors and not cached already, cache ppid
 	p.parent = ppid
+	if 0 == p.getPpid() {
+		p.setPpid(ppid)
+	}
 
 	return name, nil
 }
@@ -441,8 +445,11 @@ func (p *Process) NumThreadsWithContext(ctx context.Context) (int32, error) {
 		return 0, err
 	}
 
-	// if no errors, cache ppid
+	// if no errors and not cached already, cache ppid
 	p.parent = ppid
+	if 0 == p.getPpid() {
+		p.setPpid(ppid)
+	}
 
 	return ret, nil
 }
@@ -594,6 +601,21 @@ func (p *Process) KillWithContext(ctx context.Context) error {
 	return process.Kill()
 }
 
+// retrieve Ppid in a thread-safe manner
+func (p *Process) getPpid() int32 {
+	p.parentMutex.RLock()
+	defer p.parentMutex.RUnlock()
+	return p.parent
+}
+
+// cache Ppid in a thread-safe manner (WINDOWS ONLY)
+// see https://psutil.readthedocs.io/en/latest/#psutil.Process.ppid
+func (p *Process) setPpid(ppid int32) {
+	p.parentMutex.Lock()
+	defer p.parentMutex.Unlock()
+	p.parent = ppid
+}
+
 func getFromSnapProcess(pid int32) (int32, int32, string, error) {
 	snap, err := windows.CreateToolhelp32Snapshot(windows.TH32CS_SNAPPROCESS, uint32(pid))
 	if err != nil {