From a524096fc2fcbeae20f43bcde506b0fd6812e98e Mon Sep 17 00:00:00 2001
From: Hongtao Zhang <hongtao_zhang@outlook.com>
Date: Tue, 6 Feb 2024 13:07:21 -0600
Subject: [PATCH] Use job object to control child process

---
 Flow.Launcher.Core/Flow.Launcher.Core.csproj  |  1 +
 .../Plugin/ProcessStreamPluginV2.cs           | 24 ++++++++++++++-----
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/Flow.Launcher.Core/Flow.Launcher.Core.csproj b/Flow.Launcher.Core/Flow.Launcher.Core.csproj
index 5cd09d407..68e7c115c 100644
--- a/Flow.Launcher.Core/Flow.Launcher.Core.csproj
+++ b/Flow.Launcher.Core/Flow.Launcher.Core.csproj
@@ -55,6 +55,7 @@
   <ItemGroup>
     <PackageReference Include="Droplex" Version="1.7.0" />
     <PackageReference Include="FSharp.Core" Version="7.0.401" />
+    <PackageReference Include="Meziantou.Framework.Win32.Jobs" Version="3.2.1" />
     <PackageReference Include="Microsoft.IO.RecyclableMemoryStream" Version="2.3.2" />
     <PackageReference Include="squirrel.windows" Version="1.5.2" NoWarn="NU1701" />
     <PackageReference Include="StreamJsonRpc" Version="2.17.8" />
diff --git a/Flow.Launcher.Core/Plugin/ProcessStreamPluginV2.cs b/Flow.Launcher.Core/Plugin/ProcessStreamPluginV2.cs
index a476f06e9..16f9dfbf9 100644
--- a/Flow.Launcher.Core/Plugin/ProcessStreamPluginV2.cs
+++ b/Flow.Launcher.Core/Plugin/ProcessStreamPluginV2.cs
@@ -5,12 +5,24 @@ using System.IO.Pipelines;
 using System.Threading.Tasks;
 using Flow.Launcher.Infrastructure;
 using Flow.Launcher.Plugin;
+using Meziantou.Framework.Win32;
 using Nerdbank.Streams;
 
 namespace Flow.Launcher.Core.Plugin
 {
     internal abstract class ProcessStreamPluginV2 : JsonRPCPluginV2
     {
+        private static JobObject _jobObject = new JobObject();
+
+        static ProcessStreamPluginV2()
+        {
+            _jobObject.SetLimits(new JobObjectLimits()
+            {
+                Flags = JobObjectLimitFlags.KillOnJobClose | JobObjectLimitFlags.DieOnUnhandledException
+            });
+            
+            _jobObject.AssignProcess(Process.GetCurrentProcess());
+        }
 
         public override string SupportedLanguage { get; set; }
         protected sealed override IDuplexPipe ClientPipe { get; set; }
@@ -30,22 +42,22 @@ namespace Flow.Launcher.Core.Plugin
 
             ClientProcess = Process.Start(StartInfo);
             ArgumentNullException.ThrowIfNull(ClientProcess);
-
+            
             SetupPipe(ClientProcess);
 
             ErrorStream = ClientProcess.StandardError;
 
             await base.InitAsync(context);
         }
-        
+
         private void SetupPipe(Process process)
         {
             var (reader, writer) = (PipeReader.Create(process.StandardOutput.BaseStream),
                 PipeWriter.Create(process.StandardInput.BaseStream));
             ClientPipe = new DuplexPipe(reader, writer);
         }
-        
-        
+
+
         public override async Task ReloadDataAsync()
         {
             var oldProcess = ClientProcess;
@@ -57,8 +69,8 @@ namespace Flow.Launcher.Core.Plugin
             await oldProcess.WaitForExitAsync();
             oldProcess.Dispose();
         }
-        
-        
+
+
         public override async ValueTask DisposeAsync()
         {
             ClientProcess.Kill(true);