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

tools/win/link_limiter/limiter.cpp

+// 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.
+
+#define NOMINMAX
+#include <windows.h>
+
+#include <stdio.h>
+#include <stdlib.h>
+
+#include <algorithm>
+#include <iterator>
+#include <string>
+#include <sstream>
+#include <vector>
+
+#ifndef SHIM_EXE
+# error You must define SHIM_EXE when building (like "link.exe" or "lib.exe")
+#endif
+
+typedef std::basic_string<TCHAR> tstring;
+
+const bool g_is_debug = (_wgetenv(L"LIMITER_DEBUG") != NULL);
+const int  g_wait_time = 30 * 1000; // 30 seconds
+const LPTSTR g_pipe_name = L"\\\\.\\pipe\\LIMITER_PIPE";
+const LPTSTR g_event_name = L"Local\\LIMITER_EVENT";
+const LPTSTR g_envvar_name = L"LIMITER_MAXCONCURRENT";
+
+// Don't use stderr for errors because VS has large buffers on them, leading
+// to confusing error output.
+static void Fatal(const wchar_t* msg) {
+  wprintf(L"limiter(" SHIM_EXE L") fatal error: %s\n", msg);
+  exit(1);
+}
+
+static tstring ErrorMessageToString(DWORD err) {
+  wchar_t* msg_buf = NULL;
+  DWORD rc = FormatMessage(FORMAT_MESSAGE_ALLOCATE_BUFFER |
+                               FORMAT_MESSAGE_FROM_SYSTEM,
+                           NULL,
+                           err,
+                           MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
+                           reinterpret_cast<LPTSTR>(&msg_buf),
+                           0,
+                           NULL);
+  if (!rc)
+    return L"unknown error";
+  tstring ret(msg_buf);
+  LocalFree(msg_buf);
+  return ret;
+}
+
+static DWORD RunExe() {
+  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(' ')+1;
+  if (first_space == -1) {
+    // I'm not sure why this would ever happen, but just in case...
+    cmdline = SHIM_EXE;
+  } else {
+    cmdline = SHIM_EXE L" " + cmdline.substr(first_space);
+  }
+
+  if (!CreateProcess(NULL,
+                     const_cast<TCHAR*>(cmdline.c_str()),
+                     NULL,
+                     NULL,
+                     TRUE,
+                     0,
+                     NULL,
+                     NULL,
+                     &startup_info, &process_info)) {
+    tstring error = ErrorMessageToString(GetLastError());
+    Fatal(error.c_str());
+  }
+  CloseHandle(process_info.hThread);
+  WaitForSingleObject(process_info.hProcess, INFINITE);
+  GetExitCodeProcess(process_info.hProcess, &exit_code);
+  CloseHandle(process_info.hProcess);
+  return exit_code;
+}
+
+int CpuConcurrencyMetric() {
+  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) {
+    if (g_is_debug)
+      wprintf(L"Error while getting number of cores. Try setting the "
+        L" environment variable '%s' to (num_cores-1): %d\n",
+        g_envvar_name, last_error);
+    exit(RunExe());
+  }
+
+  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 += 1;
+    }
+  }
+
+  // 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
+};
+
+int GetMaxConcurrency(
+  ConcurrencyMetricEnum metric = CONCURRENCY_METRIC_DEFAULT)
+{
+  static int max_concurrent = -1;
+
+  if (max_concurrent == -1) {
+    const LPTSTR max_concurrent_str = _wgetenv(g_envvar_name);
+    max_concurrent = _wtoi(max_concurrent_str ? max_concurrent_str : L"0");
+
+    if (max_concurrent == 0) {
+      switch(metric) {
+      case CONCURRENCY_METRIC_ONE:
+        max_concurrent = 1;
+        break;
+      case CONCURRENCY_METRIC_CPU:
+        max_concurrent = CpuConcurrencyMetric();
+        break;
+      }
+    }
+
+    max_concurrent = std::min(std::max(max_concurrent, 1),
+                              PIPE_UNLIMITED_INSTANCES);
+  }
+
+  return max_concurrent;
+}
+
+
+// Input command line is assumed to be of the form:
+//
+// SHIM_EXE ...
+//
+// Specifically, wait for a semaphore (whose concurrency is specified by
+// LIMITER_MAXCONCURRENT), and then pass through everything once we have
+// acquired the semaphore.
+int wmain(int, wchar_t**) {
+  ULONGLONG start_time = 0, end_time = 0;
+
+  if (g_is_debug)
+    start_time = GetTickCount64();
+
+  // 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 hEvent = CreateEvent(
+    NULL, // Default security attributes
+    FALSE, // Manual reset
+    FALSE, // Initial state
+    g_event_name);
+
+  // 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 hPipe = INVALID_HANDLE_VALUE;
+  for(;;) {
+    hPipe = CreateNamedPipe(
+      g_pipe_name,
+      PIPE_ACCESS_DUPLEX,
+      PIPE_TYPE_BYTE,
+      GetMaxConcurrency(),
+      1, // nOutBufferSize
+      1, // nInBufferSize
+      0, // nDefaultTimeOut
+      NULL // Default security attributes (noinherit)
+      );
+    if (hPipe != INVALID_HANDLE_VALUE)
+      break;
+
+    DWORD error = GetLastError();
+    if (error == ERROR_PIPE_BUSY) {
+      if (hEvent) {
+        WaitForSingleObject(hEvent, g_wait_time);
+      } else {
+        Sleep(g_wait_time);
+      }
+    } else {
+      if (g_is_debug) {
+        std::wostringstream strm;
+        strm << L"Got error " << error << L": "
+             << ErrorMessageToString(error) << L"\n";
+        wprintf(strm.str().c_str());
+      }
+      return RunExe();
+    }
+  }
+
+  if (g_is_debug) {
+    end_time = GetTickCount64();
+    wprintf(L"  took %.2fs to acquire semaphore.\n",
+              (end_time - start_time) / 1000.0);
+  }
+
+  DWORD ret = RunExe();
+
+  CloseHandle(hPipe);
+  if (hEvent != NULL) {
+    SetEvent(hEvent);
+  }
+
+  return ret;
+}