Add option to collect cpu usage from processes

lib/index.ts

@@ -4,23 +4,29 @@
  *--------------------------------------------------------------------------------------------*/
 
 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
   }
 
-interface IProcessInfo {
-    pid: number;
-    ppid: number;
-    name: string;
-    memory?: number;
-}
-
+// 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 = [];
+let cpuUsageRequestInProgress = false;
+
+const requestQueue = {
+  getProcessCpuUsage: [],
+  getProcessList: [],
+  getProcessTree: []
+};
 
+/**
+ * Filters a list of processes to rootPid and its descendents and creates a tree
+ * @param processList The list of processes
+ * @param rootPid The process to use as the root
+ */
 function buildProcessTree(processList: IProcessInfo[], rootPid: number): IProcessTreeNode {
   const rootIndex = processList.findIndex(v => v.pid === rootPid);
   if (rootIndex === -1) {
@@ -37,9 +43,83 @@ function buildProcessTree(processList: IProcessInfo[], rootPid: number): IProces
   };
 }
 
+/**
+ * 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]);
+}
+
+/**
+ * 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 {
+  // Push the request to the queue
+  requestQueue.getProcessList.push({
+    callback: callback,
+    rootPid: rootPid
+  });
+
+  // Only make a new request if there is not currently a request in progress.
+  // This prevents too many requests from being made, there is also a crash that
+  // can occur when performing multiple calls to CreateToolhelp32Snapshot at
+  // once.
+  if (!requestInProgress) {
+    requestInProgress = true;
+    native.getProcessList((processList) => {
+      requestQueue.getProcessList.forEach(r => {
+        r.callback(filterProcessList(r.rootPid, processList));
+      });
+      requestQueue.getProcessList.length = 0;
+      requestInProgress = false;
+    }, flags || 0);
+  }
+}
+
+/**
+ * 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 {
+    // Push the request to the queue
+    requestQueue.getProcessCpuUsage.push({
+      callback: callback
+    });
+
+    if (!cpuUsageRequestInProgress) {
+      cpuUsageRequestInProgress = true;
+      native.getProcessCpuUsage(processList, (processListWithCpu) => {
+        requestQueue.getProcessCpuUsage.forEach(r => {
+          r.callback(processListWithCpu);
+        });
+        requestQueue.getProcessCpuUsage.length = 0;
+        cpuUsageRequestInProgress = false;
+      });
+    }
+}
+
+/**
+ * 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 {
   // Push the request to the queue
-  requestQueue.push({
+  requestQueue.getProcessTree.push({
     callback: callback,
     rootPid: rootPid
   });
@@ -51,10 +131,10 @@ export function getProcessTree(rootPid: number, callback: (tree: IProcessTreeNod
   if (!requestInProgress) {
     requestInProgress = true;
     native.getProcessList((processList) => {
-      requestQueue.forEach(r => {
+      requestQueue.getProcessTree.forEach(r => {
         r.callback(buildProcessTree(processList, r.rootPid));
       });
-      requestQueue.length = 0;
+      requestQueue.getProcessTree.length = 0;
       requestInProgress = false;
     }, flags || 0);
   }

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,75 @@ 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) => {
+    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);
+      assert.equal((list[0] as any).cpu, undefined);
+
+      getProcessCpuUsage(list, (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();
+      });
+    });
+  });
+
+});
+
 describe('getProcessTree', () => {
   let cps;
 

src/addon.cc

@@ -32,9 +32,65 @@ 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++) {
+    // Should this be cloned?
+    v8::Local<v8::Object> object = processes->Get(Nan::New<v8::Integer>(i))->ToObject();
+    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

@@ -31,6 +31,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 +49,43 @@ void GetProcessMemoryUsage(ProcessInfo process_info[1024], uint32_t* process_cou
   if (GetProcessMemoryInfo(hProcess, &pmc, sizeof(pmc))) {
     process_info[*process_count].memory = (DWORD)pmc.WorkingSetSize;
   }
-
+}
+
+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);
+      }
+  }
+
   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;
@@ -18,11 +25,14 @@ struct ProcessInfo {
 
 enum ProcessDataFlags {
   NONE = 0,
-  MEMORY = 1
+  MEMORY = 1,
+  CPU = 2
 };
 
 void GetRawProcessList(ProcessInfo process_info[1024], uint32_t* process_count, DWORD* flags);
 
 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_