spawning an executable now only waits for that particular process

added waitForAllUSVFSProcesses() to OrganizerCore, used when closing MO

7 files changed, 233 insertions, 90 deletions

diff --git a/src/envmodule.cpp b/src/envmodule.cpp
index 3f1f8912..abbe02e5 100644
--- a/src/envmodule.cpp
+++ b/src/envmodule.cpp
@@ -408,7 +408,8 @@ std::vector<Module> getLoadedModules()
 }
 
 
-std::vector<Process> getRunningProcesses()
+template <class F>
+void forEachRunningProcess(F&& f)
 {
   HandlePtr snapshot(CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0));
 
@@ -416,7 +417,7 @@ std::vector<Process> getRunningProcesses()
   {
     const auto e = GetLastError();
     log::error("CreateToolhelp32Snapshot() failed, {}", formatSystemMessage(e));
-    return {};
+    return;
   }
 
   PROCESSENTRY32 entry = {};
@@ -427,16 +428,14 @@ std::vector<Process> getRunningProcesses()
   if (!Process32First(snapshot.get(), &entry)) {
     const auto e = GetLastError();
     log::error("Process32First() failed, {}", formatSystemMessage(e));
-    return {};
+    return;
   }
 
-  std::vector<Process> v;
-
   for (;;)
   {
-    v.push_back(Process(
-      entry.th32ProcessID,
-      QString::fromStdWString(entry.szExeFile)));
+    if (!f(entry)) {
+      break;
+    }
 
     // next process
     if (!Process32Next(snapshot.get(), &entry))
@@ -450,8 +449,66 @@ std::vector<Process> getRunningProcesses()
       break;
     }
   }
+}
+
+std::vector<Process> getRunningProcesses()
+{
+  std::vector<Process> v;
+
+  forEachRunningProcess([&](auto&& entry) {
+    v.push_back(Process(
+      entry.th32ProcessID,
+      QString::fromStdWString(entry.szExeFile)));
+
+    return true;
+  });
 
   return v;
 }
 
+QString getProcessName(HANDLE process)
+{
+  const QString badName = "unknown";
+
+  if (process == 0 || process == INVALID_HANDLE_VALUE) {
+    return badName;
+  }
+
+  const DWORD bufferSize = MAX_PATH;
+  wchar_t buffer[bufferSize + 1] = {};
+
+  const auto realSize = ::GetProcessImageFileNameW(process, buffer, bufferSize);
+
+  if (realSize == 0) {
+    const auto e = ::GetLastError();
+    log::error("GetProcessImageFileNameW() failed, {}", formatSystemMessage(e));
+    return badName;
+  }
+
+  auto s = QString::fromWCharArray(buffer, realSize);
+
+  const auto lastSlash = s.lastIndexOf("\\");
+  if (lastSlash != -1) {
+    s = s.mid(lastSlash + 1);
+  }
+
+  return s;
+}
+
+DWORD getProcessParentID(DWORD pid)
+{
+  DWORD ppid = 0;
+
+  forEachRunningProcess([&](auto&& entry) {
+    if (entry.th32ProcessID == pid) {
+      ppid = entry.th32ParentProcessID;
+      return false;
+    }
+
+    return true;
+  });
+
+  return ppid;
+}
+
 } // namespace
diff --git a/src/envmodule.h b/src/envmodule.h
index deb7520f..212f6f7b 100644
--- a/src/envmodule.h
+++ b/src/envmodule.h
@@ -120,6 +120,9 @@ private:
 std::vector<Process> getRunningProcesses();
 std::vector<Module> getLoadedModules();
 
+QString getProcessName(HANDLE process);
+DWORD getProcessParentID(DWORD pid);
+
 } // namespace env
 
 #endif // ENV_MODULE_H
diff --git a/src/mainwindow.cpp b/src/mainwindow.cpp
index 0181a335..bbb63333 100644
--- a/src/mainwindow.cpp
+++ b/src/mainwindow.cpp
@@ -1311,16 +1311,10 @@ bool MainWindow::canExit()
     }
   }
 
-  std::vector<QString> hiddenList;
-  hiddenList.push_back(QFileInfo(QCoreApplication::applicationFilePath()).fileName());
-  HANDLE injected_process_still_running = m_OrganizerCore.findAndOpenAUSVFSProcess(hiddenList, GetCurrentProcessId());
-  if (injected_process_still_running != INVALID_HANDLE_VALUE)
-  {
-    m_exitAfterWait = true;
-    m_OrganizerCore.waitForApplication(injected_process_still_running);
-    if (!m_exitAfterWait) { // if operation cancelled
-      return false;
-    }
+  m_exitAfterWait = true;
+  m_OrganizerCore.waitForAllUSVFSProcessesWithLock();
+  if (!m_exitAfterWait) { // if operation cancelled
+    return false;
   }
 
   setCursor(Qt::WaitCursor);
diff --git a/src/organizercore.cpp b/src/organizercore.cpp
index d3f4a83c..2a228fda 100644
--- a/src/organizercore.cpp
+++ b/src/organizercore.cpp
@@ -34,6 +34,7 @@
 #include "instancemanager.h"
 #include <scriptextender.h>
 #include "previewdialog.h"
+#include "envmodule.h"
 
 #include <QApplication>
 #include <QCoreApplication>
@@ -74,47 +75,6 @@ using namespace MOBase;
 //static
 CrashDumpsType OrganizerCore::m_globalCrashDumpsType = CrashDumpsType::None;
 
-static std::wstring getProcessName(HANDLE process)
-{
-  wchar_t buffer[MAX_PATH];
-  const wchar_t *fileName = L"unknown";
-
-  if (process == nullptr) return fileName;
-
-  if (::GetProcessImageFileNameW(process, buffer, MAX_PATH) != 0) {
-    fileName = wcsrchr(buffer, L'\\');
-    if (fileName == nullptr) {
-      fileName = buffer;
-    }
-    else {
-      fileName += 1;
-    }
-  }
-
-  return fileName;
-}
-
-// Get parent PID for the given process, return 0 on failure
-static DWORD getProcessParentID(DWORD pid)
-{
-  HANDLE th = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
-  PROCESSENTRY32 pe = { 0 };
-  pe.dwSize = sizeof(PROCESSENTRY32);
-
-  DWORD res = 0;
-  if (Process32First(th, &pe))
-    do {
-      if (pe.th32ProcessID == pid) {
-        res = pe.th32ParentProcessID;
-        break;
-      }
-    } while (Process32Next(th, &pe));
-
-  CloseHandle(th);
-
-  return res;
-}
-
 template <typename InputIterator>
 QStringList toStringList(InputIterator current, InputIterator end)
 {
@@ -1348,35 +1308,46 @@ HANDLE OrganizerCore::spawnAndWait(
   return handle;
 }
 
-bool OrganizerCore::waitForProcessCompletionWithLock(
-  HANDLE handle, LPDWORD exitCode)
+void OrganizerCore::withLock(std::function<void (ILockedWaitingForProcess*)> f)
 {
-  if (Settings::instance().interface().lockGUI()) {
-    std::unique_ptr<LockedDialog> dlg;
-    ILockedWaitingForProcess* uilock = nullptr;
+  std::unique_ptr<LockedDialog> dlg;
+  ILockedWaitingForProcess* uilock = nullptr;
+
+  if (m_MainWindow != nullptr) {
+    uilock = m_MainWindow->lock();
+  }
+  else {
+    // i.e. when running command line shortcuts there is no user interface
+    dlg.reset(new LockedDialog);
+    dlg->show();
+    dlg->setEnabled(true);
+    uilock = dlg.get();
+  }
 
+  ON_BLOCK_EXIT([&]() {
     if (m_MainWindow != nullptr) {
-      uilock = m_MainWindow->lock();
-    }
-    else {
-      // i.e. when running command line shortcuts there is no user interface
-      dlg.reset(new LockedDialog);
-      dlg->show();
-      dlg->setEnabled(true);
-      uilock = dlg.get();
-    }
+      m_MainWindow->unlock();
+    } });
 
-    ON_BLOCK_EXIT([&]() {
-      if (m_MainWindow != nullptr) {
-        m_MainWindow->unlock();
-      } });
+  f(uilock);
+}
+
+bool OrganizerCore::waitForProcessCompletionWithLock(
+  HANDLE handle, LPDWORD exitCode)
+{
+  if (!Settings::instance().interface().lockGUI()) {
+    return true;
+  }
+
+  bool r = false;
 
+  withLock([&](auto* uilock) {
     DWORD ignoreExitCode;
-    waitForProcessCompletion(handle, exitCode ? exitCode : &ignoreExitCode, uilock);
+    r = waitForProcessCompletion(handle, exitCode ? exitCode : &ignoreExitCode, uilock);
     cycleDiagnostics();
-  }
+  });
 
-  return handle;
+  return r;
 }
 
 bool OrganizerCore::waitForApplication(HANDLE handle, LPDWORD exitCode)
@@ -1393,30 +1364,64 @@ bool OrganizerCore::waitForApplication(HANDLE handle, LPDWORD exitCode)
     if (m_MainWindow != nullptr) {
       m_MainWindow->unlock();
     } });
+
   return waitForProcessCompletion(handle, exitCode, uilock);
 }
 
-bool OrganizerCore::waitForProcessCompletion(HANDLE handle, LPDWORD exitCode, ILockedWaitingForProcess* uilock)
+bool OrganizerCore::waitForProcessCompletion(
+  HANDLE handle, LPDWORD exitCode, ILockedWaitingForProcess* uilock)
+{
+  const auto r = spawn::waitForProcess(handle, exitCode, uilock);
+
+  switch (r)
+  {
+    case spawn::WaitResults::Completed:  // fall-through
+    case spawn::WaitResults::Unlocked:
+      return true;
+
+    case spawn::WaitResults::Error:  // fall-through
+    default:
+      return false;
+  }
+}
+
+bool OrganizerCore::waitForAllUSVFSProcessesWithLock()
 {
+  bool r = false;
+
+  withLock([&](auto* uilock) {
+    r = waitForAllUSVFSProcesses(uilock);
+  });
+
+  return r;
+}
+
+bool OrganizerCore::waitForAllUSVFSProcesses(ILockedWaitingForProcess* uilock)
+{
+  // Certain process names we wish to "hide" for aesthetic reason:
+  std::vector<QString> hiddenList;
+  hiddenList.push_back(QFileInfo(QCoreApplication::applicationFilePath()).fileName());
+
   bool originalHandle = true;
   bool newHandle = true;
   bool uiunlocked = false;
 
+  HANDLE handle = findAndOpenAUSVFSProcess(hiddenList, GetCurrentProcessId());
+  DWORD* exitCode = nullptr;
   DWORD currentPID = 0;
   QString processName;
+
   auto waitForChildUntil = GetTickCount64();
   if (handle != INVALID_HANDLE_VALUE) {
     currentPID = GetProcessId(handle);
-    processName = QString::fromStdWString(getProcessName(handle));
+    processName = env::getProcessName(handle);
   }
 
-  // Certain process names we wish to "hide" for aesthetic reason:
   bool waitingOnHidden = false;
-  std::vector<QString> hiddenList;
-  hiddenList.push_back(QFileInfo(QCoreApplication::applicationFilePath()).fileName());
   for (QString hide : hiddenList)
     if (processName.contains(hide, Qt::CaseInsensitive))
       waitingOnHidden = true;
+
   // The main reason for adding the hidden list is to hide the MO proxy we use to spawn virtualized processes.
   // On the one hand we want to display the real executable without it feeling laggy, on the other we don't want
   // to requery processes all the time if for some reason we are waiting on hidden processes and find no "unhidden"
@@ -1489,7 +1494,7 @@ bool OrganizerCore::waitForProcessCompletion(HANDLE handle, LPDWORD exitCode, IL
       newHandle = handle != INVALID_HANDLE_VALUE;
       if (newHandle) {
         currentPID = GetProcessId(handle);
-        processName = QString::fromStdWString(getProcessName(handle));
+        processName = env::getProcessName(handle);
         for (QString hide : hiddenList)
           if (processName.contains(hide, Qt::CaseInsensitive))
             waitingOnHidden = true;
@@ -1541,13 +1546,13 @@ HANDLE OrganizerCore::findAndOpenAUSVFSProcess(const std::vector<QString>& hidde
       continue;
     }
 
-    QString pname = QString::fromStdWString(getProcessName(handle));
+    QString pname = env::getProcessName(handle);
     bool phidden = false;
     for (auto hide : hiddenList)
       if (pname.contains(hide, Qt::CaseInsensitive))
         phidden = true;
 
-    bool pprefered = preferedParentPid && getProcessParentID(pids[i]) == preferedParentPid;
+    bool pprefered = preferedParentPid && env::getProcessParentID(pids[i]) == preferedParentPid;
 
     if (best_match == INVALID_HANDLE_VALUE || best_match_hidden || (!phidden && pprefered)) {
       if (best_match != INVALID_HANDLE_VALUE)
diff --git a/src/organizercore.h b/src/organizercore.h
index 3d3c7325..ffdb6830 100644
--- a/src/organizercore.h
+++ b/src/organizercore.h
@@ -167,6 +167,8 @@ public:
     const QString &profile, const QString &forcedCustomOverwrite = "",

     bool ignoreCustomOverwrite = false);

 

+  bool waitForApplication(HANDLE processHandle, LPDWORD exitCode = nullptr);

+  bool waitForAllUSVFSProcessesWithLock();

 

   void loginSuccessfulUpdate(bool necessary);

   void loginFailedUpdate(const QString &message);

@@ -221,8 +223,6 @@ public:
   DownloadManager *downloadManager();

   PluginList *pluginList();

   ModList *modList();

-  bool waitForApplication(HANDLE processHandle, LPDWORD exitCode = nullptr);

-  HANDLE findAndOpenAUSVFSProcess(const std::vector<QString>& hiddenList, DWORD preferedParentPid);

   bool onModInstalled(const std::function<void (const QString &)> &func);

   bool onAboutToRun(const std::function<bool (const QString &)> &func);

   bool onFinishedRun(const std::function<void (const QString &, unsigned int)> &func);

@@ -311,6 +311,13 @@ private:
   bool waitForProcessCompletion(

     HANDLE handle, LPDWORD exitCode, ILockedWaitingForProcess* uilock);

 

+  bool waitForAllUSVFSProcesses(ILockedWaitingForProcess* uilock);

+

+  void withLock(std::function<void (ILockedWaitingForProcess*)> f);

+

+  HANDLE findAndOpenAUSVFSProcess(

+    const std::vector<QString>& hiddenList, DWORD preferedParentPid);

+

 private slots:

 

   void directory_refreshed();

diff --git a/src/spawn.cpp b/src/spawn.cpp
index fe1e9e3e..f0b3b2c7 100644
--- a/src/spawn.cpp
+++ b/src/spawn.cpp
@@ -28,6 +28,7 @@ along with Mod Organizer.  If not, see <http://www.gnu.org/licenses/>.
 #include "settings.h"

 #include "settingsdialogworkarounds.h"

 #include <iplugingame.h>

+#include <ilockedwaitingforprocess.h>

 #include <errorcodes.h>

 #include <report.h>

 #include <log.h>

@@ -1093,6 +1094,68 @@ FileExecutionContext getFileExecutionContext(
   return {{}, {}, FileExecutionTypes::Other};

 }

 

+WaitResults waitForProcess(HANDLE handle, DWORD* exitCode, ILockedWaitingForProcess* uilock)

+{

+  if (handle == INVALID_HANDLE_VALUE) {

+    return WaitResults::Error;

+  }

+

+  const DWORD pid = ::GetProcessId(handle);

+  const QString processName = QString("%1 (%2)")

+    .arg(env::getProcessName(handle))

+    .arg(pid);

+

+  if (uilock)

+    uilock->setProcessName(processName);

+

+  constexpr DWORD INPUT_EVENT = WAIT_OBJECT_0 + 1;

+  DWORD res = WAIT_TIMEOUT;

+

+  log::debug(

+    "waiting for process completion '{}' ({})",

+    processName, pid);

+

+  for (;;) {

+    // Wait for a an event on the handle, a key press, mouse click or timeout

+    const auto res = MsgWaitForMultipleObjects(

+      1, &handle, FALSE, 50, QS_KEY | QS_MOUSEBUTTON);

+

+    if (res == WAIT_FAILED) {

+      // error

+      const auto e = ::GetLastError();

+

+      log::error(

+        "failed waiting for process completion '{}' ({}), {}",

+        processName, pid, formatSystemMessage(e));

+

+      return WaitResults::Error;

+    } else if (res == WAIT_OBJECT_0) {

+      // completed

+      log::debug("process '{}' ({}) completed", processName, pid);

+

+      if (exitCode) {

+        if (!::GetExitCodeProcess(handle, exitCode)) {

+          const auto e = ::GetLastError();

+          log::warn(

+            "failed to get exit code of process '{}' ({}): {}",

+            processName, pid, formatSystemMessage(e));

+        }

+      }

+

+      return WaitResults::Completed;

+    }

+

+    // keep processing events so the app doesn't appear dead

+    QCoreApplication::sendPostedEvents();

+    QCoreApplication::processEvents();

+

+    if (uilock && uilock->unlockForced()) {

+      log::debug("waiting for process '{}' ({}) aborted by UI", processName, pid);

+      return WaitResults::Unlocked;

+    }

+  }

+}

+

 } // namespace

 

 

diff --git a/src/spawn.h b/src/spawn.h
index 866e1795..441cad2c 100644
--- a/src/spawn.h
+++ b/src/spawn.h
@@ -27,6 +27,8 @@ along with Mod Organizer.  If not, see <http://www.gnu.org/licenses/>.
 #include <QDir>

 

 class Settings;

+class ILockedWaitingForProcess;

+

 namespace MOBase { class IPluginGame; }

 

 namespace spawn

@@ -84,6 +86,7 @@ public:
   ~SpawnedProcess();

 

   HANDLE releaseHandle();

+  void wait();

 

 private:

   HANDLE m_handle;

@@ -122,6 +125,17 @@ QString findJavaInstallation(const QString& jarFile);
 FileExecutionContext getFileExecutionContext(

   QWidget* parent, const QFileInfo& target);

 

+

+enum class WaitResults

+{

+  Completed = 1,

+  Error,

+  Unlocked

+};

+

+WaitResults waitForProcess(

+  HANDLE handle, DWORD* exitCode, ILockedWaitingForProcess* uilock);

+

 } // namespace