[Macros] Recovery after exectuable plugin crash

When executble plugins crashes 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

@@ -23,35 +23,69 @@
 namespace swift {
 
 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.
+  llvm::Error spawnIfNeeded();
+
   /// Send a message to the plugin.
   llvm::Error sendMessage(llvm::StringRef message) const;
 
@@ -63,7 +97,17 @@ class LoadedExecutablePlugin {
     this->cleanup = cleanup;
   }
 
-  llvm::sys::procid_t getPid() { return pid; }
+  /// Add "on reconnect" callback.
+  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 +122,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/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

@@ -94,8 +94,40 @@ PluginRegistry::loadExecutablePlugin(StringRef path) {
                                    "not executable");
   }
 
+  plugin = std::unique_ptr<LoadedExecutablePlugin>(
+      new LoadedExecutablePlugin(path, stat.getLastModificationTime()));
+
+  // Launch here to see if it's actually executable, and diagnose if necessary.
+  if (auto error = plugin->spawnIfNeeded()) {
+    plugin.reset();
+    return error;
+  }
+
+  return plugin.get();
+}
+
+llvm::Error LoadedExecutablePlugin::spawnIfNeeded() {
+  llvm::sys::fs::file_status stat;
+  if (auto err = llvm::sys::fs::status(ExecutablePath, stat)) {
+    return llvm::errorCodeToError(err);
+  }
+
+  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' call is too heavy.
+    // PluginRegistry::loadExecutablePlugin() checks it and replace this object
+    // itself if the plugin is updated. If the process is "stale" and we
+    // relaunch the
+
+    // 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 +139,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));
 
-  return plugin.get();
+  // Call "on reconnect" callbacks.
+  for (auto *cb : onReconnect) {
+    (*cb)();
+  }
+
+  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 +193,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;
 
@@ -179,6 +219,7 @@ ssize_t LoadedExecutablePlugin::write(const void *buf, size_t nbyte) const {
 }
 
 llvm::Error LoadedExecutablePlugin::sendMessage(llvm::StringRef message) const {
+  assert(isAlive());
   ssize_t writtenSize = 0;
 
   const char *data = message.data();
@@ -187,15 +228,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");
   }
@@ -204,13 +247,15 @@ llvm::Error LoadedExecutablePlugin::sendMessage(llvm::StringRef message) const {
 }
 
 llvm::Expected<std::string> LoadedExecutablePlugin::waitForNextMessage() const {
+  assert(isAlive());
   ssize_t readSize = 0;
 
   // 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 +269,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");
     }

lib/ASTGen/Sources/ASTGen/PluginHost.swift

@@ -38,30 +38,37 @@ public func _deinitializePlugin(
 }
 
 /// Load the library plugin in the plugin server.
+/// This should be called inside lock.
 @_cdecl("swift_ASTGen_pluginServerLoadLibraryPlugin")
 func swift_ASTGen_pluginServerLoadLibraryPlugin(
   opaqueHandle: UnsafeMutableRawPointer,
   libraryPath: UnsafePointer<Int8>,
   moduleName: UnsafePointer<Int8>,
-  cxxDiagnosticEngine: UnsafeMutablePointer<UInt8>
+  cxxDiagnosticEngine: UnsafeMutablePointer<UInt8>?
 ) -> Bool {
   let plugin =  CompilerPlugin(opaqueHandle: opaqueHandle)
   assert(plugin.capability?.features.contains(.loadPluginLibrary) == true)
   let libraryPath = String(cString: libraryPath)
   let moduleName = String(cString: moduleName)
-  let diagEngine = PluginDiagnosticsEngine(cxxDiagnosticEngine: cxxDiagnosticEngine)
+
+  let diagEngine: PluginDiagnosticsEngine?
+  if let cxxDiagnosticEngine = cxxDiagnosticEngine {
+    diagEngine = PluginDiagnosticsEngine(cxxDiagnosticEngine: cxxDiagnosticEngine)
+  } else {
+    diagEngine = nil
+  }
 
   do {
-    let result = try plugin.sendMessageAndWait(
+    let result = try plugin.sendMessageAndWaitWithoutLock(
       .loadPluginLibrary(libraryPath: libraryPath, moduleName: moduleName)
     )
     guard case .loadPluginLibraryResult(let loaded, let diagnostics) = result else {
       throw PluginError.invalidReponseKind
     }
-    diagEngine.emit(diagnostics);
+    diagEngine?.emit(diagnostics);
     return loaded
   } catch {
-    diagEngine.diagnose(error: error)
+    diagEngine?.diagnose(error: error)
     return false
   }
 }
@@ -121,28 +128,31 @@ struct CompilerPlugin {
     return try LLVMJSON.decode(PluginToHostMessage.self, from: data)
   }
 
+  func sendMessageAndWaitWithoutLock(_ message: HostToPluginMessage) throws -> PluginToHostMessage {
+    try sendMessage(message)
+    return try waitForNextMessage()
+  }
+
   func sendMessageAndWait(_ message: HostToPluginMessage) throws -> PluginToHostMessage {
     try self.withLock {
-      try sendMessage(message)
-      return try waitForNextMessage()
+      guard !Plugin_spawnIfNeeded(opaqueHandle) else {
+        throw PluginError.failedToSendMessage
+      }
+      return try sendMessageAndWaitWithoutLock(message);
     }
   }
 
+  /// Initialize the plugin. This should be called inside lock.
   func initialize() {
-    // Don't use `sendMessageAndWait` because we want to keep the lock until
-    // setting the returned value.
     do {
-      try self.withLock {
-        // Get capability.
-        try self.sendMessage(.getCapability)
-        let response = try self.waitForNextMessage()
-        guard case .getCapabilityResult(let capability) = response else {
-          throw PluginError.invalidReponseKind
-        }
-        let ptr = UnsafeMutablePointer<Capability>.allocate(capacity: 1)
-        ptr.initialize(to: .init(capability))
-        Plugin_setCapability(opaqueHandle, UnsafeRawPointer(ptr))
+      // Get capability.
+      let response = try self.sendMessageAndWaitWithoutLock(.getCapability)
+      guard case .getCapabilityResult(let capability) = response else {
+        throw PluginError.invalidReponseKind
       }
+      let ptr = UnsafeMutablePointer<Capability>.allocate(capacity: 1)
+      ptr.initialize(to: .init(capability))
+      Plugin_setCapability(opaqueHandle, UnsafeRawPointer(ptr))
     } catch {
       assertionFailure(String(describing: error))
       return