Implement a tool called link_limiter to impose a global restriction on how many

tools/win/link_limiter/limiter.cc

+// 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
+                     const_cast<TCHAR*>(cmdline.c_str()),
+                     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);
+}
+