Merge pull request #68280 from compnerd/execution

AST,Basic: update support for process execution on Windows

Author: compnerd

include/swift/AST/PluginRegistry.h

@@ -52,14 +52,13 @@ class LoadedExecutablePlugin {
 
   /// Represents the current process of the executable plugin.
   struct PluginProcess {
-    const llvm::sys::procid_t pid;
-    const int inputFileDescriptor;
-    const int outputFileDescriptor;
+    const llvm::sys::ProcessInfo process;
+    const int input;
+    const int output;
     bool isStale = false;
 
-    PluginProcess(llvm::sys::procid_t pid, int inputFileDescriptor,
-                  int outputFileDescriptor);
-
+    PluginProcess(llvm::sys::ProcessInfo process, int input, int output)
+        : process(process), input(input), output(output) {}
     ~PluginProcess();
 
     ssize_t write(const void *buf, size_t nbyte) const;
@@ -138,7 +137,8 @@ class LoadedExecutablePlugin {
     llvm::erase_value(onReconnect, fn);
   }
 
-  llvm::sys::procid_t getPid() { return Process->pid; }
+  llvm::sys::procid_t getPid() { return Process->process.Pid; }
+  llvm::sys::process_t getProcess() { return Process->process.Process; }
 
   NullTerminatedStringRef getExecutablePath() {
     return {ExecutablePath.c_str(), ExecutablePath.size()};

include/swift/Basic/Program.h

@@ -40,14 +40,12 @@ int ExecuteInPlace(const char *Program, const char **args,
                    const char **env = nullptr);
 
 struct ChildProcessInfo {
-  llvm::sys::procid_t Pid;
-  int WriteFileDescriptor;
-  int ReadFileDescriptor;
+  llvm::sys::ProcessInfo ProcessInfo;
+  int Write;
+  int Read;
 
-  ChildProcessInfo(llvm::sys::procid_t Pid, int WriteFileDescriptor,
-                   int ReadFileDescriptor)
-      : Pid(Pid), WriteFileDescriptor(WriteFileDescriptor),
-        ReadFileDescriptor(ReadFileDescriptor) {}
+  ChildProcessInfo(llvm::sys::ProcessInfo ProcessInfo, int Write, int Read)
+      : ProcessInfo(ProcessInfo), Write(Write), Read(Read) {}
 };
 
 /// This function executes the program using the argument provided.

lib/AST/PluginLoader.cpp

@@ -41,8 +41,13 @@ PluginRegistry *PluginLoader::getRegistry() {
 static StringRef pluginModuleNameStringFromPath(StringRef path) {
   // Plugin library must be named 'lib${module name}(.dylib|.so|.dll)'.
   // FIXME: Shared library prefix might be different between platforms.
+#if defined(_WIN32)
+  constexpr StringRef libPrefix{};
+  constexpr StringRef libSuffix = ".dll";
+#else
   constexpr StringRef libPrefix = "lib";
   constexpr StringRef libSuffix = LTDL_SHLIB_EXT;
+#endif
 
   StringRef filename = llvm::sys::path::filename(path);
   if (filename.starts_with(libPrefix) && filename.ends_with(libSuffix)) {

lib/AST/PluginRegistry.cpp

@@ -75,7 +75,9 @@ void *LoadedLibraryPlugin::getAddressOfSymbol(const char *symbolName) {
   auto &cached = resolvedSymbols[symbolName];
   if (cached)
     return cached;
-#if !defined(_WIN32)
+#if defined(_WIN32)
+  cached = GetProcAddress(static_cast<HMODULE>(handle), symbolName);
+#else
   cached = dlsym(handle, symbolName);
 #endif
   return cached;
@@ -153,9 +155,8 @@ llvm::Error LoadedExecutablePlugin::spawnIfNeeded() {
     return llvm::errorCodeToError(childInfo.getError());
   }
 
-  Process = std::make_unique<PluginProcess>(childInfo->Pid,
-                                            childInfo->ReadFileDescriptor,
-                                            childInfo->WriteFileDescriptor);
+  Process = std::make_unique<PluginProcess>(childInfo->ProcessInfo,
+                                            childInfo->Read, childInfo->Write);
 
   // Call "on reconnect" callbacks.
   for (auto *callback : onReconnect) {
@@ -165,15 +166,15 @@ llvm::Error LoadedExecutablePlugin::spawnIfNeeded() {
   return llvm::Error::success();
 }
 
-LoadedExecutablePlugin::PluginProcess::PluginProcess(llvm::sys::procid_t pid,
-                                                     int inputFileDescriptor,
-                                                     int outputFileDescriptor)
-    : pid(pid), inputFileDescriptor(inputFileDescriptor),
-      outputFileDescriptor(outputFileDescriptor) {}
-
 LoadedExecutablePlugin::PluginProcess::~PluginProcess() {
-  close(inputFileDescriptor);
-  close(outputFileDescriptor);
+#if defined(_WIN32)
+  _close(input);
+  _close(output);
+  CloseHandle(process.Process);
+#else
+  close(input);
+  close(output);
+#endif
 }
 
 LoadedExecutablePlugin::~LoadedExecutablePlugin() {
@@ -184,6 +185,17 @@ LoadedExecutablePlugin::~LoadedExecutablePlugin() {
 
 ssize_t LoadedExecutablePlugin::PluginProcess::read(void *buf,
                                                     size_t nbyte) const {
+#if defined(_WIN32)
+  size_t nread = 0;
+  while (nread < nbyte) {
+    int n = _read(input, static_cast<char*>(buf) + nread,
+                  std::min(static_cast<size_t>(UINT32_MAX), nbyte - nread));
+    if (n <= 0)
+      break;
+    nread += n;
+  }
+  return nread;
+#else
   ssize_t bytesToRead = nbyte;
   void *ptr = buf;
 
@@ -195,7 +207,7 @@ ssize_t LoadedExecutablePlugin::PluginProcess::read(void *buf,
 
   while (bytesToRead > 0) {
     ssize_t readingSize = std::min(ssize_t(INT32_MAX), bytesToRead);
-    ssize_t readSize = ::read(inputFileDescriptor, ptr, readingSize);
+    ssize_t readSize = ::read(input, ptr, readingSize);
     if (readSize <= 0) {
       // 0: EOF (the plugin exited?), -1: error (e.g. broken pipe.)
       // FIXME: Mark the plugin 'stale' and relaunch later.
@@ -206,10 +218,22 @@ ssize_t LoadedExecutablePlugin::PluginProcess::read(void *buf,
   }
 
   return nbyte - bytesToRead;
+#endif
 }
 
 ssize_t LoadedExecutablePlugin::PluginProcess::write(const void *buf,
                                                      size_t nbyte) const {
+#if defined(_WIN32)
+  size_t nwritten = 0;
+  while (nwritten < nbyte) {
+    int n = _write(output, static_cast<const char *>(buf) + nwritten,
+                   std::min(static_cast<size_t>(UINT32_MAX), nbyte - nwritten));
+    if (n <= 0)
+      break;
+    nwritten += n;
+  }
+  return nwritten;
+#else
   ssize_t bytesToWrite = nbyte;
   const void *ptr = buf;
 
@@ -221,7 +245,7 @@ ssize_t LoadedExecutablePlugin::PluginProcess::write(const void *buf,
 
   while (bytesToWrite > 0) {
     ssize_t writingSize = std::min(ssize_t(INT32_MAX), bytesToWrite);
-    ssize_t writtenSize = ::write(outputFileDescriptor, ptr, writingSize);
+    ssize_t writtenSize = ::write(output, ptr, writingSize);
     if (writtenSize <= 0) {
       // -1: error (e.g. broken pipe,)
       // FIXME: Mark the plugin 'stale' and relaunch later.
@@ -231,13 +255,14 @@ ssize_t LoadedExecutablePlugin::PluginProcess::write(const void *buf,
     bytesToWrite -= writtenSize;
   }
   return nbyte - bytesToWrite;
+#endif
 }
 
 llvm::Error LoadedExecutablePlugin::sendMessage(llvm::StringRef message) const {
   ssize_t writtenSize = 0;
 
   if (dumpMessaging) {
-    llvm::dbgs() << "->(plugin:" << Process->pid << ") " << message << "\n";
+    llvm::dbgs() << "->(plugin:" << Process->process.Pid << ") " << message << '\n';
   }
 
   const char *data = message.data();
@@ -297,7 +322,7 @@ llvm::Expected<std::string> LoadedExecutablePlugin::waitForNextMessage() const {
   }
 
   if (dumpMessaging) {
-    llvm::dbgs() << "<-(plugin:" << Process->pid << ") " << message << "\n";
+    llvm::dbgs() << "<-(plugin:" << Process->process.Pid << ") " << message << "\n";
   }
 
   return message;

lib/Basic/Program.cpp

@@ -18,7 +18,11 @@
 #include "llvm/Config/config.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/Program.h"
+#if defined(_WIN32)
+#include "llvm/Support/Windows/WindowsSupport.h"
+#endif
 
+#include <memory>
 #include <system_error>
 
 #if HAVE_POSIX_SPAWN
@@ -29,6 +33,12 @@
 #include <unistd.h>
 #endif
 
+#if defined(_WIN32)
+#define WIN32_LEAN_AND_MEAN
+#include <Windows.h>
+#include <io.h>
+#endif
+
 using namespace swift;
 
 int swift::ExecuteInPlace(const char *Program, const char **args,
@@ -175,7 +185,100 @@ swift::ExecuteWithPipe(llvm::StringRef program,
 #endif
   close(p1.read);
   close(p2.write);
-  return ChildProcessInfo(pid, p1.write, p2.read);
+  llvm::sys::ProcessInfo proc;
+  proc.Pid = pid;
+  return ChildProcessInfo(proc, p1.write, p2.read);
+}
+
+#elif defined(_WIN32)
+
+llvm::ErrorOr<swift::ChildProcessInfo>
+swift::ExecuteWithPipe(llvm::StringRef program,
+                       llvm::ArrayRef<llvm::StringRef> args,
+                       llvm::Optional<llvm::ArrayRef<llvm::StringRef>> env) {
+  using unique_handle = std::unique_ptr<void, decltype(&CloseHandle)>;
+  enum { PI_READ, PI_WRITE };
+
+  unique_handle input[2] = {
+    {INVALID_HANDLE_VALUE, CloseHandle},
+    {INVALID_HANDLE_VALUE, CloseHandle},
+  };
+  unique_handle output[2] = {
+    {INVALID_HANDLE_VALUE, CloseHandle},
+    {INVALID_HANDLE_VALUE, CloseHandle},
+  };
+  unique_handle error{INVALID_HANDLE_VALUE, CloseHandle};
+  HANDLE hRead = INVALID_HANDLE_VALUE, hWrite = INVALID_HANDLE_VALUE;
+  SECURITY_ATTRIBUTES saAttrs{sizeof(SECURITY_ATTRIBUTES), NULL, TRUE};
+
+  if (!CreatePipe(&hRead, &hWrite, &saAttrs, 0))
+    return std::error_code(GetLastError(), std::system_category());
+  output[PI_READ].reset(hRead);
+  output[PI_WRITE].reset(hWrite);
+
+  if (!SetHandleInformation(output[PI_READ].get(), HANDLE_FLAG_INHERIT, FALSE))
+    return std::error_code(GetLastError(), std::system_category());
+
+  if (!CreatePipe(&hRead, &hWrite, &saAttrs, 0))
+    return std::error_code(GetLastError(), std::system_category());
+  input[PI_READ].reset(hRead);
+  input[PI_WRITE].reset(hWrite);
+
+  if (!SetHandleInformation(input[PI_WRITE].get(), HANDLE_FLAG_INHERIT, FALSE))
+    return std::error_code(GetLastError(), std::system_category());
+
+  if (!DuplicateHandle(GetCurrentProcess(), GetStdHandle(STD_ERROR_HANDLE),
+                       GetCurrentProcess(), &hWrite, DUPLICATE_SAME_ACCESS,
+                       TRUE, DUPLICATE_SAME_ACCESS))
+    return std::error_code(GetLastError(), std::system_category());
+  error.reset(hWrite);
+
+  STARTUPINFO si = {0};
+  si.cb = sizeof(si);
+  si.hStdInput = input[PI_READ].get();
+  si.hStdOutput = output[PI_WRITE].get();
+  si.hStdError = error.get();
+  si.dwFlags = STARTF_USESTDHANDLES;
+
+  llvm::SmallVector<wchar_t, MAX_PATH> executable;
+  if (std::error_code ec = llvm::sys::windows::widenPath(program, executable))
+    return ec;
+
+  std::vector<StringRef> components;
+  components.push_back(program);
+  components.assign(args.begin(), args.end());
+  llvm::ErrorOr<std::wstring> commandline =
+      llvm::sys::flattenWindowsCommandLine(components);
+  if (!commandline)
+    return commandline.getError();
+
+  std::vector<wchar_t> command(commandline->size() + 1, 0);
+  std::copy(commandline->begin(), commandline->end(), command.begin());
+
+  PROCESS_INFORMATION pi = {0};
+  if (!CreateProcessW(executable.data(),
+                      command.data(), nullptr, nullptr, TRUE, 0, nullptr,
+                      nullptr, &si, &pi))
+    return std::error_code(GetLastError(), std::system_category());
+
+  unique_handle hThread{pi.hThread, CloseHandle};
+  unique_handle hProcess{pi.hProcess, CloseHandle};
+
+  int ifd = _open_osfhandle(reinterpret_cast<intptr_t>(input[PI_WRITE].get()), 0);
+  if (ifd < 0)
+    return std::error_code(errno, std::system_category());
+  input[PI_WRITE].release();
+
+  int ofd = _open_osfhandle(reinterpret_cast<intptr_t>(output[PI_READ].get()), 0);
+  if (ofd < 0) {
+    _close(ifd);
+    return std::error_code(errno, std::system_category());
+  }
+  output[PI_READ].release();
+
+  llvm::sys::ProcessInfo proc;
+  proc.Process = pi.hProcess;
+  return ChildProcessInfo(proc, ifd, ofd);
 }
 
 #else // HAVE_UNISTD_H