From 6b539051d2b215e59e65ff8a057cce3dd0eefad5 Mon Sep 17 00:00:00 2001
From: Lomanic <lomanic@hotmail.fr>
Date: Sun, 18 Nov 2018 20:11:55 +0100
Subject: [PATCH] [process][windows] Use win32 API in process.Children()
 instead of slow WMI call

The CreateToolhelp32Snapshot+Process32First+Process32Next combo already
iterates over all processes, so it would be inefficient to enumerate all
processes with process.Processes() and then calling p.Ppid() on each of
them: we just use this combo to get all processes and their ppid in a
single iteration.

This is faster by a factor of 25 compared to the previous WMI call.
---
 process/process_windows.go | 34 +++++++++++++++++++++++-----------
 1 file changed, 23 insertions(+), 11 deletions(-)

diff --git a/process/process_windows.go b/process/process_windows.go
index 82aed37e3..0d27a99a7 100644
--- a/process/process_windows.go
+++ b/process/process_windows.go
@@ -250,6 +250,7 @@ func (p *Process) Status() (string, error) {
 func (p *Process) StatusWithContext(ctx context.Context) (string, error) {
 	return "", common.ErrNotImplementedError
 }
+
 func (p *Process) Username() (string, error) {
 	return p.UsernameWithContext(context.Background())
 }
@@ -456,22 +457,33 @@ 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))
-	err := common.WMIQueryWithContext(ctx, query, &dst)
-	if err != nil {
-		return nil, err
+	out := []*Process{}
+	snap := w32.CreateToolhelp32Snapshot(w32.TH32CS_SNAPPROCESS, uint32(0))
+	if snap == 0 {
+		return out, windows.GetLastError()
+	}
+	defer w32.CloseHandle(snap)
+	var pe32 w32.PROCESSENTRY32
+	pe32.DwSize = uint32(unsafe.Sizeof(pe32))
+	if w32.Process32First(snap, &pe32) == false {
+		return out, windows.GetLastError()
 	}
 
-	out := []*Process{}
-	for _, proc := range dst {
-		p, err := NewProcess(int32(proc.ProcessID))
-		if err != nil {
-			continue
+	if pe32.Th32ParentProcessID == uint32(p.Pid) {
+		p, err := NewProcess(int32(pe32.Th32ProcessID))
+		if err == nil {
+			out = append(out, p)
 		}
-		out = append(out, p)
 	}
 
+	for w32.Process32Next(snap, &pe32) {
+		if pe32.Th32ParentProcessID == uint32(p.Pid) {
+			p, err := NewProcess(int32(pe32.Th32ProcessID))
+			if err == nil {
+				out = append(out, p)
+			}
+		}
+	}
 	return out, nil
 }