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