Kill child processes of a Kernel process when existing Kernel Proc

package.json

@@ -2230,6 +2230,7 @@
         "node-fetch": "^2.6.7",
         "node-stream-zip": "^1.6.0",
         "pdfkit": "^0.13.0",
+        "pidtree": "^0.6.0",
         "portfinder": "^1.0.25",
         "re-resizable": "~6.5.5",
         "react": "^16.5.2",

src/kernels/raw/launcher/kernelLauncher.node.ts

@@ -36,6 +36,7 @@ import { getDisplayPathFromLocalFile } from '../../../platform/common/platform/f
 import { noop } from '../../../platform/common/utils/misc';
 import { sendKernelTelemetryWhenDone } from '../../telemetry/sendKernelTelemetryEvent';
 import { PythonKernelInterruptDaemon } from '../finder/pythonKernelInterruptDaemon.node';
+import { IPlatformService } from '../../../platform/common/platform/types';
 
 const PortFormatString = `kernelLauncherPortStart_{0}.tmp`;
 // Launches and returns a kernel process given a resource or python interpreter.
@@ -60,7 +61,8 @@ export class KernelLauncher implements IKernelLauncher {
         @inject(IPythonExecutionFactory) private readonly pythonExecFactory: IPythonExecutionFactory,
         @inject(IConfigurationService) private readonly configService: IConfigurationService,
         @inject(JupyterPaths) private readonly jupyterPaths: JupyterPaths,
-        @inject(PythonKernelInterruptDaemon) private readonly pythonKernelInterruptDaemon: PythonKernelInterruptDaemon
+        @inject(PythonKernelInterruptDaemon) private readonly pythonKernelInterruptDaemon: PythonKernelInterruptDaemon,
+        @inject(IPlatformService) private readonly platformService: IPlatformService
     ) {}
 
     public static async cleanupStartPort() {
@@ -213,7 +215,8 @@ export class KernelLauncher implements IKernelLauncher {
             outputChannel,
             jupyterSettings,
             this.jupyterPaths,
-            this.pythonKernelInterruptDaemon
+            this.pythonKernelInterruptDaemon,
+            this.platformService
         );
 
         try {
@@ -222,7 +225,7 @@ export class KernelLauncher implements IKernelLauncher {
                 createPromiseFromCancellation({ token: cancelToken, cancelAction: 'reject' })
             ]);
         } catch (ex) {
-            kernelProcess.dispose();
+            await kernelProcess.dispose();
             Cancellation.throwIfCanceled(cancelToken);
             throw ex;
         }

src/kernels/raw/launcher/kernelProcess.node.ts

@@ -49,7 +49,7 @@ import {
     IProcessService
 } from '../../../platform/common/process/types.node';
 import { Resource, IOutputChannel, IJupyterSettings } from '../../../platform/common/types';
-import { createDeferred } from '../../../platform/common/utils/async';
+import { createDeferred, sleep } from '../../../platform/common/utils/async';
 import { DataScience } from '../../../platform/common/utils/localize';
 import { noop, swallowExceptions } from '../../../platform/common/utils/misc';
 import { KernelDiedError } from '../../errors/kernelDiedError';
@@ -59,6 +59,9 @@ import { captureTelemetry, Telemetry } from '../../../telemetry';
 import { Interrupter, PythonKernelInterruptDaemon } from '../finder/pythonKernelInterruptDaemon.node';
 import { TraceOptions } from '../../../platform/logging/types';
 import { JupyterPaths } from '../finder/jupyterPaths.node';
+import { ProcessService } from '../../../platform/common/process/proc.node';
+import { IPlatformService } from '../../../platform/common/platform/types';
+import pidtree from 'pidtree';
 
 // Launches and disposes a kernel process given a kernelspec and a resource or python interpreter.
 // Exposes connection information and the process itself.
@@ -105,7 +108,8 @@ export class KernelProcess implements IKernelProcess {
         private readonly outputChannel: IOutputChannel | undefined,
         private readonly jupyterSettings: IJupyterSettings,
         private readonly jupyterPaths: JupyterPaths,
-        private readonly pythonKernelInterruptDaemon: PythonKernelInterruptDaemon
+        private readonly pythonKernelInterruptDaemon: PythonKernelInterruptDaemon,
+        private readonly platform: IPlatformService
     ) {
         this._kernelConnectionMetadata = kernelConnectionMetadata;
     }
@@ -246,7 +250,7 @@ export class KernelProcess implements IKernelProcess {
                 traceError(stderrProc || stderr);
             }
             // Make sure to dispose if we never connect.
-            this.dispose();
+            await this.dispose();
 
             if (!cancelToken?.isCancellationRequested && e instanceof BaseError) {
                 throw e;
@@ -272,12 +276,16 @@ export class KernelProcess implements IKernelProcess {
         }
     }
 
-    public dispose(): void {
+    public async dispose(): Promise<void> {
         if (this.disposed) {
             return;
         }
         traceVerbose('Dispose Kernel process');
         this.disposed = true;
+        await Promise.race([
+            sleep(1_000), // Wait for a max of 1s, we don't want to delay killing the kernel process.
+            this.killChildProcesses(this._process?.pid).catch(noop)
+        ]);
         swallowExceptions(() => {
             this.interrupter?.dispose().ignoreErrors();
             this._process?.kill(); // NOSONAR
@@ -292,6 +300,35 @@ export class KernelProcess implements IKernelProcess {
         });
     }
 
+    private async killChildProcesses(pid?: number) {
+        // Do not remove this code, in in unit tests we end up running this,
+        // then we run into the danger of kill all of the processes on the machine.
+        // because calling `pidtree` without a pid will return all pids and hence everything ends up getting killed.
+        if (!pid || !ProcessService.isAlive(pid)) {
+            return;
+        }
+        try {
+            if (this.platform.isWindows) {
+                const windir = process.env['WINDIR'] || 'C:\\Windows';
+                const TASK_KILL = path.join(windir, 'System32', 'taskkill.exe');
+                await new ProcessService().exec(TASK_KILL, ['/F', '/T', '/PID', pid.toString()]);
+            } else {
+                await new Promise<void>((resolve) => {
+                    pidtree(pid, (ex: unknown, pids: number[]) => {
+                        if (ex) {
+                            traceWarning(`Failed to kill children for ${pid}`, ex);
+                        } else {
+                            pids.forEach((procId) => ProcessService.kill(procId));
+                        }
+                        resolve();
+                    });
+                });
+            }
+        } catch (ex) {
+            traceWarning(`Failed to kill children for ${pid}`, ex);
+        }
+    }
+
     private sendToOutput(data: string) {
         if (this.outputChannel) {
             this.outputChannel.append(data);

src/kernels/raw/session/rawJupyterSession.node.ts

@@ -294,7 +294,7 @@ export class RawJupyterSession extends BaseJupyterSession implements IRawKernelC
             traceVerbose('Successfully waited for Raw Session to be ready in postStartRawSession');
         } catch (ex) {
             traceError('Failed waiting for Raw Session to be ready', ex);
-            process.dispose();
+            await process.dispose();
             result.dispose().catch(noop);
             if (isCancellationError(ex) || options.token.isCancellationRequested) {
                 throw new CancellationError();
@@ -329,7 +329,7 @@ export class RawJupyterSession extends BaseJupyterSession implements IRawKernelC
                 ]);
             } catch (ex) {
                 traceError('Failed to request kernel info', ex);
-                process.dispose();
+                await process.dispose();
                 result.dispose().catch(noop);
                 throw ex;
             } finally {

src/kernels/raw/session/rawSession.node.ts

@@ -108,7 +108,7 @@ export class RawSession implements ISessionWithSocket {
                 .shutdown()
                 .catch((ex) => traceWarning(`Failed to shutdown kernel, ${this.kernelConnectionMetadata.id}`, ex));
             this._kernel.dispose();
-            this.kernelProcess.dispose();
+            await this.kernelProcess.dispose();
         }
         try {
             this._kernel.statusChanged.disconnect(this.onKernelStatus, this);

src/kernels/raw/types.ts

@@ -38,7 +38,7 @@ export interface IKernelConnection {
     kernel_name?: string;
 }
 
-export interface IKernelProcess extends IDisposable {
+export interface IKernelProcess extends IAsyncDisposable {
     readonly connection: Readonly<IKernelConnection>;
     readonly kernelConnectionMetadata: Readonly<LocalKernelSpecConnectionMetadata | PythonKernelConnectionMetadata>;
     /**

src/test/datascience/kernelProcess.unit.test.ts

@@ -42,6 +42,7 @@ import { waitForCondition } from '../common.node';
 import { uriEquals } from './helpers';
 import { IS_REMOTE_NATIVE_TEST } from '../constants';
 import { traceInfo } from '../../platform/logging';
+import { IPlatformService } from '../../platform/common/platform/types';
 
 suite('kernel Process', () => {
     let kernelProcess: KernelProcess;
@@ -131,6 +132,8 @@ suite('kernel Process', () => {
             filePath: 'connection.json'
         });
         when(jupyterPaths.getRuntimeDir()).thenResolve();
+        const platform = mock<IPlatformService>();
+        when(platform.isWindows).thenReturn(false);
         kernelProcess = new KernelProcess(
             instance(processServiceFactory),
             connection,
@@ -143,7 +146,8 @@ suite('kernel Process', () => {
             instance(outputChannel),
             instance(jupyterSettings),
             instance(jupyterPaths),
-            instance(daemon)
+            instance(daemon),
+            instance(platform)
         );
     });
     teardown(() => {
@@ -477,6 +481,9 @@ suite('DataScience - Kernel Process', () => {
             interrupt: () => Promise.resolve(),
             handle: 1
         });
+        const platform = mock<IPlatformService>();
+        when(platform.isWindows).thenReturn(false);
+
         return new KernelProcess(
             instance(processExecutionFactory),
             instance(connection),
@@ -489,7 +496,8 @@ suite('DataScience - Kernel Process', () => {
             undefined,
             instance(settings),
             instance(jupyterPaths),
-            instance(interruptDaemon)
+            instance(interruptDaemon),
+            instance(platform)
         );
     }
     test('Launch from kernelspec (linux)', async function () {