[Macros] Recovery after executable plugin crash

When executable plugins crashed or somehow decided to exit, the compiler
should relaunch the plugin executable before sending another message.

Author: rintaro

include/swift/AST/CASTBridging.h

@@ -307,6 +307,9 @@ void Plugin_lock(PluginHandle handle);
 /// Unlock the plugin.
 void Plugin_unlock(PluginHandle handle);
 
+/// Launch the plugin if it's not running.
+_Bool Plugin_spawnIfNeeded(PluginHandle handle);
+
 /// Sends the message to the plugin, returns true if there was an error.
 /// Clients should receive the response  by \c Plugin_waitForNextMessage .
 _Bool Plugin_sendMessage(PluginHandle handle, const BridgedData data);

include/swift/AST/PluginRegistry.h

@@ -22,36 +22,79 @@
 
 namespace swift {
 
+/// Represent a "resolved" exectuable plugin.
+///
+/// Plugin clients usually deal with this object to communicate with the actual
+/// plugin implementation.
+/// This object has a file path of the plugin executable, and is responsible to
+/// launch it and manages the process. When the plugin process crashes, this
+/// should automatically relaunch the process so the clients can keep using this
+/// object as the interface.
 class LoadedExecutablePlugin {
-  const llvm::sys::procid_t pid;
+
+  /// Represents the current process of the executable plugin.
+  struct PluginProcess {
+    const llvm::sys::procid_t pid;
+    const int inputFileDescriptor;
+    const int outputFileDescriptor;
+    bool isStale = false;
+
+    PluginProcess(llvm::sys::procid_t pid, int inputFileDescriptor,
+                  int outputFileDescriptor);
+
+    ~PluginProcess();
+
+    ssize_t write(const void *buf, size_t nbyte) const;
+    ssize_t read(void *buf, size_t nbyte) const;
+  };
+
+  /// Launched current process.
+  std::unique_ptr<PluginProcess> Process;
+
+  /// Path to the plugin executable.
+  const std::string ExecutablePath;
+
+  /// Last modification time of the `ExecutablePath` when this is initialized.
   const llvm::sys::TimePoint<> LastModificationTime;
-  const int inputFileDescriptor;
-  const int outputFileDescriptor;
 
   /// Opaque value of the protocol capability of the pluugin. This is a
   /// value from ASTGen.
   const void *capability = nullptr;
 
+  /// Callbacks to be called when the connection is restored.
+  llvm::SmallVector<std::function<void(void)> *, 0> onReconnect;
+
   /// Cleanup function to call ASTGen.
   std::function<void(void)> cleanup;
 
   std::mutex mtx;
 
-  ssize_t write(const void *buf, size_t nbyte) const;
-  ssize_t read(void *buf, size_t nbyte) const;
-
 public:
-  LoadedExecutablePlugin(llvm::sys::procid_t pid,
-                         llvm::sys::TimePoint<> LastModificationTime,
-                         int inputFileDescriptor, int outputFileDescriptor);
+  LoadedExecutablePlugin(llvm::StringRef ExecutablePath,
+                         llvm::sys::TimePoint<> LastModificationTime)
+      : ExecutablePath(ExecutablePath),
+        LastModificationTime(LastModificationTime){};
   ~LoadedExecutablePlugin();
+
+  /// The last modification time of 'ExecutablePath' when this object is
+  /// created.
   llvm::sys::TimePoint<> getLastModificationTime() const {
     return LastModificationTime;
   }
 
+  /// Indicates that the current process is usable.
+  bool isAlive() const { return Process != nullptr && !Process->isStale; }
+
+  /// Mark the current process "stale".
+  void setStale() const { Process->isStale = true; }
+
   void lock() { mtx.lock(); }
   void unlock() { mtx.unlock(); }
 
+  // Launch the plugin if it's not already running, or it's stale. Return an
+  // error if it's fails to execute it.
+  llvm::Error spawnIfNeeded();
+
   /// Send a message to the plugin.
   llvm::Error sendMessage(llvm::StringRef message) const;
 
@@ -63,7 +106,18 @@ class LoadedExecutablePlugin {
     this->cleanup = cleanup;
   }
 
-  llvm::sys::procid_t getPid() { return pid; }
+  /// Add "on reconnect" callback.
+  /// These callbacks are called when `spawnIfNeeded()` relaunched the plugin.
+  void addOnReconnect(std::function<void(void)> *fn) {
+    onReconnect.push_back(fn);
+  }
+
+  /// Remove "on reconnect" callback.
+  void removeOnReconnect(std::function<void(void)> *fn) {
+    llvm::erase_value(onReconnect, fn);
+  }
+
+  llvm::sys::procid_t getPid() { return Process->pid; }
 
   const void *getCapability() { return capability; };
   void setCapability(const void *newValue) { capability = newValue; };
@@ -78,7 +132,12 @@ class PluginRegistry {
       LoadedPluginExecutables;
 
 public:
+  /// Load a dynamic link library specified by \p path.
+  /// If \p path plugin is already loaded, this returns the cached object.
   llvm::Expected<void *> loadLibraryPlugin(llvm::StringRef path);
+
+  /// Load an executable plugin specified by \p path .
+  /// If \p path plugin is already loaded, this returns the cached object.
   llvm::Expected<LoadedExecutablePlugin *>
   loadExecutablePlugin(llvm::StringRef path);
 

lib/AST/ASTContext.cpp

@@ -528,6 +528,9 @@ struct ASTContext::Implementation {
   /// NOTE: Do not reference this directly. Use ASTContext::getPluginRegistry().
   PluginRegistry *Plugins = nullptr;
 
+  /// `Plugins` storage if this ASTContext owns it.
+  std::unique_ptr<PluginRegistry> OwnedPluginRegistry = nullptr;
+
   /// Cache of loaded symbols.
   llvm::StringMap<void *> LoadedSymbols;
 
@@ -6255,9 +6258,7 @@ PluginRegistry *ASTContext::getPluginRegistry() const {
   // Create a new one if it hasn't been set.
   if (!registry) {
     registry = new PluginRegistry();
-    const_cast<ASTContext *>(this)->addCleanup([registry]{
-      delete registry;
-    });
+    getImpl().OwnedPluginRegistry.reset(registry);
   }
 
   assert(registry != nullptr);

lib/AST/CASTBridging.cpp

@@ -649,6 +649,14 @@ void Plugin_unlock(PluginHandle handle) {
   plugin->unlock();
 }
 
+bool Plugin_spawnIfNeeded(PluginHandle handle) {
+  auto *plugin = static_cast<LoadedExecutablePlugin *>(handle);
+  auto error = plugin->spawnIfNeeded();
+  bool hadError(error);
+  llvm::consumeError(std::move(error));
+  return hadError;
+}
+
 bool Plugin_sendMessage(PluginHandle handle, const BridgedData data) {
   auto *plugin = static_cast<LoadedExecutablePlugin *>(handle);
   StringRef message(data.baseAddress, data.size);

lib/AST/PluginRegistry.cpp

@@ -38,9 +38,6 @@
 #include <io.h>
 #endif
 
-extern "C" const void *swift_ASTGen_getCompilerPluginCapability(void *handle);
-extern "C" void swift_ASTGen_destroyCompilerPluginCapability(void *value);
-
 using namespace swift;
 
 llvm::Expected<void *> PluginRegistry::loadLibraryPlugin(StringRef path) {
@@ -74,14 +71,14 @@ PluginRegistry::loadExecutablePlugin(StringRef path) {
   }
 
   // See if the plugin is already loaded.
-  auto &plugin = LoadedPluginExecutables[path];
-  if (plugin) {
+  auto &storage = LoadedPluginExecutables[path];
+  if (storage) {
     // See if the loaded one is still usable.
-    if (plugin->getLastModificationTime() == stat.getLastModificationTime())
-      return plugin.get();
+    if (storage->getLastModificationTime() == stat.getLastModificationTime())
+      return storage.get();
 
     // The plugin is updated. Close the previously opened plugin.
-    plugin = nullptr;
+    storage = nullptr;
   }
 
   if (!llvm::sys::fs::exists(stat)) {
@@ -94,8 +91,35 @@ PluginRegistry::loadExecutablePlugin(StringRef path) {
                                    "not executable");
   }
 
+  auto plugin = std::unique_ptr<LoadedExecutablePlugin>(
+      new LoadedExecutablePlugin(path, stat.getLastModificationTime()));
+
+  // Launch here to see if it's actually executable, and diagnose (by returning
+  // an error) if necessary.
+  if (auto error = plugin->spawnIfNeeded()) {
+    return std::move(error);
+  }
+
+  storage = std::move(plugin);
+  return storage.get();
+}
+
+llvm::Error LoadedExecutablePlugin::spawnIfNeeded() {
+  if (Process) {
+    // See if the loaded one is still usable.
+    if (!Process->isStale)
+      return llvm::Error::success();
+
+    // NOTE: We don't check the mtime here because 'stat(2)' call is too heavy.
+    // PluginRegistry::loadExecutablePlugin() checks it and replace this object
+    // itself if the plugin is updated.
+
+    // The plugin is stale. Discard the previously opened process.
+    Process.reset();
+  }
+
   // Create command line arguments.
-  SmallVector<StringRef, 4> command{path};
+  SmallVector<StringRef, 4> command{ExecutablePath};
 
   // Apply sandboxing.
   llvm::BumpPtrAllocator Allocator;
@@ -107,29 +131,36 @@ PluginRegistry::loadExecutablePlugin(StringRef path) {
     return llvm::errorCodeToError(childInfo.getError());
   }
 
-  plugin = std::unique_ptr<LoadedExecutablePlugin>(new LoadedExecutablePlugin(
-      childInfo->Pid, stat.getLastModificationTime(),
-      childInfo->ReadFileDescriptor, childInfo->WriteFileDescriptor));
+  Process = std::unique_ptr<PluginProcess>(
+      new PluginProcess(childInfo->Pid, childInfo->ReadFileDescriptor,
+                        childInfo->WriteFileDescriptor));
+
+  // Call "on reconnect" callbacks.
+  for (auto *callback : onReconnect) {
+    (*callback)();
+  }
 
-  return plugin.get();
+  return llvm::Error::success();
 }
 
-LoadedExecutablePlugin::LoadedExecutablePlugin(
-    llvm::sys::procid_t pid, llvm::sys::TimePoint<> LastModificationTime,
-    int inputFileDescriptor, int outputFileDescriptor)
-    : pid(pid), LastModificationTime(LastModificationTime),
-      inputFileDescriptor(inputFileDescriptor),
+LoadedExecutablePlugin::PluginProcess::PluginProcess(llvm::sys::procid_t pid,
+                                                     int inputFileDescriptor,
+                                                     int outputFileDescriptor)
+    : pid(pid), inputFileDescriptor(inputFileDescriptor),
       outputFileDescriptor(outputFileDescriptor) {}
 
-LoadedExecutablePlugin::~LoadedExecutablePlugin() {
+LoadedExecutablePlugin::PluginProcess::~PluginProcess() {
   close(inputFileDescriptor);
   close(outputFileDescriptor);
+}
 
+LoadedExecutablePlugin::~LoadedExecutablePlugin() {
   // Let ASTGen to cleanup things.
   this->cleanup();
 }
 
-ssize_t LoadedExecutablePlugin::read(void *buf, size_t nbyte) const {
+ssize_t LoadedExecutablePlugin::PluginProcess::read(void *buf,
+                                                    size_t nbyte) const {
   ssize_t bytesToRead = nbyte;
   void *ptr = buf;
 
@@ -154,7 +185,8 @@ ssize_t LoadedExecutablePlugin::read(void *buf, size_t nbyte) const {
   return nbyte - bytesToRead;
 }
 
-ssize_t LoadedExecutablePlugin::write(const void *buf, size_t nbyte) const {
+ssize_t LoadedExecutablePlugin::PluginProcess::write(const void *buf,
+                                                     size_t nbyte) const {
   ssize_t bytesToWrite = nbyte;
   const void *ptr = buf;
 
@@ -187,15 +219,17 @@ llvm::Error LoadedExecutablePlugin::sendMessage(llvm::StringRef message) const {
   // Write header (message size).
   uint64_t header = llvm::support::endian::byte_swap(
       uint64_t(size), llvm::support::endianness::little);
-  writtenSize = write(&header, sizeof(header));
+  writtenSize = Process->write(&header, sizeof(header));
   if (writtenSize != sizeof(header)) {
+    setStale();
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "failed to write plugin message header");
   }
 
   // Write message.
-  writtenSize = write(data, size);
+  writtenSize = Process->write(data, size);
   if (writtenSize != ssize_t(size)) {
+    setStale();
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "failed to write plugin message data");
   }
@@ -208,9 +242,10 @@ llvm::Expected<std::string> LoadedExecutablePlugin::waitForNextMessage() const {
 
   // Read header (message size).
   uint64_t header;
-  readSize = read(&header, sizeof(header));
+  readSize = Process->read(&header, sizeof(header));
 
   if (readSize != sizeof(header)) {
+    setStale();
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "failed to read plugin message header");
   }
@@ -224,8 +259,9 @@ llvm::Expected<std::string> LoadedExecutablePlugin::waitForNextMessage() const {
   auto sizeToRead = size;
   while (sizeToRead > 0) {
     char buffer[4096];
-    readSize = read(buffer, std::min(sizeof(buffer), sizeToRead));
+    readSize = Process->read(buffer, std::min(sizeof(buffer), sizeToRead));
     if (readSize == 0) {
+      setStale();
       return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                      "failed to read plugin message data");
     }