| // Copyright (c) 2012 The Chromium Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| #include <stdio.h> |
| #include <stdlib.h> |
| |
| #define NOMINMAX |
| #include <windows.h> |
| |
| #include <algorithm> |
| #include <iterator> |
| #include <sstream> |
| #include <string> |
| #include <vector> |
| |
| typedef std::basic_string<TCHAR> tstring; |
| |
| namespace { |
| const bool g_is_debug = (_wgetenv(L"LIMITER_DEBUG") != NULL); |
| } |
| |
| // Don't use stderr for errors because VS has large buffers on them, leading |
| // to confusing error output. |
| static void Error(const wchar_t* msg, ...) { |
| tstring new_msg = tstring(L"limiter fatal error: ") + msg + L"\n"; |
| va_list args; |
| va_start(args, msg); |
| vwprintf(new_msg.c_str(), args); |
| va_end(args); |
| } |
| |
| static void Warn(const wchar_t* msg, ...) { |
| if (!g_is_debug) |
| return; |
| tstring new_msg = tstring(L"limiter warning: ") + msg + L"\n"; |
| va_list args; |
| va_start(args, msg); |
| vwprintf(new_msg.c_str(), args); |
| va_end(args); |
| } |
| |
| static tstring ErrorMessageToString(DWORD err) { |
| TCHAR* msg_buf = NULL; |
| DWORD rc = FormatMessage( |
| FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM, |
| NULL, // lpSource |
| err, |
| MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT), |
| reinterpret_cast<LPTSTR>(&msg_buf), |
| 0, // nSize |
| NULL); // Arguments |
| if (!rc) |
| return L"unknown error"; |
| tstring ret(msg_buf); |
| LocalFree(msg_buf); |
| return ret; |
| } |
| |
| static DWORD RunExe(const tstring& exe_name) { |
| STARTUPINFO startup_info = { sizeof(STARTUPINFO) }; |
| PROCESS_INFORMATION process_info; |
| DWORD exit_code; |
| |
| GetStartupInfo(&startup_info); |
| tstring cmdline = tstring(GetCommandLine()); |
| |
| size_t first_space = cmdline.find(' '); |
| if (first_space == -1) { |
| // I'm not sure why this would ever happen, but just in case... |
| cmdline = exe_name; |
| } else { |
| cmdline = exe_name + cmdline.substr(first_space); |
| } |
| |
| if (!CreateProcess(NULL, // lpApplicationName |
| &cmdline[0], |
| NULL, // lpProcessAttributes |
| NULL, // lpThreadAttributes |
| TRUE, // bInheritHandles |
| 0, // dwCreationFlags, |
| NULL, // lpEnvironment, |
| NULL, // lpCurrentDirectory, |
| &startup_info, |
| &process_info)) { |
| Error(L"Error in CreateProcess[%s]: %s", |
| cmdline.c_str(), ErrorMessageToString(GetLastError()).c_str()); |
| return MAXDWORD; |
| } |
| CloseHandle(process_info.hThread); |
| WaitForSingleObject(process_info.hProcess, INFINITE); |
| GetExitCodeProcess(process_info.hProcess, &exit_code); |
| CloseHandle(process_info.hProcess); |
| return exit_code; |
| } |
| |
| // Returns 0 if there was an error |
| static int CpuConcurrencyMetric(const tstring& envvar_name) { |
| int max_concurrent = 0; |
| std::vector<char> buffer(1); |
| BOOL ok = false; |
| DWORD last_error = 0; |
| do { |
| DWORD bufsize = buffer.size(); |
| ok = GetLogicalProcessorInformation( |
| reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION>(&buffer[0]), |
| &bufsize); |
| last_error = GetLastError(); |
| if (!ok && last_error == ERROR_INSUFFICIENT_BUFFER && |
| bufsize > buffer.size()) { |
| buffer.resize(bufsize); |
| } |
| } while (!ok && last_error == ERROR_INSUFFICIENT_BUFFER); |
| |
| if (!ok) { |
| Warn(L"Error while getting number of cores. Try setting the " |
| L" environment variable '%s' to (num_cores - 1): %s", |
| envvar_name.c_str(), ErrorMessageToString(last_error).c_str()); |
| return 0; |
| } |
| |
| PSYSTEM_LOGICAL_PROCESSOR_INFORMATION pproc_info = |
| reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION>(&buffer[0]); |
| int num_entries = buffer.size() / |
| sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION); |
| |
| for (int i = 0; i < num_entries; ++i) { |
| SYSTEM_LOGICAL_PROCESSOR_INFORMATION& info = pproc_info[i]; |
| if (info.Relationship == RelationProcessorCore) { |
| ++max_concurrent; |
| } |
| } |
| |
| // Leave one core for other tasks |
| return max_concurrent - 1; |
| } |
| |
| // TODO(defaults): Create a better heuristic than # of CPUs. It seems likely |
| // that the right value will, in fact, be based on the memory capacity of the |
| // machine, not on the number of CPUs. |
| enum ConcurrencyMetricEnum { |
| CONCURRENCY_METRIC_ONE, |
| CONCURRENCY_METRIC_CPU, |
| CONCURRENCY_METRIC_DEFAULT = CONCURRENCY_METRIC_CPU |
| }; |
| |
| static int GetMaxConcurrency(const tstring& base_pipename, |
| ConcurrencyMetricEnum metric) { |
| static int max_concurrent = -1; |
| |
| if (max_concurrent == -1) { |
| tstring envvar_name = base_pipename + L"_MAXCONCURRENCY"; |
| |
| const LPTSTR max_concurrent_str = _wgetenv(envvar_name.c_str()); |
| max_concurrent = max_concurrent_str ? _wtoi(max_concurrent_str) : 0; |
| |
| if (max_concurrent == 0) { |
| switch (metric) { |
| case CONCURRENCY_METRIC_CPU: |
| max_concurrent = CpuConcurrencyMetric(envvar_name); |
| if (max_concurrent) |
| break; |
| // else fall through |
| case CONCURRENCY_METRIC_ONE: |
| max_concurrent = 1; |
| break; |
| } |
| } |
| |
| max_concurrent = std::min(std::max(max_concurrent, 1), |
| PIPE_UNLIMITED_INSTANCES); |
| } |
| |
| return max_concurrent; |
| } |
| |
| static HANDLE WaitForPipe(const tstring& pipename, |
| HANDLE event, |
| int max_concurrency) { |
| // We're using a named pipe instead of a semaphore so the Kernel can clean up |
| // after us if we crash while holding onto the pipe (A real semaphore will |
| // not release on process termination). |
| HANDLE pipe = INVALID_HANDLE_VALUE; |
| for (;;) { |
| pipe = CreateNamedPipe( |
| pipename.c_str(), |
| PIPE_ACCESS_DUPLEX, |
| PIPE_TYPE_BYTE, |
| max_concurrency, |
| 1, // nOutBufferSize |
| 1, // nInBufferSize |
| 0, // nDefaultTimeOut |
| NULL); // Default security attributes (noinherit) |
| if (pipe != INVALID_HANDLE_VALUE) |
| break; |
| |
| DWORD error = GetLastError(); |
| if (error == ERROR_PIPE_BUSY) { |
| if (event) { |
| WaitForSingleObject(event, 60 * 1000 /* ms */); |
| } else { |
| // TODO(iannucci): Maybe we should error out here instead of falling |
| // back to a sleep-poll |
| Sleep(5 * 1000 /* ms */); |
| } |
| } else { |
| Warn(L"Got error %d while waiting for pipe: %s", error, |
| ErrorMessageToString(error).c_str()); |
| return INVALID_HANDLE_VALUE; |
| } |
| } |
| |
| return pipe; |
| } |
| |
| static int WaitAndRun(const tstring& shimmed_exe, |
| const tstring& base_pipename) { |
| ULONGLONG start_time = 0, end_time = 0; |
| tstring pipename = L"\\\\.\\pipe\\" + base_pipename; |
| tstring event_name = L"Local\\EVENT_" + base_pipename; |
| |
| // This event lets us do better than strict polling, but we don't rely on it |
| // (in case a process crashes before signalling the event). |
| HANDLE event = CreateEvent( |
| NULL, // Default security attributes |
| FALSE, // Manual reset |
| FALSE, // Initial state |
| event_name.c_str()); |
| |
| if (g_is_debug) |
| start_time = GetTickCount64(); |
| |
| HANDLE pipe = |
| WaitForPipe(pipename, event, |
| GetMaxConcurrency(base_pipename, CONCURRENCY_METRIC_DEFAULT)); |
| |
| if (g_is_debug) { |
| end_time = GetTickCount64(); |
| wprintf(L" took %.2fs to acquire semaphore.\n", |
| (end_time - start_time) / 1000.0); |
| } |
| |
| DWORD ret = RunExe(shimmed_exe); |
| |
| if (pipe != INVALID_HANDLE_VALUE) |
| CloseHandle(pipe); |
| if (event != NULL) |
| SetEvent(event); |
| |
| return ret; |
| } |
| |
| void Usage(const tstring& msg) { |
| tstring usage(msg); |
| usage += L"\n" |
| L"Usage: SHIMED_NAME__SEMAPHORE_NAME\n" |
| L"\n" |
| L" SHIMMED_NAME - ex. 'link.exe' or 'lib.exe'\n" |
| L" - can be exe, bat, or com\n" |
| L" - must exist in PATH\n" |
| L"\n" |
| L" SEMAPHORE_NAME - ex. 'SOME_NAME' or 'GROOVY_SEMAPHORE'\n" |
| L"\n" |
| L" Example:\n" |
| L" link.exe__LINK_LIMITER.exe\n" |
| L" lib.exe__LINK_LIMITER.exe\n" |
| L" * Both will limit on the same semaphore\n" |
| L"\n" |
| L" link.exe__LINK_LIMITER.exe\n" |
| L" lib.exe__LIB_LIMITER.exe\n" |
| L" * Both will limit on independent semaphores\n" |
| L"\n" |
| L" This program is meant to be run after renaming it into the\n" |
| L" above format. Once you have done so, executing it will block\n" |
| L" on the availability of the semaphore SEMAPHORE_NAME. Once\n" |
| L" the semaphore is obtained, it will execute SHIMMED_NAME, \n" |
| L" passing through all arguments as-is.\n" |
| L"\n" |
| L" The maximum concurrency can be manually set by setting the\n" |
| L" environment variable <SEMAPHORE_NAME>_MAXCONCURRENCY to an\n" |
| L" integer value (1, 254).\n" |
| L" * This value must be set the same for ALL invocations.\n" |
| L" * If the value is not set, it defaults to (num_cores-1).\n" |
| L"\n" |
| L" The semaphore is automatically released when the program\n" |
| L" completes normally, OR if the program crashes (or even if\n" |
| L" limiter itself crashes).\n"; |
| Error(usage.c_str()); |
| exit(-1); |
| } |
| |
| // Input command line is assumed to be of the form: |
| // |
| // thing.exe__PIPE_NAME.exe ... |
| // |
| // Specifically, wait for a semaphore (whose concurrency is specified by |
| // LIMITER_MAXCONCURRENT), and then pass through everything once we have |
| // acquired the semaphore. |
| // |
| // argv[0] is parsed for: |
| // * exe_to_shim_including_extension.exe |
| // * This could also be a bat or com. Anything that CreateProcess will |
| // accept. |
| // * "__" |
| // * We search for this separator from the end of argv[0], so the exe name |
| // could contain a double underscore if necessary. |
| // * PIPE_NAME |
| // * Can only contain single underscores, not a double underscore. |
| // * i.e. HELLO_WORLD_PIPE will work, but HELLO__WORLD_PIPE will not. |
| // * This would allow the shimmed exe to contain arbitrary numbers of |
| // underscores. We control the pipe name, but not necessarily the thing |
| // we're shimming. |
| // |
| int wmain(int, wchar_t** argv) { |
| tstring shimmed_plus_pipename = argv[0]; |
| size_t last_slash = shimmed_plus_pipename.find_last_of(L"/\\"); |
| if (last_slash != tstring::npos) { |
| shimmed_plus_pipename = shimmed_plus_pipename.substr(last_slash + 1); |
| } |
| |
| size_t separator = shimmed_plus_pipename.rfind(L"__"); |
| if (separator == tstring::npos) { |
| Usage(L"Cannot parse argv[0]. No '__' found. " |
| L"Should be like '[...(\\|/)]link.exe__PIPE_NAME.exe'"); |
| } |
| tstring shimmed_exe = shimmed_plus_pipename.substr(0, separator); |
| tstring base_pipename = shimmed_plus_pipename.substr(separator + 2); |
| |
| size_t dot = base_pipename.find(L'.'); |
| if (dot == tstring::npos) { |
| Usage(L"Expected an executable extension in argv[0]. No '.' found."); |
| } |
| base_pipename = base_pipename.substr(0, dot); |
| |
| return WaitAndRun(shimmed_exe, base_pipename); |
| } |
| |