1 #include <windows.h> 2 #include <assert.h> 3 #include <psapi.h> 4 #include <stdio.h> 5 #include <tchar.h> 6 #include <time.h> 7 #include <tlhelp32.h> 8 #include "Shlwapi.h" 9 10 #pragma comment(lib, "psapi.lib") 11 #pragma comment(lib, "shlwapi.lib") 12 13 int gQueryInterval = 5; // seconds 14 time_t gDuration = 0; // seconds 15 LPTSTR gCommandLine; 16 17 HRESULT ProcessArgs(int argc, TCHAR *argv[]); 18 HRESULT PrintUsage(); 19 void UseImage(void (functionForQueryType(HANDLE))); 20 void QueryContinuously(HANDLE hProcess); 21 int EvalProcesses(HANDLE hProcess); 22 time_t ElapsedTime(time_t startTime); 23 24 int __cdecl _tmain (int argc, TCHAR *argv[]) 25 { 26 HRESULT result = ProcessArgs(argc, argv); 27 if (FAILED(result)) 28 return result; 29 30 UseImage(QueryContinuously); 31 return S_OK; 32 } 33 34 HRESULT ProcessArgs(int argc, TCHAR *argv[]) 35 { 36 LPTSTR argument; 37 for( int count = 1; count < argc; count++ ) { 38 argument = argv[count] ; 39 if (wcsstr(argument, _T("-h")) || wcsstr(argument, _T("--help"))) 40 return PrintUsage(); 41 else if (wcsstr(argument, _T("--exe"))) { 42 gCommandLine = argv[++count]; 43 } else if (wcsstr(argument, _T("-i")) || 44 wcsstr(argument, _T("--interval"))) { 45 gQueryInterval = _wtoi(argv[++count]); 46 if (gQueryInterval < 1) { 47 printf("ERROR: invalid interval\n"); 48 return E_INVALIDARG; 49 } 50 } else if (wcsstr(argument, _T("-d")) || 51 wcsstr(argument, _T("--duration"))) { 52 gDuration = _wtoi(argv[++count]); 53 if (gDuration < 1) { 54 printf("ERROR: invalid duration\n"); 55 return E_INVALIDARG; 56 } 57 } else { 58 _tprintf(_T("ERROR: unrecognized argument \"%s\"\n"), (LPCTSTR)argument); 59 return PrintUsage(); 60 } 61 } 62 if (argc < 2 || !wcslen(gCommandLine) ) { 63 printf("ERROR: executable path is required\n"); 64 return PrintUsage(); 65 } 66 return S_OK; 67 } 68 69 HRESULT PrintUsage() 70 { 71 printf("record-memory-win --exe EXE_PATH\n"); 72 printf(" Launch an executable and print the memory usage (in Private Bytes)\n"); 73 printf(" of the process.\n\n"); 74 printf("Usage:\n"); 75 printf("-h [--help] : Print usage\n"); 76 printf("--exe arg : Launch specified image. Required\n"); 77 printf("-i [--interval] arg : Print memory usage every arg seconds. Default: 5 seconds\n"); 78 printf("-d [--duration] arg : Run for up to arg seconds. Default: no limit\n\n"); 79 printf("Examples:\n"); 80 printf(" record-memory-win --exe \"C:\\Program Files\\Safari\\Safari.exe /newprocess\"\n"); 81 printf(" record-memory-win --exe \"Safari.exe /newprocess\" -i 10 -d 7200\n"); 82 printf(" NOTE: Close all other browser intances to ensure launching in a new process\n"); 83 printf(" Or, pass the /newprocess (or equivalent) argument to the browser\n"); 84 return E_FAIL; 85 } 86 87 unsigned int getMemoryInfo(DWORD processID) 88 { 89 unsigned int memInfo = 0; 90 HANDLE hProcess; 91 PROCESS_MEMORY_COUNTERS_EX pmc; 92 93 hProcess = OpenProcess( PROCESS_QUERY_INFORMATION | 94 PROCESS_VM_READ, 95 FALSE, processID ); 96 if (NULL == hProcess) 97 return 0; 98 99 if (GetProcessMemoryInfo( hProcess, (PPROCESS_MEMORY_COUNTERS)&pmc, sizeof(pmc))) { 100 memInfo = (pmc.PrivateUsage); 101 } 102 103 CloseHandle( hProcess ); 104 return memInfo; 105 } 106 107 void printProcessInfo(DWORD processID) 108 { 109 TCHAR szProcessName[MAX_PATH] = TEXT("<unknown>"); 110 111 // Get a handle to the process. 112 HANDLE hProcess = OpenProcess( PROCESS_QUERY_INFORMATION | 113 PROCESS_VM_READ, 114 FALSE, processID ); 115 116 // Get the process name. 117 if (NULL != hProcess) { 118 HMODULE hMod; // An array that receives the list of module handles. 119 DWORD cbNeeded; //The number of bytes required to store all module handles in the Module array 120 121 if (EnumProcessModules(hProcess, &hMod, sizeof(hMod), &cbNeeded)) { 122 GetModuleBaseName(hProcess, hMod, szProcessName, 123 sizeof(szProcessName)/sizeof(TCHAR)); 124 } 125 } 126 127 // Print the process name and identifier of matching strings, ignoring case 128 _tprintf(TEXT("%s (PID: %u)\n"), szProcessName, processID); 129 130 // Release the handle to the process. 131 CloseHandle( hProcess ); 132 } 133 134 int evalProcesses(HANDLE hProcess) 135 { 136 if (NULL == hProcess) 137 return 0; 138 139 unsigned int totalMemUsage = 0; 140 DWORD processID = GetProcessId(hProcess); 141 142 HANDLE hProcessSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0); 143 144 PROCESSENTRY32 processEntry = { 0 }; 145 processEntry.dwSize = sizeof(PROCESSENTRY32); 146 147 // Retrieves information about the first process encountered in a system snapshot 148 if(Process32First(hProcessSnapshot, &processEntry)) { 149 do { 150 // if th32processID = processID, we are the parent process! 151 // if th32ParentProcessID = processID, we are a child process! 152 if ((processEntry.th32ProcessID == processID) || (processEntry.th32ParentProcessID == processID)) { 153 unsigned int procMemUsage = 0; 154 // Record parent process memory 155 procMemUsage = getMemoryInfo(processEntry.th32ProcessID); 156 totalMemUsage += procMemUsage; 157 } 158 // Retrieves information about the next process recorded in a system snapshot. 159 } while(Process32Next(hProcessSnapshot, &processEntry)); 160 } 161 162 CloseHandle(hProcessSnapshot); 163 return totalMemUsage; 164 } 165 166 167 void UseImage(void (functionForQueryType(HANDLE))) 168 { 169 STARTUPINFO si = {0}; 170 si.cb = sizeof(STARTUPINFO); 171 PROCESS_INFORMATION pi = {0}; 172 173 // Start the child process. 174 if(!CreateProcess( NULL, // No module name (use command line) 175 gCommandLine, // Command line 176 NULL, // Process handle not inheritable 177 NULL, // Thread handle not inheritable 178 FALSE, // Set handle inheritance to FALSE 179 0, // No creation flags 180 NULL, // Use parent's environment block 181 NULL, // Use parent's starting directory 182 &si, // Pointer to STARTUPINFO structure 183 &pi )) // Pointer to PROCESS_INFORMATION structure 184 printf("CreateProcess failed (%d)\n", GetLastError()); 185 else { 186 printf("Created process with id: %d\n", pi.dwProcessId); 187 functionForQueryType(pi.hProcess); 188 // Close process and thread handles. 189 CloseHandle( pi.hProcess ); 190 CloseHandle( pi.hThread ); 191 } 192 } 193 194 void QueryContinuously(HANDLE hProcess) 195 { 196 Sleep(2000); // give the process some time to launch 197 bool pastDuration = false; 198 time_t startTime = time(NULL); 199 unsigned int memUsage = evalProcesses(hProcess); 200 while(memUsage && !pastDuration) { 201 printf( "%u\n", memUsage ); 202 Sleep(gQueryInterval*1000); 203 memUsage = evalProcesses(hProcess); 204 pastDuration = gDuration > 0 ? ElapsedTime(startTime) > gDuration : false; 205 } 206 } 207 208 // returns elapsed time in seconds 209 time_t ElapsedTime(time_t startTime) 210 { 211 time_t currentTime = time(NULL); 212 return currentTime - startTime; 213 } 214
