From 2cb84e0056c0237368938a448bcaf6dad5087000 Mon Sep 17 00:00:00 2001
From: Bernd Hufmann <bernd.hufmann@ericsson.com>
Date: Fri, 18 Aug 2023 15:59:17 -0400
Subject: [PATCH] Fix stopping trace server during shutdown

Fixes #13

Signed-off-by: Bernd Hufmann <bernd.hufmann@ericsson.com>
---
 src/extension.ts    |  4 ++--
 src/trace-server.ts | 22 +++++++++++++++++++++-
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/src/extension.ts b/src/extension.ts
index 1720ee0..93c469e 100644
--- a/src/extension.ts
+++ b/src/extension.ts
@@ -14,8 +14,8 @@ export function activate(context: vscode.ExtensionContext) {
     activation = context;
 }
 
-export function deactivate() {
-    server.stopOrReset(activation);
+export async function deactivate() {
+    await server.shutdown(activation);
 }
 
 function registerStopOrReset(context: vscode.ExtensionContext | undefined): vscode.Disposable {
diff --git a/src/trace-server.ts b/src/trace-server.ts
index 02cc1dc..5b8cc53 100644
--- a/src/trace-server.ts
+++ b/src/trace-server.ts
@@ -13,6 +13,7 @@ const millis = 10000;
 const none = -1;
 const prefix = 'Trace Server';
 const suffix = ' failure or so.';
+const SHUTDOWN_DELAY = 2000;
 
 export class TraceServer {
     private server: ChildProcess | undefined;
@@ -71,11 +72,30 @@ export class TraceServer {
         this.server = undefined;
     }
 
+    async shutdown(context: vscode.ExtensionContext) {
+        // Take the pid from the server instead from workspace state because
+        // during shutdown the workspace can't be queried anymore
+        const pid = this.server ? this.server.pid : undefined;
+        if (!pid) {
+            return;
+        }
+        if (pid) {
+            treeKill(pid);
+            // Allow the treeKill to finish collecting and killing all
+            // spawned processes.
+            await new Promise(resolve => setTimeout(resolve, SHUTDOWN_DELAY));
+
+            // Try to reset pid in workspace state
+            await context?.workspaceState.update(key, none);
+        }
+    }
+
     async startIfStopped(context: vscode.ExtensionContext | undefined) {
+        
         const pid = context?.workspaceState.get(key);
+        console.log("Pid: " + pid);
         const stopped = !pid || pid === none;
         const foreigner = await this.isUp();
-
         if (stopped && !foreigner) {
             await this.start(context);
         } else if (foreigner) {