tools/win/link_limiter/limiter.cc

   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
  17 typedef std::basic_string<TCHAR> tstring;
  18
  19 namespace {
  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.
  25 static 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
  33 static 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
  43 static 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
  60 static 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
  98 static 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.
 141 enum ConcurrencyMetricEnum {
 142   CONCURRENCY_METRIC_ONE,
 143   CONCURRENCY_METRIC_CPU,
 144   CONCURRENCY_METRIC_DEFAULT = CONCURRENCY_METRIC_CPU
 145 };
 146
 147 static 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
 177 static 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
 216 static 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
 253 void 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 //
 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'.');
 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