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 "Shlwapi.h" 8 9 #pragma comment(lib, "psapi.lib") 10 #pragma comment(lib, "shlwapi.lib") 11 12 bool gSingleProcess = true; 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 time_t ElapsedTime(time_t startTime); 22 unsigned int OneQuery(HANDLE hProcess); 23 unsigned int OneQueryMP(HANDLE hProcess); 24 25 int __cdecl _tmain (int argc, TCHAR *argv[]) 26 { 27 HRESULT result = ProcessArgs(argc, argv); 28 if (FAILED(result)) 29 return result; 30 31 UseImage(QueryContinuously); 32 return S_OK; 33 } 34 35 HRESULT ProcessArgs(int argc, TCHAR *argv[]) 36 { 37 LPTSTR argument; 38 for( int count = 1; count < argc; count++ ) { 39 argument = argv[count] ; 40 if (wcsstr(argument, _T("-h")) || 41 wcsstr(argument, _T("--help"))) 42 return PrintUsage(); 43 else if (wcsstr(argument, _T("--exe"))) { 44 gCommandLine = argv[++count]; 45 if (wcsstr(gCommandLine, _T("chrome.exe"))) 46 gSingleProcess = false; 47 } else if (wcsstr(argument, _T("-i")) || 48 wcsstr(argument, _T("--interval"))) { 49 gQueryInterval = _wtoi(argv[++count]); 50 if (gQueryInterval < 1) { 51 printf("ERROR: invalid interval\n"); 52 return E_INVALIDARG; 53 } 54 } else if (wcsstr(argument, _T("-d")) || 55 wcsstr(argument, _T("--duration"))) { 56 gDuration = _wtoi(argv[++count]); 57 if (gDuration < 1) { 58 printf("ERROR: invalid duration\n"); 59 return E_INVALIDARG; 60 } 61 } else { 62 _tprintf(_T("ERROR: unrecognized argument \"%s\"\n"), (LPCTSTR)argument); 63 return PrintUsage(); 64 } 65 } 66 if (argc < 2 || !wcslen(gCommandLine) ) { 67 printf("ERROR: executable path is required\n"); 68 return PrintUsage(); 69 } 70 return S_OK; 71 } 72 73 HRESULT PrintUsage() 74 { 75 printf("record-memory-win --exe EXE_PATH\n"); 76 printf(" Launch an executable and print the memory usage (in Private Bytes)\n"); 77 printf(" of the process.\n\n"); 78 printf("Usage:\n"); 79 printf("-h [--help] : Print usage\n"); 80 printf("--exe arg : Launch specified image. Required\n"); 81 printf("-i [--interval] arg : Print memory usage every arg seconds. Default: 5 seconds\n"); 82 printf("-d [--duration] arg : Run for up to arg seconds. Default: no limit\n\n"); 83 printf("Examples:\n"); 84 printf(" record-memory-win --exe \"C:\\Program Files\\Safari\\Safari.exe\"\n"); 85 printf(" record-memory-win --exe Safari.exe -i 10 -d 7200\n"); 86 return E_FAIL; 87 } 88 89 void UseImage(void (functionForQueryType(HANDLE))) 90 { 91 STARTUPINFO si = {0}; 92 si.cb = sizeof(STARTUPINFO); 93 PROCESS_INFORMATION pi = {0}; 94 95 // Start the child process. 96 if(!CreateProcess( NULL, // No module name (use command line) 97 gCommandLine, // Command line 98 NULL, // Process handle not inheritable 99 NULL, // Thread handle not inheritable 100 FALSE, // Set handle inheritance to FALSE 101 0, // No creation flags 102 NULL, // Use parent's environment block 103 NULL, // Use parent's starting directory 104 &si, // Pointer to STARTUPINFO structure 105 &pi )) // Pointer to PROCESS_INFORMATION structure 106 printf("CreateProcess failed (%d)\n", GetLastError()); 107 else { 108 printf("Created process\n"); 109 functionForQueryType(pi.hProcess); 110 // Close process and thread handles. 111 CloseHandle( pi.hProcess ); 112 CloseHandle( pi.hThread ); 113 } 114 } 115 116 void QueryContinuously(HANDLE hProcess) 117 { 118 Sleep(2000); // give the process some time to launch 119 bool pastDuration = false; 120 time_t startTime = time(NULL); 121 unsigned int memUsage = gSingleProcess ? OneQuery(hProcess) : OneQueryMP(hProcess); 122 while(memUsage && !pastDuration) { 123 printf( "%u\n", memUsage ); 124 Sleep(gQueryInterval*1000); 125 memUsage = gSingleProcess ? OneQuery(hProcess) : OneQueryMP(hProcess); 126 pastDuration = gDuration > 0 ? ElapsedTime(startTime) > gDuration : false; 127 } 128 } 129 130 // returns elapsed time in seconds 131 time_t ElapsedTime(time_t startTime) 132 { 133 time_t currentTime = time(NULL); 134 return currentTime - startTime; 135 } 136 137 // returns Commit Size (Private Bytes) in bytes 138 unsigned int OneQuery(HANDLE hProcess) 139 { 140 PROCESS_MEMORY_COUNTERS_EX pmc; 141 if (NULL == hProcess) 142 return 0; 143 if (GetProcessMemoryInfo(hProcess, (PPROCESS_MEMORY_COUNTERS)&pmc, sizeof(pmc))) 144 return (unsigned)pmc.PrivateUsage; 145 return 0; 146 } 147 148 // returns Commit Size (Private Bytes) in bytes for multi-process executables 149 unsigned int OneQueryMP(HANDLE hProcess) 150 { 151 unsigned int memUsage = 0; 152 TCHAR monitoredProcessName[MAX_PATH]; 153 GetProcessImageFileName(hProcess, monitoredProcessName, sizeof(monitoredProcessName)/sizeof(TCHAR)); 154 LPTSTR shortProcessName = PathFindFileName(monitoredProcessName); 155 DWORD aProcesses[1024], cbNeeded, cProcesses; 156 HANDLE hFoundProcess; 157 if (!EnumProcesses(aProcesses, sizeof(aProcesses), &cbNeeded)) 158 return 0; 159 160 // Calculate how many process identifiers were returned. 161 cProcesses = cbNeeded / sizeof(DWORD); 162 // find existing process 163 for (unsigned int i = 0; i < cProcesses; i++) 164 if (aProcesses[i] != 0) { 165 DWORD retVal = 0; 166 TCHAR foundProcessName[MAX_PATH]; 167 168 // Get a handle to the process. 169 hFoundProcess = OpenProcess(PROCESS_QUERY_INFORMATION | 170 PROCESS_VM_READ, 171 FALSE, aProcesses[i]); 172 173 // Get the process name. 174 if (NULL != hFoundProcess) { 175 HMODULE hMod; 176 DWORD cbNeeded; 177 178 if (EnumProcessModules(hFoundProcess, &hMod, sizeof(hMod), &cbNeeded)) { 179 GetModuleBaseName(hFoundProcess, hMod, foundProcessName, sizeof(foundProcessName)/sizeof(TCHAR)); 180 if (wcsstr(foundProcessName, shortProcessName)) 181 memUsage += OneQuery(hFoundProcess); 182 } 183 } 184 CloseHandle(hFoundProcess); 185 } 186 return memUsage; 187 } 188
