PR feedback

Don't leak process handles.

Signed-off-by: Daniel Canter <[email protected]>

internal/jobobject/jobobject.go

@@ -557,7 +557,31 @@ func (job *JobObject) PrivateWorkingSet() (uint64, error) {
 		return 0, err
 	}
 
-	queryProcessWorkingSet := func(h windows.Handle) (uint64, error) {
+	openAndQueryWorkingSet := func(pid uint32) (uint64, error) {
+		h, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, pid)
+		if err != nil {
+			// Continue to the next if OpenProcess doesn't return a valid handle (fails). Handles a
+			// case where one of the pids in the job exited before we open.
+			return 0, nil
+		}
+		defer func() {
+			_ = windows.Close(h)
+		}()
+		// Check if the process is actually running in the job still. There's a small chance
+		// that the process could have exited and had its pid re-used between grabbing the pids
+		// in the job and opening the handle to it above.
+		var inJob int32
+		if err := winapi.IsProcessInJob(h, job.handle, &inJob); err != nil {
+			// This shouldn't fail unless we have incorrect access rights which we control
+			// here so probably best to error out if this failed.
+			return 0, err
+		}
+		// Don't report stats for this process as it's not running in the job. This shouldn't be
+		// an error condition though.
+		if inJob == 0 {
+			return 0, nil
+		}
+
 		var vmCounters winapi.VM_COUNTERS_EX2
 		status := winapi.NtQueryInformationProcess(
 			h,
@@ -574,29 +598,11 @@ func (job *JobObject) PrivateWorkingSet() (uint64, error) {
 
 	var jobWorkingSetSize uint64
 	for _, pid := range pids {
-		h, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, pid)
+		workingSet, err := openAndQueryWorkingSet(pid)
 		if err != nil {
-			// Continue to the next if OpenProcess doesn't return a valid handle (fails). Handles a
-			// case where one of the pids in the job exited before we open.
-			continue
-		}
-		// Check if the process is actually running in the job still. There's a small chance
-		// that the process could have exited and had its pid re-used between grabbing the pids
-		// in the job and opening the handle to it above.
-		var inJob int32
-		if err := winapi.IsProcessInJob(h, job.handle, &inJob); err != nil {
-			windows.Close(h)
-			// This shouldn't fail unless we have incorrect access rights which we control
-			// here so probably best to error out if this failed.
 			return 0, err
 		}
-		if inJob == 1 {
-			workingSet, err := queryProcessWorkingSet(h)
-			if err != nil {
-				return 0, err
-			}
-			jobWorkingSetSize += workingSet
-		}
+		jobWorkingSetSize += workingSet
 	}
 
 	return jobWorkingSetSize, nil