Tools/record-memory-win/main.cpp

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