From 4dbd9610044ff24c9710d688a96f7f635dbef10d Mon Sep 17 00:00:00 2001
From: Gregg Miskelly <greggm@microsoft.com>
Date: Thu, 21 Sep 2023 17:40:37 -0700
Subject: [PATCH] Wait for attach complete (#6415)

This PR adds code to the Roslyn unit test debugging code to wait for
attach complete.
---
 src/lsptoolshost/roslynLanguageServer.ts | 57 +++++++++++++++++++++---
 1 file changed, 52 insertions(+), 5 deletions(-)

diff --git a/src/lsptoolshost/roslynLanguageServer.ts b/src/lsptoolshost/roslynLanguageServer.ts
index 6e0a0a18c..c676dc921 100644
--- a/src/lsptoolshost/roslynLanguageServer.ts
+++ b/src/lsptoolshost/roslynLanguageServer.ts
@@ -54,6 +54,7 @@ import { registerRazorCommands } from './razorCommands';
 import { registerOnAutoInsert } from './onAutoInsert';
 import { commonOptions, languageServerOptions, omnisharpOptions } from '../shared/options';
 import { NamedPipeInformation } from './roslynProtocol';
+import { IDisposable } from '../disposable';
 
 let _channel: vscode.OutputChannel;
 let _traceChannel: vscode.OutputChannel;
@@ -615,6 +616,46 @@ export class RoslynLanguageServer {
         );
     }
 
+    // eslint-disable-next-line @typescript-eslint/promise-function-async
+    private WaitForAttachCompleteAsync(attachRequestId: string): Promise<boolean> {
+        return new Promise<boolean>((resolve) => {
+            let didTerminateRegistation: IDisposable | null = null;
+            let customEventReg: IDisposable | null = null;
+            let isComplete = false;
+
+            const fire = (result: boolean) => {
+                if (isComplete === false) {
+                    isComplete = true;
+                    didTerminateRegistation?.dispose();
+                    customEventReg?.dispose();
+                    resolve(result);
+                }
+            };
+
+            didTerminateRegistation = vscode.debug.onDidTerminateDebugSession((session: vscode.DebugSession) => {
+                if (session.configuration.attachRequestId !== attachRequestId) {
+                    return;
+                }
+
+                fire(false);
+            });
+
+            customEventReg = vscode.debug.onDidReceiveDebugSessionCustomEvent(
+                (event: vscode.DebugSessionCustomEvent) => {
+                    if (event.session.configuration.attachRequestId !== attachRequestId) {
+                        return;
+                    }
+
+                    if (event.event !== 'attachComplete') {
+                        return;
+                    }
+
+                    fire(true);
+                }
+            );
+        });
+    }
+
     private registerDebuggerAttach() {
         this._languageClient.onRequest<RoslynProtocol.DebugAttachParams, RoslynProtocol.DebugAttachResult, void>(
             RoslynProtocol.DebugAttachRequest.type,
@@ -622,16 +663,22 @@ export class RoslynLanguageServer {
                 const debugOptions = commonOptions.unitTestDebuggingOptions;
                 const debugConfiguration: vscode.DebugConfiguration = {
                     ...debugOptions,
-                    name: '.NET Core Attach',
+                    name: '.NET Debug Unit test',
                     type: 'coreclr',
                     request: 'attach',
                     processId: request.processId,
+                    attachRequestId: randomUUID(),
                 };
 
-                const result = await vscode.debug.startDebugging(undefined, debugConfiguration, undefined);
-                return {
-                    didAttach: result,
-                };
+                const waitCompletePromise = this.WaitForAttachCompleteAsync(debugConfiguration.attachRequestId);
+
+                let success = await vscode.debug.startDebugging(undefined, debugConfiguration, undefined);
+                if (!success) {
+                    return { didAttach: false };
+                }
+
+                success = await waitCompletePromise;
+                return { didAttach: success };
             }
         );
     }