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
12bool gSingleProcess = true;
13int gQueryInterval = 5; // seconds
14time_t gDuration = 0;   // seconds
15LPTSTR gCommandLine;
16
17HRESULT ProcessArgs(int argc, TCHAR *argv[]);
18HRESULT PrintUsage();
19void UseImage(void (functionForQueryType(HANDLE)));
20void QueryContinuously(HANDLE hProcess);
21time_t ElapsedTime(time_t startTime);
22unsigned int OneQuery(HANDLE hProcess);
23unsigned int OneQueryMP(HANDLE hProcess);
24
25int __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
35HRESULT 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
73HRESULT 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
89void 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
116void 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
131time_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
138unsigned 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
149unsigned 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