Add option to collect cpu usage from processes

lib/index.ts

@@ -4,24 +4,32 @@
  *--------------------------------------------------------------------------------------------*/
 
 const native = require('../build/Release/windows_process_tree.node');
-import { IProcessTreeNode } from 'windows-process-tree';
+import { IProcessInfo, IProcessTreeNode, IProcessCpuInfo } from 'windows-process-tree';
 
 export enum ProcessDataFlag {
-    None = 0,
-    Memory = 1
-  }
+  None = 0,
+  Memory = 1
+}
 
-interface IProcessInfo {
-    pid: number;
-    ppid: number;
-    name: string;
-    memory?: number;
+interface  IRequest {
+  callback: (processes: IProcessTreeNode | IProcessInfo[]) => void;
+  rootPid: number;
 }
 
+type RequestQueue = IRequest[];
+
+// requestInProgress is used for any function that uses CreateToolhelp32Snapshot, as multiple calls
+// to this cannot be done at the same time.
 let requestInProgress = false;
-const requestQueue = [];
+const processListRequestQueue: RequestQueue = [];
+const processTreeRequestQueue: RequestQueue = [];
 
-function buildProcessTree(processList: IProcessInfo[], rootPid: number): IProcessTreeNode {
+/**
+ * Filters a list of processes to rootPid and its descendents and creates a tree
+ * @param rootPid The process to use as the root
+ * @param processList The list of processes
+ */
+function buildProcessTree( rootPid: number, processList: IProcessInfo[]): IProcessTreeNode {
   const rootIndex = processList.findIndex(v => v.pid === rootPid);
   if (rootIndex === -1) {
     return undefined;
@@ -33,15 +41,36 @@ function buildProcessTree(processList: IProcessInfo[], rootPid: number): IProces
     pid: rootProcess.pid,
     name: rootProcess.name,
     memory: rootProcess.memory,
-    children: childIndexes.map(c => buildProcessTree(processList, c.pid))
+    children: childIndexes.map(c => buildProcessTree(c.pid, processList))
   };
 }
 
-export function getProcessTree(rootPid: number, callback: (tree: IProcessTreeNode) => void, flags?: ProcessDataFlag): void {
-  // Push the request to the queue
-  requestQueue.push({
+/**
+ * Filters processList to contain the process with rootPid and all of its descendants
+ * @param rootPid The root pid
+ * @param processList The list of all processes
+ */
+function filterProcessList(rootPid: number, processList: IProcessInfo[]): IProcessInfo[] {
+  const rootIndex = processList.findIndex(v => v.pid === rootPid);
+  if (rootIndex === -1) {
+    return undefined;
+  }
+
+  const rootProcess = processList[rootIndex];
+  const childIndexes = processList.filter(v => v.ppid === rootPid);
+  return childIndexes.map(c => filterProcessList(c.pid, processList)).reduce((prev, current) => prev.concat(current), [rootProcess]);
+}
+
+function getRawProcessList(
+  pid: number,
+  queue: RequestQueue,
+  callback: (processList: IProcessInfo[] | IProcessTreeNode) => void,
+  filter: (pid: number, processList: IProcessInfo[]) => IProcessInfo[] | IProcessTreeNode,
+  flags?: ProcessDataFlag
+): void {
+  queue.push({
     callback: callback,
-    rootPid: rootPid
+    rootPid: pid
   });
 
   // Only make a new request if there is not currently a request in progress.
@@ -50,12 +79,41 @@ export function getProcessTree(rootPid: number, callback: (tree: IProcessTreeNod
   // once.
   if (!requestInProgress) {
     requestInProgress = true;
-    native.getProcessList((processList) => {
-      requestQueue.forEach(r => {
-        r.callback(buildProcessTree(processList, r.rootPid));
+    native.getProcessList((processList: IProcessInfo[]) => {
+      queue.forEach(r => {
+        r.callback(filter(r.rootPid, processList));
       });
-      requestQueue.length = 0;
+      queue.length = 0;
       requestInProgress = false;
     }, flags || 0);
   }
 }
+
+/**
+ * Returns a list of processes containing the rootPid process and all of its descendants
+ * @param rootPid The pid of the process of interest
+ * @param callback The callback to use with the returned set of processes
+ * @param flags The flags for what process data should be included
+ */
+export function getProcessList(rootPid: number, callback: (processList: IProcessInfo[]) => void, flags?: ProcessDataFlag): void {
+  getRawProcessList(rootPid, processListRequestQueue, callback, filterProcessList, flags);
+}
+
+/**
+ * Returns the list of processes annotated with cpu usage information
+ * @param processList The list of processes
+ * @param callback The callback to use with the returned list of processes
+ */
+export function getProcessCpuUsage(processList: IProcessInfo[], callback: (tree: IProcessCpuInfo[]) => void): void {
+  native.getProcessCpuUsage(processList, (processListWithCpu) => callback(processListWithCpu));
+}
+
+/**
+ * Returns a tree of processes with rootPid as the root
+ * @param rootPid The pid of the process that will be the root of the tree
+ * @param callback The callback to use with the returned list of processes
+ * @param flags Flags indicating what process data should be written on each node
+ */
+export function getProcessTree(rootPid: number, callback: (tree: IProcessTreeNode) => void, flags?: ProcessDataFlag): void {
+  getRawProcessList(rootPid, processTreeRequestQueue, callback, buildProcessTree, flags);
+}

lib/test.ts

@@ -5,7 +5,7 @@
 
 import * as assert from 'assert';
 import * as child_process from 'child_process';
-import { getProcessTree, ProcessDataFlag } from './index';
+import { getProcessTree, getProcessList, getProcessCpuUsage, ProcessDataFlag } from './index';
 
 const native = require('../build/Release/windows_process_tree.node');
 
@@ -21,7 +21,7 @@ function pollUntil(makePromise: () => Promise<boolean>, cb: () => void, interval
   });
 }
 
-describe('getProcessList', () => {
+describe('getRawProcessList', () => {
   it('should throw if arguments are not provided', (done) => {
     assert.throws(() => native.getProcessList());
     done();
@@ -70,6 +70,80 @@ describe('getProcessList', () => {
   });
 });
 
+describe('getProcessList', () => {
+  let cps;
+
+  beforeEach(() => {
+    cps = [];
+  });
+
+  afterEach(() => {
+    cps.forEach(cp => {
+      cp.kill();
+    });
+  });
+
+  it('should return a list containing this process', (done) => {
+    getProcessList(process.pid, (list) => {
+      assert.equal(list.length, 1);
+      assert.equal(list[0].name, 'node.exe');
+      assert.equal(list[0].pid, process.pid);
+      assert.equal(list[0].memory, undefined);
+      done();
+    });
+  });
+
+  it('should return a list containing this process\'s memory if the flag is set', done => {
+    getProcessList(process.pid, (list) => {
+      assert.equal(list.length, 1);
+      assert.equal(list[0].name, 'node.exe');
+      assert.equal(list[0].pid, process.pid);
+      assert.equal(typeof list[0].memory, 'number');
+      done();
+    }, ProcessDataFlag.Memory);
+  });
+
+  it('should return a tree containing this process\'s child processes', done => {
+    cps.push(child_process.spawn('cmd.exe'));
+    pollUntil(() => {
+      return new Promise((resolve) => {
+        getProcessList(process.pid, (list) => {
+          resolve(list.length === 2 && list[0].pid === process.pid && list[1].pid === cps[0].pid);
+        });
+      });
+    }, () => done(), 20, 500);
+  });
+});
+
+describe('getProcessCpuUsage', () => {
+
+  it('should get process cpu usage', (done) => {
+      getProcessCpuUsage([{ pid: process.pid, ppid: process.ppid, name: 'node.exe' }], (annotatedList) => {
+        assert.equal(annotatedList.length, 1);
+        assert.equal(annotatedList[0].name, 'node.exe');
+        assert.equal(annotatedList[0].pid, process.pid);
+        assert.equal(annotatedList[0].memory, undefined);
+        assert.equal(typeof annotatedList[0].cpu, 'number');
+        assert.equal(0 <= annotatedList[0].cpu && annotatedList[0].cpu <= 100, true);
+        done();
+    });
+  });
+
+  it('should handle multiple calls gracefully', function (done: MochaDone): void {
+    this.timeout(3000);
+
+    let counter = 0;
+    const callback = (list) => {
+      assert.notEqual(list.find(p => p.pid === process.pid), undefined);
+      if (++counter === 2) {
+        done();
+      }
+    };
+    getProcessCpuUsage([{ pid: process.pid, ppid: process.ppid, name: 'node.exe' }], callback);
+    getProcessCpuUsage([{ pid: process.pid, ppid: process.ppid, name: 'node.exe' }], callback);
+  });
+});
+
 describe('getProcessTree', () => {
   let cps;
 

src/addon.cc

@@ -6,6 +6,7 @@
 #include <nan.h>
 #include "process.h"
 #include "worker.h"
+#include <cmath>
 
 void GetProcessList(const Nan::FunctionCallbackInfo<v8::Value>& args) {
   ProcessInfo process_info[1024];
@@ -32,9 +33,66 @@ void GetProcessList(const Nan::FunctionCallbackInfo<v8::Value>& args) {
   Nan::AsyncQueueWorker(worker);
 }
 
+void GetProcessCpuUsage(const Nan::FunctionCallbackInfo<v8::Value>& args) {
+  if (args.Length() < 2) {
+      Nan::ThrowTypeError("GetProcessCpuUsage expects two arguments.");
+      return;
+  }
+
+  if (!args[0]->IsArray()) {
+    Nan::ThrowTypeError("The first argument of GetProcessCpuUsage, callback, must be an array.");
+    return;
+  }
+
+  if (!args[1]->IsFunction()) {
+    Nan::ThrowTypeError("The second argument of GetProcessCpuUsage, flags, must be a function.");
+    return;
+  }
+
+  // Read the ProcessTreeNode JS object
+  v8::Local<v8::Array> processes = v8::Local<v8::Array>::Cast(args[0]);
+  uint32_t count = processes->Length();
+  Cpu* cpu_info = new Cpu[count];
+
+  // Read pid from each array and populate data structure to calculate CPU, take first sample of counters
+  for (uint32_t i = 0; i < count; i++) {
+    v8::Local<v8::Object> process = processes->Get(Nan::New<v8::Integer>(i))->ToObject();
+    DWORD pid = (DWORD)(process->Get(Nan::New("pid").ToLocalChecked()))->NumberValue();
+    cpu_info[i].pid = pid;
+    GetCpuUsage(cpu_info, &i, true);
+  }
+
+  // Sleep for one second
+  Sleep(1000);
+
+  // Sample counters again and complete CPU usage calculation
+  for (uint32_t i = 0; i < count; i++) {
+    GetCpuUsage(cpu_info, &i, false);
+  }
+
+  Nan::Callback *callback = new Nan::Callback(v8::Local<v8::Function>::Cast(args[1]));
+
+  v8::Local<v8::Array> result = Nan::New<v8::Array>(count);
+  for (uint32_t i = 0; i < count; i++) {
+    v8::Local<v8::Object> object = processes->Get(Nan::New<v8::Integer>(i))->ToObject();
+
+    if (!std::isnan(cpu_info[i].cpu)) {
+      Nan::Set(object, Nan::New<v8::String>("cpu").ToLocalChecked(),
+                Nan::New<v8::Number>(cpu_info[i].cpu));
+    }
+
+    Nan::Set(result, i, Nan::New<v8::Value>(object));
+  }
+
+  v8::Local<v8::Value> argv[] = { result };
+  callback->Call(1, argv);
+}
+
 void Init(v8::Local<v8::Object> exports) {
   exports->Set(Nan::New("getProcessList").ToLocalChecked(),
                Nan::New<v8::FunctionTemplate>(GetProcessList)->GetFunction());
+  exports->Set(Nan::New("getProcessCpuUsage").ToLocalChecked(),
+               Nan::New<v8::FunctionTemplate>(GetProcessCpuUsage)->GetFunction());
 }
 
 NODE_MODULE(hello, Init)

src/process.cc

@@ -7,6 +7,7 @@
 
 #include <tlhelp32.h>
 #include <psapi.h>
+#include <limits>
 
 void GetRawProcessList(ProcessInfo process_info[1024], uint32_t* process_count, DWORD* process_data_flags) {
   *process_count = 0;
@@ -31,6 +32,7 @@ void GetRawProcessList(ProcessInfo process_info[1024], uint32_t* process_count,
       }
     } while (*process_count < 1024 && Process32Next(snapshot_handle, &process_entry));
   }
+
   CloseHandle(snapshot_handle);
 }
 
@@ -48,6 +50,49 @@ void GetProcessMemoryUsage(ProcessInfo process_info[1024], uint32_t* process_cou
   if (GetProcessMemoryInfo(hProcess, &pmc, sizeof(pmc))) {
     process_info[*process_count].memory = (DWORD)pmc.WorkingSetSize;
   }
-
+}
+
+// Per documentation, it is not recommended to add or subtract values from the FILETIME
+// structure, or to cast it to ULARGE_INTEGER as this can cause alignment faults on 64-bit Windows.
+// Copy the high and low part to a ULARGE_INTEGER and peform arithmetic on that instead.
+// See https://msdn.microsoft.com/en-us/library/windows/desktop/ms724284(v=vs.85).aspx
+ULONGLONG GetTotalTime(const FILETIME* kernelTime, const FILETIME* userTime) {
+  ULARGE_INTEGER kt, ut;
+  kt.LowPart = (*kernelTime).dwLowDateTime;
+  kt.HighPart = (*kernelTime).dwHighDateTime;
+
+  ut.LowPart = (*userTime).dwLowDateTime;
+  ut.HighPart = (*userTime).dwHighDateTime;
+
+  return kt.QuadPart + ut.QuadPart;
+}
+
+void GetCpuUsage(Cpu* cpu_info, uint32_t* process_index, BOOL first_pass) {
+  DWORD pid = cpu_info[*process_index].pid;
+  HANDLE hProcess;
+
+  hProcess = OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, false, pid);
+
+  if (hProcess == NULL) {
+    return;
+  }
+
+  FILETIME creationTime, exitTime, kernelTime, userTime;
+  FILETIME sysIdleTime, sysKernelTime, sysUserTime;
+  if (GetProcessTimes(hProcess, &creationTime, &exitTime, &kernelTime, &userTime)
+    && GetSystemTimes(&sysIdleTime, &sysKernelTime, &sysUserTime)) {
+    if (first_pass) {
+      cpu_info[*process_index].initialProcRunTime = GetTotalTime(&kernelTime, &userTime);
+      cpu_info[*process_index].initialSystemTime = GetTotalTime(&sysKernelTime, &sysUserTime);
+    } else {
+      ULONGLONG endProcTime = GetTotalTime(&kernelTime, &userTime);
+      ULONGLONG endSysTime = GetTotalTime(&sysKernelTime, &sysUserTime);
+
+      cpu_info[*process_index].cpu = 100.0 * (endProcTime - cpu_info[*process_index].initialProcRunTime) / (endSysTime - cpu_info[*process_index].initialSystemTime);
+    }
+  } else {
+    cpu_info[*process_index].cpu = std::numeric_limits<double>::quiet_NaN();
+  }
+
   CloseHandle(hProcess);
 }

src/process.h

@@ -9,6 +9,13 @@
 #include <nan.h>
 #include <windows.h>
 
+struct Cpu {
+  DWORD pid;
+  double cpu;
+  ULONGLONG initialProcRunTime;
+  ULONGLONG initialSystemTime;
+};
+
 struct ProcessInfo {
   TCHAR name[MAX_PATH];
   DWORD pid;
@@ -25,4 +32,6 @@ void GetRawProcessList(ProcessInfo process_info[1024], uint32_t* process_count,
 
 void GetProcessMemoryUsage(ProcessInfo process_info[1024], uint32_t* process_count);
 
+void GetCpuUsage(Cpu cpu_info[1024], uint32_t* process_count, BOOL first_run);
+
 #endif  // SRC_PROCESS_H_