OLD | NEW |
---|---|
(Empty) | |
1 // Copyright (c) 2012 The Chromium Authors. All rights reserved. | |
2 // Use of this source code is governed by a BSD-style license that can be | |
3 // found in the LICENSE file. | |
4 | |
5 #define NOMINMAX | |
6 #include <windows.h> | |
7 | |
8 #include <stdio.h> | |
9 #include <stdlib.h> | |
10 | |
11 #include <algorithm> | |
12 #include <iterator> | |
13 #include <string> | |
14 #include <sstream> | |
15 #include <vector> | |
cmp
2012/08/01 23:36:42
alphanum sort lines 11-15
| |
16 | |
17 typedef std::basic_string<TCHAR> tstring; | |
18 | |
19 const bool g_is_debug = (_wgetenv(L"LIMITER_DEBUG") != NULL); | |
20 const int g_wait_time = 30 * 1000; // 30 seconds | |
cmp
2012/08/01 23:36:42
probably should remove one of the extra spaces aft
| |
21 | |
22 // Don't use stderr for errors because VS has large buffers on them, leading | |
23 // to confusing error output. | |
24 static void Error(const tstring& msg, ...) { | |
25 tstring new_msg = L"limiter fatal error: "+msg+L"\n"; | |
26 va_list args; | |
27 va_start(args, msg); | |
28 vwprintf(new_msg.c_str(), args); | |
29 va_end(args); | |
30 } | |
31 static void Warn(const tstring& msg, ...) { | |
cmp
2012/08/01 23:36:42
insert an empty line before line 31
| |
32 if(!g_is_debug) | |
cmp
2012/08/01 23:36:42
one space after if
| |
33 return; | |
34 tstring new_msg = L"limiter warning: "+msg+L"\n"; | |
35 va_list args; | |
36 va_start(args, msg); | |
37 vwprintf(new_msg.c_str(), args); | |
38 va_end(args); | |
39 } | |
40 | |
41 static tstring ErrorMessageToString(DWORD err) { | |
42 TCHAR* msg_buf = NULL; | |
43 DWORD rc = FormatMessage( | |
44 FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, | |
45 NULL, | |
46 err, | |
47 MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), | |
48 reinterpret_cast<LPTSTR>(&msg_buf), | |
49 0, | |
50 NULL); | |
51 if (!rc) | |
52 return L"unknown error"; | |
53 tstring ret(msg_buf); | |
54 LocalFree(msg_buf); | |
55 return ret; | |
56 } | |
57 | |
58 static DWORD RunExe(const tstring& exe_name) { | |
59 STARTUPINFO startup_info = { sizeof(STARTUPINFO) }; | |
60 PROCESS_INFORMATION process_info; | |
61 DWORD exit_code; | |
62 | |
63 GetStartupInfo(&startup_info); | |
64 tstring cmdline = tstring(GetCommandLine()); | |
65 | |
66 size_t first_space = cmdline.find(' '); | |
67 if (first_space == -1) { | |
68 // I'm not sure why this would ever happen, but just in case... | |
69 cmdline = exe_name; | |
70 } else { | |
71 cmdline = exe_name + cmdline.substr(first_space+1); | |
72 } | |
73 | |
74 if (!CreateProcess(NULL, | |
75 const_cast<TCHAR*>(cmdline.c_str()), | |
76 NULL, | |
77 NULL, | |
78 TRUE, | |
79 0, | |
80 NULL, | |
81 NULL, | |
82 &startup_info, &process_info)) { | |
83 Error(L"Error in CreateProcess[%s]: %s", | |
84 cmdline, ErrorMessageToString(GetLastError()).c_str()); | |
85 return MAXDWORD; | |
86 } | |
87 CloseHandle(process_info.hThread); | |
88 WaitForSingleObject(process_info.hProcess, INFINITE); | |
89 GetExitCodeProcess(process_info.hProcess, &exit_code); | |
90 CloseHandle(process_info.hProcess); | |
91 return exit_code; | |
92 } | |
93 | |
94 // Returns 0 if there was an error | |
95 static int CpuConcurrencyMetric(const tstring& envvar_name) { | |
96 int max_concurrent = 0; | |
97 std::vector<char> buffer(1); | |
98 BOOL ok = false; | |
99 DWORD last_error = 0; | |
100 do { | |
101 DWORD bufsize = buffer.size(); | |
102 ok = GetLogicalProcessorInformation( | |
103 reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION>(&buffer[0]), | |
104 &bufsize); | |
105 last_error = GetLastError(); | |
106 if (!ok && last_error == ERROR_INSUFFICIENT_BUFFER && | |
107 bufsize > buffer.size()) | |
cmp
2012/08/01 23:36:42
this indent should line the first b in bufsize und
| |
108 { | |
109 buffer.resize(bufsize); | |
110 } | |
111 } while(!ok && last_error == ERROR_INSUFFICIENT_BUFFER); | |
cmp
2012/08/01 23:36:42
insert a space after while
| |
112 | |
113 if (!ok) { | |
114 Warn(L"Error while getting number of cores. Try setting the " | |
115 L" environment variable '%s' to (num_cores-1): %s", | |
116 envvar_name.c_str(), ErrorMessageToString(last_error).c_str()); | |
117 return 0; | |
118 } | |
119 | |
120 PSYSTEM_LOGICAL_PROCESSOR_INFORMATION pproc_info = | |
121 reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION>(&buffer[0]); | |
cmp
2012/08/01 23:36:42
increase indent on line 121 by 2 spaces
| |
122 int num_entries = buffer.size() / | |
123 sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION); | |
cmp
2012/08/01 23:36:42
increase indent on line 123 by 2 spaces
| |
124 | |
125 for(int i = 0; i < num_entries; ++i) { | |
cmp
2012/08/01 23:36:42
one space after for
| |
126 SYSTEM_LOGICAL_PROCESSOR_INFORMATION &info = pproc_info[i]; | |
127 if(info.Relationship == RelationProcessorCore) { | |
cmp
2012/08/01 23:36:42
one space after if
| |
128 max_concurrent += 1; | |
129 } | |
130 } | |
131 | |
132 // Leave one core for other tasks | |
133 return max_concurrent-1; | |
134 } | |
135 | |
136 | |
137 // TODO(defaults): Create a better heuristic than # of CPUs. It seems likely | |
138 // that the right value will, in fact, be based on the memory capacity of the | |
139 // machine, not on the number of CPUs. | |
140 enum ConcurrencyMetricEnum { | |
141 CONCURRENCY_METRIC_ONE, | |
142 CONCURRENCY_METRIC_CPU, | |
143 CONCURRENCY_METRIC_DEFAULT = CONCURRENCY_METRIC_CPU | |
144 }; | |
145 | |
146 static int GetMaxConcurrency( | |
147 const tstring& base_pipename, | |
148 ConcurrencyMetricEnum metric = CONCURRENCY_METRIC_DEFAULT) | |
149 { | |
150 static int max_concurrent = -1; | |
151 | |
152 if (max_concurrent == -1) { | |
cmp
2012/08/01 23:36:42
could max_concurrent ever not equal -1? i'm think
iannucci1
2012/08/01 23:52:39
It's static. Multiple calls will get the cached va
cmp
2012/08/02 00:04:23
Never mind, then.
| |
153 tstring envvar_name = base_pipename+L"_MAXCONCURRENCY"; | |
154 | |
155 const LPTSTR max_concurrent_str = _wgetenv(envvar_name.c_str()); | |
156 max_concurrent = _wtoi(max_concurrent_str ? max_concurrent_str : L"0"); | |
157 | |
158 if (max_concurrent == 0) { | |
159 switch(metric) { | |
cmp
2012/08/01 23:36:42
one space after switch
| |
160 case CONCURRENCY_METRIC_CPU: | |
cmp
2012/08/01 23:36:42
each case in a switch should be indented by 2 spac
| |
161 max_concurrent = CpuConcurrencyMetric(envvar_name); | |
162 if(max_concurrent) | |
cmp
2012/08/01 23:36:42
one space after if
| |
163 break; | |
164 // else fallthrough | |
165 case CONCURRENCY_METRIC_ONE: | |
166 max_concurrent = 1; | |
167 break; | |
168 } | |
169 } | |
170 | |
171 max_concurrent = std::min(std::max(max_concurrent, 1), | |
172 PIPE_UNLIMITED_INSTANCES); | |
173 } | |
174 | |
175 return max_concurrent; | |
176 } | |
177 | |
178 | |
cmp
2012/08/01 23:36:42
let's be consistent about 2 empty lines between fu
| |
179 static HANDLE wait_for_pipe(const tstring& pipename, | |
180 HANDLE hEvent, | |
181 int max_concurrency) | |
182 { | |
183 // We're using a named pipe instead of a semaphore so the Kernel can clean up | |
184 // after us if we crash while holding onto the pipe (A real semaphore will | |
185 // not release on process termination). | |
186 HANDLE hPipe = INVALID_HANDLE_VALUE; | |
cmp
2012/08/01 23:36:42
i believe the 'no hungarian notation' rule applies
| |
187 for(;;) { | |
cmp
2012/08/01 23:36:42
one space after for
| |
188 hPipe = CreateNamedPipe( | |
189 pipename.c_str(), | |
cmp
2012/08/01 23:36:42
these should be indented by an extra 2 spaces per
| |
190 PIPE_ACCESS_DUPLEX, | |
191 PIPE_TYPE_BYTE, | |
192 max_concurrency, | |
193 1, // nOutBufferSize | |
194 1, // nInBufferSize | |
195 0, // nDefaultTimeOut | |
196 NULL // Default security attributes (noinherit) | |
197 ); | |
cmp
2012/08/01 23:36:42
i think this line should be de-indented by 2 space
| |
198 if (hPipe != INVALID_HANDLE_VALUE) | |
199 break; | |
200 | |
201 DWORD error = GetLastError(); | |
202 if (error == ERROR_PIPE_BUSY) { | |
203 if (hEvent) { | |
204 WaitForSingleObject(hEvent, g_wait_time); | |
205 } else { | |
206 Sleep(g_wait_time); | |
207 } | |
208 } else { | |
209 Warn(L"Got error %d while waiting for pipe: %s", error, | |
210 ErrorMessageToString(error).c_str()); | |
211 return INVALID_HANDLE_VALUE; | |
212 } | |
213 } | |
214 | |
215 return hPipe; | |
216 } | |
217 | |
218 static int wait_and_run(const tstring& shimmed_exe, | |
219 const tstring& base_pipename) | |
220 { | |
221 ULONGLONG start_time = 0, end_time = 0; | |
222 tstring pipename = L"\\\\.\\pipe\\"+base_pipename; | |
223 tstring event_name = L"Local\\EVENT_"+base_pipename; | |
224 | |
225 // This event lets us do better than strict polling, but we don't rely on it | |
226 // (in case a process crashes before signalling the event). | |
227 HANDLE hEvent = CreateEvent( | |
228 NULL, // Default security attributes | |
229 FALSE, // Manual reset | |
230 FALSE, // Initial state | |
231 event_name.c_str()); | |
232 | |
233 if (g_is_debug) | |
234 start_time = GetTickCount64(); | |
235 | |
236 HANDLE hPipe = | |
cmp
2012/08/01 23:36:42
hungarian notation should not be used (see C++ sty
| |
237 wait_for_pipe(pipename, hEvent, GetMaxConcurrency(base_pipename)); | |
238 | |
239 if (g_is_debug) { | |
240 end_time = GetTickCount64(); | |
241 wprintf(L" took %.2fs to acquire semaphore.\n", | |
242 (end_time - start_time) / 1000.0); | |
243 } | |
244 | |
245 DWORD ret = RunExe(shimmed_exe); | |
246 | |
247 if (hPipe != INVALID_HANDLE_VALUE) { CloseHandle(hPipe); } | |
248 if (hEvent != NULL) { SetEvent(hEvent); } | |
249 | |
250 return ret; | |
251 } | |
252 | |
253 void usage(const tstring& msg) { | |
254 std::wostringstream strm; | |
255 strm << msg << L"\n"; | |
256 strm << L"Usage: SHIMED_NAME__SEMAPHORE_NAME\n" | |
257 << L"\n" | |
258 << L" SHIMMED_NAME - ex. 'link.exe' or 'lib.exe'\n" | |
259 << L" - can be exe, bat, or com\n" | |
260 << L" - must exist in PATH\n" | |
261 << L"\n" | |
262 << L" SEMAPHORE_NAME - ex. 'SOME_NAME' or 'GROOVY_SEMAPHORE'\n" | |
263 << L"\n" | |
264 << L" Example:\n" | |
265 << L" link.exe__LINK_LIMITER.exe\n" | |
266 << L" lib.exe__LINK_LIMITER.exe\n" | |
267 << L" * Both will limit on the same semaphore\n" | |
268 << L"\n" | |
269 << L" link.exe__LINK_LIMITER.exe\n" | |
270 << L" lib.exe__LIB_LIMITER.exe\n" | |
271 << L" * Both will limit on independent semaphores\n" | |
272 << L"\n" | |
273 << L" This program is meant to be run after renaming it into the\n" | |
274 << L" above format. Once you have done so, executing it will block\n" | |
275 << L" on the availability of the semaphore SEMAPHORE_NAME. Once\n" | |
276 << L" the semaphore is obtained, it will execute SHIMMED_NAME, \n" | |
277 << L" passing through all arguments as-is.\n" | |
278 << L"\n" | |
279 << L" The maximum concurrency can be manually set by setting the\n" | |
280 << L" environment variable <SEMAPHORE_NAME>_MAXCONCURRENCY to an\n" | |
281 << L" integer value (1, 254).\n" | |
282 << L" * This value must be set the same for ALL invocations.\n" | |
283 << L" * If the value is not set, it defaults to (num_cores-1).\n" | |
284 << L"\n" | |
285 << L" The semaphore is automatically released when the program\n" | |
286 << L" completes normally, OR if the program crashes (or even if\n" | |
287 << L" limiter itself crashes).\n"; | |
288 Error(strm.str().c_str()); | |
289 exit(-1); | |
290 } | |
291 | |
292 // Input command line is assumed to be of the form: | |
293 // | |
294 // thing.exe__PIPE_NAME.exe ... | |
295 // | |
296 // Specifically, wait for a semaphore (whose concurrency is specified by | |
297 // LIMITER_MAXCONCURRENT), and then pass through everything once we have | |
298 // acquired the semaphore. | |
299 // | |
300 // argv[0] is parsed for: | |
301 // * exe_to_shim_including_extension.exe | |
302 // * This could also be a bat or com. Anything that CreateProcess will | |
303 // accept. | |
304 // * "__" | |
305 // * We search for this separator from the end of argv[0], so the exe name | |
306 // could contain a double underscore if necessary. | |
307 // * PIPE_NAME | |
308 // * Can only contain single underscores, not a double underscore. | |
309 // * i.e. HELLO_WORLD_PIPE will work, but HELLO__WORLD_PIPE will not. | |
310 // * This would allow the shimmed exe to contain arbitrary numbers of | |
311 // underscores. We control the pipe name, but not necessarily the thing | |
312 // we're shimming. | |
313 // | |
314 int wmain(int, wchar_t** argv) { | |
315 tstring shimmed_plus_pipename = argv[0]; | |
316 size_t last_slash = shimmed_plus_pipename.find_last_of(L"/\\"); | |
317 if (last_slash != tstring::npos) { | |
318 shimmed_plus_pipename = shimmed_plus_pipename.substr(last_slash+1); | |
319 } | |
320 | |
321 size_t separator = shimmed_plus_pipename.rfind(L"__"); | |
322 if (separator == tstring::npos) { | |
323 usage(L"Cannot parse argv[0]. No '__' found. " | |
324 L"Should be like '[...(\\|/)]link.exe__PIPE_NAME.exe'"); | |
325 } | |
326 tstring shimmed_exe = shimmed_plus_pipename.substr(0, separator); | |
327 tstring base_pipename = shimmed_plus_pipename.substr(separator+2); | |
328 | |
329 size_t dot = base_pipename.find(L'.'); | |
cmp
2012/08/01 23:36:42
for consistency, use " here
iannucci1
2012/08/01 23:52:39
The implementation of find is (marginally) more ef
cmp
2012/08/02 00:04:23
Nah, character here, then. Can you document that
| |
330 if (dot == tstring::npos) { | |
331 usage(L"Hunh? No '.' in argv[0]? Are we on windows?"); | |
cmp
2012/08/01 23:36:42
can you make this more clear? (eg. 'ERROR: Requir
iannucci1
2012/08/01 23:52:39
Well... Since argv[0] must be the executable name,
cmp
2012/08/02 00:04:23
SGTM
| |
332 } | |
333 base_pipename = base_pipename.substr(0, dot); | |
334 | |
335 return wait_and_run(shimmed_exe, base_pipename); | |
336 } | |
337 | |
cmp
2012/08/01 23:36:42
nit: remove empty line
| |
OLD | NEW |