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#include <stdio.h>
6#include <stdlib.h>
7
8#define NOMINMAX
9#include <windows.h>
10
11#include <algorithm>
12#include <iterator>
13#include <sstream>
14#include <string>
15#include <vector>
16
17typedef std::basic_string<TCHAR> tstring;
18
19namespace {
20  const bool g_is_debug = (_wgetenv(L"LIMITER_DEBUG") != NULL);
21}
22
23// Don't use stderr for errors because VS has large buffers on them, leading
24// to confusing error output.
25static void Error(const wchar_t* msg, ...) {
26  tstring new_msg = tstring(L"limiter fatal error: ") + msg + L"\n";
27  va_list args;
28  va_start(args, msg);
29  vwprintf(new_msg.c_str(), args);
30  va_end(args);
31}
32
33static void Warn(const wchar_t* msg, ...) {
34  if (!g_is_debug)
35    return;
36  tstring new_msg = tstring(L"limiter warning: ") + msg + L"\n";
37  va_list args;
38  va_start(args, msg);
39  vwprintf(new_msg.c_str(), args);
40  va_end(args);
41}
42
43static tstring ErrorMessageToString(DWORD err) {
44  TCHAR* msg_buf = NULL;
45  DWORD rc = FormatMessage(
46      FORMAT_MESSAGE_ALLOCATE_BUFFER | FORMAT_MESSAGE_FROM_SYSTEM,
47      NULL,   // lpSource
48      err,
49      MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
50      reinterpret_cast<LPTSTR>(&msg_buf),
51      0,      // nSize
52      NULL);  // Arguments
53  if (!rc)
54    return L"unknown error";
55  tstring ret(msg_buf);
56  LocalFree(msg_buf);
57  return ret;
58}
59
60static DWORD RunExe(const tstring& exe_name) {
61  STARTUPINFO startup_info = { sizeof(STARTUPINFO) };
62  PROCESS_INFORMATION process_info;
63  DWORD exit_code;
64
65  GetStartupInfo(&startup_info);
66  tstring cmdline = tstring(GetCommandLine());
67
68  size_t first_space = cmdline.find(' ');
69  if (first_space == -1) {
70    // I'm not sure why this would ever happen, but just in case...
71    cmdline = exe_name;
72  } else {
73    cmdline = exe_name  + cmdline.substr(first_space);
74  }
75
76  if (!CreateProcess(NULL,  // lpApplicationName
77                     const_cast<TCHAR*>(cmdline.c_str()),
78                     NULL,  // lpProcessAttributes
79                     NULL,  // lpThreadAttributes
80                     TRUE,  // bInheritHandles
81                     0,     // dwCreationFlags,
82                     NULL,  // lpEnvironment,
83                     NULL,  // lpCurrentDirectory,
84                     &startup_info,
85                     &process_info)) {
86    Error(L"Error in CreateProcess[%s]: %s",
87          cmdline.c_str(), ErrorMessageToString(GetLastError()).c_str());
88    return MAXDWORD;
89  }
90  CloseHandle(process_info.hThread);
91  WaitForSingleObject(process_info.hProcess, INFINITE);
92  GetExitCodeProcess(process_info.hProcess, &exit_code);
93  CloseHandle(process_info.hProcess);
94  return exit_code;
95}
96
97// Returns 0 if there was an error
98static int CpuConcurrencyMetric(const tstring& envvar_name) {
99  int max_concurrent = 0;
100  std::vector<char> buffer(1);
101  BOOL ok = false;
102  DWORD last_error = 0;
103  do {
104    DWORD bufsize = buffer.size();
105    ok = GetLogicalProcessorInformation(
106        reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION>(&buffer[0]),
107        &bufsize);
108    last_error = GetLastError();
109    if (!ok && last_error == ERROR_INSUFFICIENT_BUFFER &&
110        bufsize > buffer.size()) {
111      buffer.resize(bufsize);
112    }
113  } while (!ok && last_error == ERROR_INSUFFICIENT_BUFFER);
114
115  if (!ok) {
116    Warn(L"Error while getting number of cores. Try setting the "
117         L" environment variable '%s'  to (num_cores - 1): %s",
118         envvar_name.c_str(), ErrorMessageToString(last_error).c_str());
119    return 0;
120  }
121
122  PSYSTEM_LOGICAL_PROCESSOR_INFORMATION pproc_info =
123      reinterpret_cast<PSYSTEM_LOGICAL_PROCESSOR_INFORMATION>(&buffer[0]);
124  int num_entries = buffer.size() /
125      sizeof(SYSTEM_LOGICAL_PROCESSOR_INFORMATION);
126
127  for (int i = 0; i < num_entries; ++i) {
128    SYSTEM_LOGICAL_PROCESSOR_INFORMATION& info = pproc_info[i];
129    if (info.Relationship == RelationProcessorCore) {
130      ++max_concurrent;
131    }
132  }
133
134  // Leave one core for other tasks
135  return max_concurrent - 1;
136}
137
138// TODO(defaults): Create a better heuristic than # of CPUs. It seems likely
139// that the right value will, in fact, be based on the memory capacity of the
140// machine, not on the number of CPUs.
141enum ConcurrencyMetricEnum {
142  CONCURRENCY_METRIC_ONE,
143  CONCURRENCY_METRIC_CPU,
144  CONCURRENCY_METRIC_DEFAULT = CONCURRENCY_METRIC_CPU
145};
146
147static int GetMaxConcurrency(const tstring& base_pipename,
148                             ConcurrencyMetricEnum metric) {
149  static int max_concurrent = -1;
150
151  if (max_concurrent == -1) {
152    tstring envvar_name = base_pipename + L"_MAXCONCURRENCY";
153
154    const LPTSTR max_concurrent_str = _wgetenv(envvar_name.c_str());
155    max_concurrent = max_concurrent_str ? _wtoi(max_concurrent_str) : 0;
156
157    if (max_concurrent == 0) {
158      switch (metric) {
159        case CONCURRENCY_METRIC_CPU:
160          max_concurrent = CpuConcurrencyMetric(envvar_name);
161          if (max_concurrent)
162            break;
163          // else fall through
164        case CONCURRENCY_METRIC_ONE:
165          max_concurrent = 1;
166          break;
167      }
168    }
169
170    max_concurrent = std::min(std::max(max_concurrent, 1),
171                              PIPE_UNLIMITED_INSTANCES);
172  }
173
174  return max_concurrent;
175}
176
177static HANDLE WaitForPipe(const tstring& pipename,
178                          HANDLE event,
179                          int max_concurrency) {
180  // We're using a named pipe instead of a semaphore so the Kernel can clean up
181  // after us if we crash while holding onto the pipe (A real semaphore will
182  // not release on process termination).
183  HANDLE pipe = INVALID_HANDLE_VALUE;
184  for (;;) {
185    pipe = CreateNamedPipe(
186        pipename.c_str(),
187        PIPE_ACCESS_DUPLEX,
188        PIPE_TYPE_BYTE,
189        max_concurrency,
190        1,      // nOutBufferSize
191        1,      // nInBufferSize
192        0,      // nDefaultTimeOut
193        NULL);  // Default security attributes (noinherit)
194    if (pipe != INVALID_HANDLE_VALUE)
195      break;
196
197    DWORD error = GetLastError();
198    if (error == ERROR_PIPE_BUSY) {
199      if (event) {
200        WaitForSingleObject(event, 60 * 1000 /* ms */);
201      } else {
202        // TODO(iannucci): Maybe we should error out here instead of falling
203        // back to a sleep-poll
204        Sleep(5 * 1000 /* ms */);
205      }
206    } else {
207      Warn(L"Got error %d while waiting for pipe: %s", error,
208           ErrorMessageToString(error).c_str());
209      return INVALID_HANDLE_VALUE;
210    }
211  }
212
213  return pipe;
214}
215
216static int WaitAndRun(const tstring& shimmed_exe,
217                      const tstring& base_pipename) {
218  ULONGLONG start_time = 0, end_time = 0;
219  tstring pipename = L"\\\\.\\pipe\\" + base_pipename;
220  tstring event_name = L"Local\\EVENT_" + base_pipename;
221
222  // This event lets us do better than strict polling, but we don't rely on it
223  // (in case a process crashes before signalling the event).
224  HANDLE event = CreateEvent(
225      NULL,   // Default security attributes
226      FALSE,  // Manual reset
227      FALSE,  // Initial state
228      event_name.c_str());
229
230  if (g_is_debug)
231    start_time = GetTickCount64();
232
233  HANDLE pipe =
234    WaitForPipe(pipename, event,
235                GetMaxConcurrency(base_pipename, CONCURRENCY_METRIC_DEFAULT));
236
237  if (g_is_debug) {
238    end_time = GetTickCount64();
239    wprintf(L"  took %.2fs to acquire semaphore.\n",
240        (end_time - start_time) / 1000.0);
241  }
242
243  DWORD ret = RunExe(shimmed_exe);
244
245  if (pipe != INVALID_HANDLE_VALUE)
246    CloseHandle(pipe);
247  if (event != NULL)
248    SetEvent(event);
249
250  return ret;
251}
252
253void Usage(const tstring& msg) {
254  tstring usage(msg);
255  usage += L"\n"
256           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(usage.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//
314int 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'.');
330  if (dot == tstring::npos) {
331    Usage(L"Expected an executable extension in argv[0]. No '.' found.");
332  }
333  base_pipename = base_pipename.substr(0, dot);
334
335  return WaitAndRun(shimmed_exe, base_pipename);
336}
337
338