[Runtime] Allow disabling/enabling prespecializations library per-process.

include/swift/Runtime/LibPrespecialized.h

@@ -31,9 +31,12 @@ struct LibPrespecializedData {
   uint32_t minorVersion;
 
   TargetPointer<Runtime, const void> metadataMap;
+  TargetPointer<Runtime, const void> disabledProcessesTable;
 
   static constexpr uint32_t currentMajorVersion = 1;
-  static constexpr uint32_t currentMinorVersion = 1;
+  static constexpr uint32_t currentMinorVersion = 2;
+
+  static constexpr uint32_t minorVersionWithDisabledProcessesTable = 2;
 
   // Helpers for retrieving the metadata map in-process.
   static bool stringIsNull(const char *str) { return str == nullptr; }
@@ -43,6 +46,12 @@ struct LibPrespecializedData {
   const MetadataMap *getMetadataMap() const {
     return reinterpret_cast<const MetadataMap *>(metadataMap);
   }
+
+  const char *const *getDisabledProcessesTable() const {
+    if (minorVersion < minorVersionWithDisabledProcessesTable)
+      return nullptr;
+    return reinterpret_cast<const char *const *>(disabledProcessesTable);
+  }
 };
 
 const LibPrespecializedData<InProcess> *getLibPrespecializedData();

stdlib/public/runtime/EnvironmentVariables.def

@@ -85,9 +85,19 @@ VARIABLE(SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED, bool, true,
          "Enable use of prespecializations library.")
 
 VARIABLE(SWIFT_DEBUG_LIB_PRESPECIALIZED_PATH, string, "",
-         "A path to a prespecializations library to use at runtime. In order to"
-         "be used, this library must be loaded into the process by other means"
-         "(such as DYLD_INSERT_LIBRARIES) before Swift tries to use it.")
+         "A path to a prespecializations library to use at runtime. In order "
+         "to be used, this library must be loaded into the process by other "
+         "means (such as DYLD_INSERT_LIBRARIES) before Swift tries to use it.")
+
+VARIABLE(SWIFT_DEBUG_LIB_PRESPECIALIZED_DISABLED_PROCESSES, string, "",
+         "A colon-separated list of process names where the prespecializations "
+         "library will be forcibly disabled.")
+
+VARIABLE(SWIFT_DEBUG_LIB_PRESPECIALIZED_ENABLED_PROCESSES, string, "",
+         "A colon-separated list of process names where the prespecializations "
+         "library will be forcibly enabled. This overrides the disabled "
+         "processes list in the prespecializations library, as well as the "
+         "list in SWIFT_DEBUG_LIB_PRESPECIALIZED_DISABLED_PROCESSES.")
 
 VARIABLE(SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING, bool, false,
          "Enable debug logging of prespecializations library use.")

stdlib/public/runtime/LibPrespecialized.cpp

@@ -32,9 +32,85 @@ using namespace swift;
 
 static std::atomic<bool> disablePrespecializedMetadata = false;
 
+static bool prespecializedLoggingEnabled = false;
+
+#define LOG(fmt, ...)                                                          \
+  do {                                                                         \
+    if (SWIFT_UNLIKELY(prespecializedLoggingEnabled))                          \
+      fprintf(stderr, "Prespecializations library: " fmt "\n", __VA_ARGS__);   \
+  } while (0)
+
+static bool environmentProcessListContainsProcess(const char *list,
+                                                  const char *progname) {
+  auto prognameLen = strlen(progname);
+
+  const char *cursor = list;
+  while (true) {
+    const char *next = strchr(cursor, ':');
+    if (!next) {
+      // Last entry in the list. Compare with the entire rest of the string.
+      return strcmp(progname, cursor) == 0;
+    }
+
+    // Entry at beginning or middle of the list. Compare against this substring.
+    size_t len = next - cursor;
+    if (len == prognameLen && strncmp(cursor, progname, len) == 0)
+      return true;
+
+    cursor = next + 1;
+  }
+}
+
+static bool isThisProcessEnabled(const LibPrespecializedData<InProcess> *data) {
+  extern const char *__progname;
+
+  if (!__progname)
+    return true;
+
+  auto envEnabledProcesses =
+      runtime::environment::SWIFT_DEBUG_LIB_PRESPECIALIZED_ENABLED_PROCESSES();
+  if (envEnabledProcesses && *envEnabledProcesses) {
+    if (environmentProcessListContainsProcess(envEnabledProcesses,
+                                              __progname)) {
+      LOG("Found %s in SWIFT_DEBUG_LIB_PRESPECIALIZED_ENABLED_PROCESSES, "
+          "enabling",
+          __progname);
+      return true;
+    }
+  }
+
+  auto envDisabledProcesses =
+      runtime::environment::SWIFT_DEBUG_LIB_PRESPECIALIZED_DISABLED_PROCESSES();
+  if (envDisabledProcesses && *envDisabledProcesses) {
+    if (environmentProcessListContainsProcess(envDisabledProcesses,
+                                              __progname)) {
+      LOG("Found %s in SWIFT_DEBUG_LIB_PRESPECIALIZED_DISABLED_PROCESSES, "
+          "disabling",
+          __progname);
+      return false;
+    }
+  }
+
+  if (auto *disabledProcesses = data->getDisabledProcessesTable()) {
+    auto *cursor = disabledProcesses;
+    while (auto *name = *cursor) {
+      if (strcmp(name, __progname) == 0) {
+        LOG("Found %s in disabled processes list, disabling", name);
+        return false;
+      }
+      cursor++;
+    }
+  }
+
+  return true;
+}
+
 static const LibPrespecializedData<InProcess> *findLibPrespecialized() {
-  if (!runtime::environment::SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED())
+  if (!runtime::environment::SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED()) {
+    LOG("Disabling, SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED = %d",
+        runtime::environment::SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED());
     return nullptr;
+  }
 
   const void *dataPtr = nullptr;
 #if USE_DLOPEN
@@ -51,15 +127,22 @@ static const LibPrespecializedData<InProcess> *findLibPrespecialized() {
     }
 
     dataPtr = dlsym(handle, LIB_PRESPECIALIZED_TOP_LEVEL_SYMBOL_NAME);
+    LOG("Loaded custom library from %s, found dataPtr %p", path, dataPtr);
   }
 #if DYLD_GET_SWIFT_PRESPECIALIZED_DATA_DEFINED
   else if (SWIFT_RUNTIME_WEAK_CHECK(_dyld_get_swift_prespecialized_data)) {
     // Disable the prespecializations library if anything in the shared cache is
     // overridden. Eventually we want to be cleverer and only disable the
     // prespecializations that have been invalidated, but we'll start with the
     // simplest approach.
-    if (!dyld_shared_cache_some_image_overridden())
+    if (!dyld_shared_cache_some_image_overridden()) {
       dataPtr = SWIFT_RUNTIME_WEAK_USE(_dyld_get_swift_prespecialized_data());
+      LOG("Got dataPtr %p from _dyld_get_swift_prespecialized_data", dataPtr);
+    } else {
+      LOG("Not calling _dyld_get_swift_prespecialized_data "
+          "dyld_shared_cache_some_image_overridden = %d",
+          dyld_shared_cache_some_image_overridden());
+    }
   }
 #endif
 #endif
@@ -70,8 +153,16 @@ static const LibPrespecializedData<InProcess> *findLibPrespecialized() {
   auto *data =
       reinterpret_cast<const LibPrespecializedData<InProcess> *>(dataPtr);
   if (data->majorVersion !=
-      LibPrespecializedData<InProcess>::currentMajorVersion)
+      LibPrespecializedData<InProcess>::currentMajorVersion) {
+    LOG("Unknown major version %" PRIu32 ", disabling", data->majorVersion);
     return nullptr;
+  }
+
+  if (!isThisProcessEnabled(data))
+    return nullptr;
+
+  LOG("Returning data %p, major version %" PRIu32 " minor %" PRIu32, data,
+      data->majorVersion, data->minorVersion);
 
   return data;
 }
@@ -85,13 +176,12 @@ struct LibPrespecializedState {
     }
   };
 
-  bool loggingEnabled;
   const LibPrespecializedData<InProcess> *data;
   AddressRange sharedCacheRange{0, 0};
   AddressRange metadataAllocatorInitialPoolRange{0, 0};
 
   LibPrespecializedState() {
-    loggingEnabled =
+    prespecializedLoggingEnabled =
         runtime::environment::SWIFT_DEBUG_ENABLE_LIB_PRESPECIALIZED_LOGGING();
     data = findLibPrespecialized();
 
@@ -112,12 +202,6 @@ struct LibPrespecializedState {
 
 static Lazy<LibPrespecializedState> LibPrespecialized;
 
-#define LOG(fmt, ...)                                                          \
-  do {                                                                         \
-    if (SWIFT_UNLIKELY(prespecialized.loggingEnabled))                         \
-      fprintf(stderr, "Prespecializations library: " fmt "\n", __VA_ARGS__);   \
-  } while (0)
-
 const LibPrespecializedData<InProcess> *swift::getLibPrespecializedData() {
   return SWIFT_LAZY_CONSTANT(findLibPrespecialized());
 }