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>
+
+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
+
+static std::string ErrorMessageToString(DWORD err) {
+  char* msg_buf = NULL;
+  DWORD rc = FormatMessageA(FORMAT_MESSAGE_ALLOCATE_BUFFER |

[nsylvain, 2012/08/01 18:53:46]

FormatMessageA is sad a little bit. We can't use the default?

[iannucci, 2012/08/01 19:09:53]

Because I use it for formatting exception messages (and because the std::
exception classes only take narrow strings (and because I didn't want to make a
custom exception class)), I opted to use the ansi variant here.

Now that I think about it, if someone runs this on a non-ansi-friendly box, this
would potentially do the wrong thing. I'll suck it up and make a custom wide
exception so I can do unicode everywhere :)

+                               FORMAT_MESSAGE_FROM_SYSTEM,
+                           NULL,

[nsylvain, 2012/08/01 18:53:46]

spacing is off by one

[iannucci, 2012/08/01 19:09:53]

So it is... ugly :(

+                           err,
+                           MAKELANGID(LANG_NEUTRAL, SUBLANG_DEFAULT),
+                           reinterpret_cast<LPSTR>(&msg_buf),
+                           0,
+                           NULL);
+  if (!rc)
+    return "unknown error";
+  std::string 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+1);
+  }
+
+  if (!CreateProcess(NULL,
+                     const_cast<TCHAR*>(cmdline.c_str()),
+                     NULL,
+                     NULL,
+                     TRUE,
+                     0,
+                     NULL,
+                     NULL,
+                     &startup_info, &process_info)) {
+    std::string error = ErrorMessageToString(GetLastError());

[nsylvain, 2012/08/01 18:53:46]

Little overkill maybe. I'm sure people can just look up the error number.

[iannucci, 2012/08/01 19:09:53]

Yeah probably. Again, this was a holdover from supalink, and it's one less
"google msdn error ######" search I have to do :P

+    std::string narrow_cmdline(cmdline.begin(), cmdline.end());
+    std::ostringstream strm;
+    strm << "Error in CreateProcess["<< narrow_cmdline <<"]: " << error;
+    throw std::runtime_error(strm.str().c_str());
+  }
+  CloseHandle(process_info.hThread);
+  WaitForSingleObject(process_info.hProcess, INFINITE);
+  GetExitCodeProcess(process_info.hProcess, &exit_code);
+  CloseHandle(process_info.hProcess);
+  return exit_code;
+}
+
+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) {
+    std::string narrow_envvar_name(envvar_name.begin(), envvar_name.end());
+    std::ostringstream strm;
+    strm << "Error while getting number of cores. Try setting the "
+         << " environment variable '" << narrow_envvar_name << '\''
+         << " to (num_cores-1): " << ErrorMessageToString(last_error) << "\n";

[nsylvain, 2012/08/01 18:53:46]

same here I guess.

+    throw std::runtime_error(strm.str());
+  }
+
+  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
+};
+
+static int GetMaxConcurrency(
+  const tstring& base_pipename,
+  ConcurrencyMetricEnum metric = CONCURRENCY_METRIC_DEFAULT)
+{
+  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 = _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(envvar_name);
+        break;
+      }
+    }
+
+    max_concurrent = std::min(std::max(max_concurrent, 1),
+                              PIPE_UNLIMITED_INSTANCES);
+  }
+
+  return max_concurrent;
+}
+
+
+static HANDLE wait_for_pipe(const tstring& pipename,
+                            HANDLE hEvent,
+                            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 hPipe = INVALID_HANDLE_VALUE;
+  for(;;) {
+    hPipe = CreateNamedPipe(
+      pipename.c_str(),
+      PIPE_ACCESS_DUPLEX,
+      PIPE_TYPE_BYTE,
+      max_concurrency,
+      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 {
+      std::ostringstream strm;
+      strm << L"Got error " << error << L": "
+           << ErrorMessageToString(error) << L"\n";
+      throw std::logic_error(strm.str());
+    }
+  }
+
+  return hPipe;
+}
+
+static int wait_and_run(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 hEvent = CreateEvent(
+    NULL, // Default security attributes
+    FALSE, // Manual reset
+    FALSE, // Initial state
+    event_name.c_str());
+
+  HANDLE hPipe = INVALID_HANDLE_VALUE;
+  DWORD ret;
+  try {
+    if (g_is_debug)
+      start_time = GetTickCount64();
+    hPipe = wait_for_pipe(pipename, hEvent, GetMaxConcurrency(base_pipename));
+    if (g_is_debug) {
+      end_time = GetTickCount64();
+      wprintf(L"  took %.2fs to acquire semaphore.\n",
+        (end_time - start_time) / 1000.0);
+    }
+    ret = RunExe(shimmed_exe);
+  } catch (...) {
+    if (hPipe != INVALID_HANDLE_VALUE) { CloseHandle(hPipe); }
+    if (hEvent != NULL) { SetEvent(hEvent); }
+    throw;
+  }
+
+  if (hPipe != INVALID_HANDLE_VALUE) { CloseHandle(hPipe); }
+  if (hEvent != NULL) { SetEvent(hEvent); }
+
+  return ret;
+}
+
+void usage(const std::string& msg) {
+  std::ostringstream strm;
+  strm << msg << "\n";
+  strm << "Usage: SHIMED_NAME__SEMAPHORE_NAME\n"
+       << "\n"
+       << "  SHIMMED_NAME   - ex. 'link.exe' or 'lib.exe'\n"
+       << "                 - can be exe, bat, or com\n"
+       << "                 - must exist in PATH\n"
+       << "\n"
+       << "  SEMAPHORE_NAME - ex. 'SOME_NAME' or 'GROOVY_SEMAPHORE'\n"
+       << "\n"
+       << "  Example:\n"
+       << "    link.exe__LINK_LIMITER.exe\n"
+       << "    lib.exe__LINK_LIMITER.exe\n"
+       << "      * Both will limit on the same semaphore\n"
+       << "\n"
+       << "    link.exe__LINK_LIMITER.exe\n"
+       << "    lib.exe__LIB_LIMITER.exe\n"
+       << "      * Both will limit on independent semaphores\n"
+       << "\n"
+       << "  This program is meant to be run after renaming it into the\n"
+       << "  above format. Once you have done so, executing it will block\n"
+       << "  on the availability of the semaphore SEMAPHORE_NAME. Once\n"
+       << "  the semaphore is obtained, it will execute SHIMMED_NAME, \n"
+       << "  passing through all arguments as-is.\n"
+       << "\n"
+       << "  The maximum concurrency can be manually set by setting the\n"
+       << "  environment variable <SEMAPHORE_NAME>_MAXCONCURRENCY to an\n"
+       << "  integer value (1, 254).\n"
+       << "    * This value must be set the same for ALL invocations.\n"
+       << "    * If the value is not set, it defaults to (num_cores-1).\n"
+       << "\n"
+       << "  The semaphore is automatically released when the program\n"
+       << "  completes normally, OR if the program crashes (or even if\n"
+       << "  limiter itself crashes).\n";
+  throw std::logic_error(strm.str().c_str());
+}
+
+static int inner_main(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("Cannot parse argv[0]. No '__' found. "
+          "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("Hunh? No '.' in argv[0]? Are we on windows?");
+  }
+  base_pipename = base_pipename.substr(0, dot);
+
+  return wait_and_run(shimmed_exe, base_pipename);
+}
+
+// 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 argc, wchar_t** argv) {
+  try {
+    return inner_main(argc, argv);
+  } catch(const std::exception &ex) {
+    // Don't use stderr for errors because VS has large buffers on them, leading
+    // to confusing error output.
+    printf("limiter fatal error: %s\n", ex.what());
+    return 1;
+  }
+}