Add Breakpoint Label frame to optimize debug stepping performance (#2190)

* Add Breakpoint label frame to optimize debug stepping performance
* Incorporate PR Review Feedback
* Add Tests

Co-authored-by: Patrick Meinecke <SeeminglyScience@users.noreply.github.com>

Author: JustinGrote

src/PowerShellEditorServices/Server/PsesDebugServer.cs

@@ -117,6 +117,7 @@ public async Task StartAsync()
                         response.SupportsHitConditionalBreakpoints = true;
                         response.SupportsLogPoints = true;
                         response.SupportsSetVariable = true;
+                        response.SupportsDelayedStackTraceLoading = true;
 
                         return Task.CompletedTask;
                     });

src/PowerShellEditorServices/Services/DebugAdapter/DebugService.cs

@@ -6,8 +6,6 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Management.Automation;
-using System.Management.Automation.Language;
-using System.Reflection;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.Extensions.Logging;
@@ -50,7 +48,6 @@ internal class DebugService
         private VariableContainerDetails scriptScopeVariables;
         private VariableContainerDetails localScopeVariables;
         private StackFrameDetails[] stackFrameDetails;
-        private readonly PropertyInfo invocationTypeScriptPositionProperty;
 
         private readonly SemaphoreSlim debugInfoHandle = AsyncUtils.CreateSimpleLockingSemaphore();
         #endregion
@@ -75,6 +72,11 @@ internal class DebugService
         /// </summary>
         public DebuggerStoppedEventArgs CurrentDebuggerStoppedEventArgs { get; private set; }
 
+        /// <summary>
+        /// Returns a task that completes when script frames and variables have completed population
+        /// </summary>
+        public Task StackFramesAndVariablesFetched { get; private set; }
+
         /// <summary>
         /// Tracks whether we are running <c>Debug-Runspace</c> in an out-of-process runspace.
         /// </summary>
@@ -111,12 +113,6 @@ public DebugService(
             _debugContext.DebuggerResuming += OnDebuggerResuming;
             _debugContext.BreakpointUpdated += OnBreakpointUpdated;
             _remoteFileManager = remoteFileManager;
-
-            invocationTypeScriptPositionProperty =
-                typeof(InvocationInfo)
-                    .GetProperty(
-                        "ScriptPosition",
-                        BindingFlags.NonPublic | BindingFlags.Instance);
         }
 
         #endregion
@@ -981,8 +977,8 @@ await _executionService.ExecutePSCommandAsync<PSObject>(
                     }
                 }
 
-                // Get call stack and variables.
-                await FetchStackFramesAndVariablesAsync(noScriptName ? localScriptPath : null).ConfigureAwait(false);
+                // Begin call stack and variables fetch. We don't need to block here.
+                StackFramesAndVariablesFetched = FetchStackFramesAndVariablesAsync(noScriptName ? localScriptPath : null);
 
                 // If this is a remote connection and the debugger stopped at a line
                 // in a script file, get the file contents
@@ -996,53 +992,6 @@ await _remoteFileManager.FetchRemoteFileAsync(
                             _psesHost.CurrentRunspace).ConfigureAwait(false);
                 }
 
-                if (stackFrameDetails.Length > 0)
-                {
-                    // Augment the top stack frame with details from the stop event
-                    if (invocationTypeScriptPositionProperty.GetValue(e.InvocationInfo) is IScriptExtent scriptExtent)
-                    {
-                        StackFrameDetails targetFrame = stackFrameDetails[0];
-
-                        // Certain context changes (like stepping into the default value expression
-                        // of a parameter) do not create a call stack frame. In order to represent
-                        // this context change we create a fake call stack frame.
-                        if (!string.IsNullOrEmpty(scriptExtent.File)
-                            && !PathUtils.IsPathEqual(scriptExtent.File, targetFrame.ScriptPath))
-                        {
-                            await debugInfoHandle.WaitAsync().ConfigureAwait(false);
-                            try
-                            {
-                                targetFrame = new StackFrameDetails
-                                {
-                                    ScriptPath = scriptExtent.File,
-                                    // Just use the last frame's variables since we don't have a
-                                    // good way to get real values.
-                                    AutoVariables = targetFrame.AutoVariables,
-                                    CommandVariables = targetFrame.CommandVariables,
-                                    // Ideally we'd get a real value here but since there's no real
-                                    // call stack frame for this, we'd need to replicate a lot of
-                                    // engine code.
-                                    FunctionName = "<ScriptBlock>",
-                                };
-
-                                StackFrameDetails[] newFrames = new StackFrameDetails[stackFrameDetails.Length + 1];
-                                newFrames[0] = targetFrame;
-                                stackFrameDetails.CopyTo(newFrames, 1);
-                                stackFrameDetails = newFrames;
-                            }
-                            finally
-                            {
-                                debugInfoHandle.Release();
-                            }
-                        }
-
-                        targetFrame.StartLineNumber = scriptExtent.StartLineNumber;
-                        targetFrame.EndLineNumber = scriptExtent.EndLineNumber;
-                        targetFrame.StartColumnNumber = scriptExtent.StartColumnNumber;
-                        targetFrame.EndColumnNumber = scriptExtent.EndColumnNumber;
-                    }
-                }
-
                 CurrentDebuggerStoppedEventArgs =
                     new DebuggerStoppedEventArgs(
                         e,

src/PowerShellEditorServices/Services/DebugAdapter/Handlers/ScopesHandler.cs

@@ -18,16 +18,27 @@ internal class ScopesHandler : IScopesHandler
 
         public ScopesHandler(DebugService debugService) => _debugService = debugService;
 
+        /// <summary>
+        /// Retrieves the variable scopes (containers) for the currently selected stack frame. Variables details are fetched via a separate request.
+        /// </summary>
         public Task<ScopesResponse> Handle(ScopesArguments request, CancellationToken cancellationToken)
         {
-            VariableScope[] variableScopes =
-                _debugService.GetVariableScopes(
-                    (int)request.FrameId);
+            // HACK: The StackTraceHandler injects an artificial label frame as the first frame as a performance optimization, so when scopes are requested by the client, we need to adjust the frame index accordingly to match the underlying PowerShell frame, so when the client clicks on the label (or hit the default breakpoint), they get variables populated from the top of the PowerShell stackframe. If the client dives deeper, we need to reflect that as well (though 90% of debug users don't actually investigate this)
+            // VSCode Frame 0 (Label) -> PowerShell StackFrame 0 (for convenience)
+            // VSCode Frame 1 (First Real PS Frame) -> Also PowerShell StackFrame 0
+            // VSCode Frame 2 -> PowerShell StackFrame 1
+            // VSCode Frame 3 -> PowerShell StackFrame 2
+            // etc.
+            int powershellFrameId = request.FrameId == 0 ? 0 : (int)request.FrameId - 1;
+
+            VariableScope[] variableScopes = _debugService.GetVariableScopes(powershellFrameId);
 
             return Task.FromResult(new ScopesResponse
             {
-                Scopes = new Container<Scope>(variableScopes
-                    .Select(LspDebugUtils.CreateScope))
+                Scopes = new Container<Scope>(
+                    variableScopes
+                    .Select(LspDebugUtils.CreateScope)
+                )
             });
         }
     }

src/PowerShellEditorServices/Services/DebugAdapter/Handlers/StackTraceHandler.cs

@@ -1,64 +1,133 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
+#nullable enable
 
 using System;
 using System.Collections.Generic;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Management.Automation;
 using Microsoft.PowerShell.EditorServices.Services;
-using Microsoft.PowerShell.EditorServices.Services.DebugAdapter;
-using Microsoft.PowerShell.EditorServices.Utility;
 using OmniSharp.Extensions.DebugAdapter.Protocol.Models;
 using OmniSharp.Extensions.DebugAdapter.Protocol.Requests;
+using Microsoft.PowerShell.EditorServices.Services.DebugAdapter;
+using System.Linq;
+using OmniSharp.Extensions.JsonRpc;
+
+namespace Microsoft.PowerShell.EditorServices.Handlers;
 
-namespace Microsoft.PowerShell.EditorServices.Handlers
+internal class StackTraceHandler(DebugService debugService) : IStackTraceHandler
 {
-    internal class StackTraceHandler : IStackTraceHandler
+    /// <summary>
+    /// Because we don't know the size of the stacktrace beforehand, we will tell the client that there are more frames available, this is effectively a paging size, as the client should request this many frames after the first one.
+    /// </summary>
+    private const int INITIAL_PAGE_SIZE = 20;
+
+    public async Task<StackTraceResponse> Handle(StackTraceArguments request, CancellationToken cancellationToken)
     {
-        private readonly DebugService _debugService;
+        if (!debugService.IsDebuggerStopped)
+        {
+            throw new RpcErrorException(0, null!, "Stacktrace was requested while we are not stopped at a breakpoint. This is a violation of the DAP protocol, and is probably a bug.");
+        }
 
-        public StackTraceHandler(DebugService debugService) => _debugService = debugService;
+        // Adapting to int to let us use LINQ, realistically if you have a stacktrace larger than this that the client is requesting, you have bigger problems...
+        int skip = Convert.ToInt32(request.StartFrame ?? 0);
+        int take = Convert.ToInt32(request.Levels ?? 0);
 
-        public async Task<StackTraceResponse> Handle(StackTraceArguments request, CancellationToken cancellationToken)
-        {
-            StackFrameDetails[] stackFrameDetails = await _debugService.GetStackFramesAsync(cancellationToken).ConfigureAwait(false);
+        // We generate a label for the breakpoint and can return that immediately if the client is supporting DelayedStackTraceLoading.
+        InvocationInfo invocationInfo = debugService.CurrentDebuggerStoppedEventArgs?.OriginalEvent?.InvocationInfo
+            ?? throw new RpcErrorException(0, null!, "InvocationInfo was not available on CurrentDebuggerStoppedEvent args. This is a bug and you should report it.");
 
-            // Handle a rare race condition where the adapter requests stack frames before they've
-            // begun building.
-            if (stackFrameDetails is null)
-            {
-                return new StackTraceResponse
-                {
-                    StackFrames = Array.Empty<StackFrame>(),
-                    TotalFrames = 0
-                };
-            }
-
-            List<StackFrame> newStackFrames = new();
-
-            long startFrameIndex = request.StartFrame ?? 0;
-            long maxFrameCount = stackFrameDetails.Length;
-
-            // If the number of requested levels == 0 (or null), that means get all stack frames
-            // after the specified startFrame index. Otherwise get all the stack frames.
-            long requestedFrameCount = request.Levels ?? 0;
-            if (requestedFrameCount > 0)
-            {
-                maxFrameCount = Math.Min(maxFrameCount, startFrameIndex + requestedFrameCount);
-            }
+        StackFrame breakpointLabel = CreateBreakpointLabel(invocationInfo);
 
-            for (long i = startFrameIndex; i < maxFrameCount; i++)
+        if (skip == 0 && take == 1) // This indicates the client is doing an initial fetch, so we want to return quickly to unblock the UI and wait on the remaining stack frames for the subsequent requests.
+        {
+            return new StackTraceResponse()
             {
-                // Create the new StackFrame object with an ID that can
-                // be referenced back to the current list of stack frames
-                newStackFrames.Add(LspDebugUtils.CreateStackFrame(stackFrameDetails[i], id: i));
-            }
+                StackFrames = new StackFrame[] { breakpointLabel },
+                TotalFrames = INITIAL_PAGE_SIZE //Indicate to the client that there are more frames available
+            };
+        }
+
+        // Wait until the stack frames and variables have been fetched.
+        await debugService.StackFramesAndVariablesFetched.ConfigureAwait(false);
+
+        StackFrameDetails[] stackFrameDetails = await debugService.GetStackFramesAsync(cancellationToken)
+            .ConfigureAwait(false);
 
+        // Handle a rare race condition where the adapter requests stack frames before they've
+        // begun building.
+        if (stackFrameDetails is null)
+        {
             return new StackTraceResponse
             {
-                StackFrames = newStackFrames,
-                TotalFrames = newStackFrames.Count
+                StackFrames = Array.Empty<StackFrame>(),
+                TotalFrames = 0
+            };
+        }
+
+        List<StackFrame> newStackFrames = new();
+        if (skip == 0)
+        {
+            newStackFrames.Add(breakpointLabel);
+        }
+
+        newStackFrames.AddRange(
+            stackFrameDetails
+            .Skip(skip != 0 ? skip - 1 : skip)
+            .Take(take != 0 ? take - 1 : take)
+            .Select((frame, index) => CreateStackFrame(frame, index + 1))
+        );
+
+        return new StackTraceResponse
+        {
+            StackFrames = newStackFrames,
+            TotalFrames = newStackFrames.Count
+        };
+    }
+
+    public static StackFrame CreateStackFrame(StackFrameDetails stackFrame, long id)
+    {
+        SourcePresentationHint sourcePresentationHint =
+            stackFrame.IsExternalCode ? SourcePresentationHint.Deemphasize : SourcePresentationHint.Normal;
+
+        // When debugging an interactive session, the ScriptPath is <No File> which is not a valid source file.
+        // We need to make sure the user can't open the file associated with this stack frame.
+        // It will generate a VSCode error in this case.
+        Source? source = null;
+        if (!stackFrame.ScriptPath.Contains("<No File>"))
+        {
+            source = new Source
+            {
+                Path = stackFrame.ScriptPath,
+                PresentationHint = sourcePresentationHint
             };
         }
+
+        return new StackFrame
+        {
+            Id = id,
+            Name = (source is not null) ? stackFrame.FunctionName : "Interactive Session",
+            Line = (source is not null) ? stackFrame.StartLineNumber : 0,
+            EndLine = stackFrame.EndLineNumber,
+            Column = (source is not null) ? stackFrame.StartColumnNumber : 0,
+            EndColumn = stackFrame.EndColumnNumber,
+            Source = source
+        };
     }
+
+    public static StackFrame CreateBreakpointLabel(InvocationInfo invocationInfo, int id = 0) => new()
+    {
+        Name = "<Breakpoint>",
+        Id = id,
+        Source = new()
+        {
+            Path = invocationInfo.ScriptName
+        },
+        Line = invocationInfo.ScriptLineNumber,
+        Column = invocationInfo.OffsetInLine,
+        PresentationHint = StackFramePresentationHint.Label
+    };
+
 }
+

src/PowerShellEditorServices/Utility/LspDebugUtils.cs

@@ -56,38 +56,6 @@ public static Breakpoint CreateBreakpoint(
             };
         }
 
-        public static StackFrame CreateStackFrame(
-            StackFrameDetails stackFrame,
-            long id)
-        {
-            SourcePresentationHint sourcePresentationHint =
-                stackFrame.IsExternalCode ? SourcePresentationHint.Deemphasize : SourcePresentationHint.Normal;
-
-            // When debugging an interactive session, the ScriptPath is <No File> which is not a valid source file.
-            // We need to make sure the user can't open the file associated with this stack frame.
-            // It will generate a VSCode error in this case.
-            Source source = null;
-            if (!stackFrame.ScriptPath.Contains("<"))
-            {
-                source = new Source
-                {
-                    Path = stackFrame.ScriptPath,
-                    PresentationHint = sourcePresentationHint
-                };
-            }
-
-            return new StackFrame
-            {
-                Id = id,
-                Name = (source != null) ? stackFrame.FunctionName : "Interactive Session",
-                Line = (source != null) ? stackFrame.StartLineNumber : 0,
-                EndLine = stackFrame.EndLineNumber,
-                Column = (source != null) ? stackFrame.StartColumnNumber : 0,
-                EndColumn = stackFrame.EndColumnNumber,
-                Source = source
-            };
-        }
-
         public static Scope CreateScope(VariableScope scope)
         {
             return new Scope

test/PowerShellEditorServices.Test.E2E/DebugAdapterProtocolMessageTests.cs

@@ -2,22 +2,31 @@
 // Licensed under the MIT License.
 
 using System;
+using System.Diagnostics;
 using System.IO;
 using System.Linq;
 using System.Runtime.InteropServices;
 using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.Extensions.Logging;
-using Microsoft.Extensions.Logging.Debug;
 using OmniSharp.Extensions.DebugAdapter.Client;
+using DapStackFrame = OmniSharp.Extensions.DebugAdapter.Protocol.Models.StackFrame;
+using OmniSharp.Extensions.DebugAdapter.Protocol.Events;
 using OmniSharp.Extensions.DebugAdapter.Protocol.Models;
 using OmniSharp.Extensions.DebugAdapter.Protocol.Requests;
+using OmniSharp.Extensions.JsonRpc.Server;
 using Xunit;
 using Xunit.Abstractions;
+using Microsoft.Extensions.Logging.Abstractions;
 
 namespace PowerShellEditorServices.Test.E2E
 {
+    public class XunitOutputTraceListener(ITestOutputHelper output) : TraceListener
+    {
+        public override void Write(string message) => output.WriteLine(message);
+        public override void WriteLine(string message) => output.WriteLine(message);
+    }
+
     [Trait("Category", "DAP")]
     public class DebugAdapterProtocolMessageTests : IAsyncLifetime, IDisposable
     {
@@ -28,16 +37,26 @@ public class DebugAdapterProtocolMessageTests : IAsyncLifetime, IDisposable
         private DebugAdapterClient PsesDebugAdapterClient;
         private PsesStdioProcess _psesProcess;
 
+        /// <summary>
+        /// Completes when the debug adapter is started.
+        /// </summary>
         public TaskCompletionSource<object> Started { get; } = new TaskCompletionSource<object>();
-
+        /// <summary>
+        /// Completes when the first breakpoint is reached.
+        /// </summary>
+        public TaskCompletionSource<StoppedEvent> Stopped { get; } = new TaskCompletionSource<StoppedEvent>();
+
+        /// <summary>
+        /// Constructor. The ITestOutputHelper is injected by xUnit and used to write diagnostic logs.
+        /// </summary>
+        /// <param name="output"></param>
         public DebugAdapterProtocolMessageTests(ITestOutputHelper output) => _output = output;
 
         public async Task InitializeAsync()
         {
-            LoggerFactory debugLoggerFactory = new();
-            debugLoggerFactory.AddProvider(new DebugLoggerProvider());
+            // NOTE: To see debug logger output, add this line to your test
 
-            _psesProcess = new PsesStdioProcess(debugLoggerFactory, true);
+            _psesProcess = new PsesStdioProcess(new NullLoggerFactory(), true);
             await _psesProcess.Start();
 
             TaskCompletionSource<bool> initialized = new();
@@ -53,18 +72,20 @@ public async Task InitializeAsync()
                 options
                     .WithInput(_psesProcess.OutputStream)
                     .WithOutput(_psesProcess.InputStream)
-                    .ConfigureLogging(builder =>
-                        builder
-                        .AddDebug()
-                        .SetMinimumLevel(LogLevel.Trace)
-                    )
                     // The OnStarted delegate gets run when we receive the _Initialized_ event from the server:
                     // https://microsoft.github.io/debug-adapter-protocol/specification#Events_Initialized
                     .OnStarted((_, _) =>
                     {
                         Started.SetResult(true);
                         return Task.CompletedTask;
                     })
+                    // We use this to create a task we can await to test debugging after a breakpoint has been received.
+                    .OnNotification<StoppedEvent>(null, (stoppedEvent, _) =>
+                    {
+                        Console.WriteLine("StoppedEvent received");
+                        Stopped.SetResult(stoppedEvent);
+                        return Task.CompletedTask;
+                    })
                     // The OnInitialized delegate gets run when we first receive the _Initialize_ response:
                     // https://microsoft.github.io/debug-adapter-protocol/specification#Requests_Initialize
                     .OnInitialized((_, _, _, _) =>
@@ -263,6 +284,83 @@ public async Task CanSetBreakpointsAsync()
                 (i) => Assert.Equal("after breakpoint", i));
         }
 
+        [SkippableFact]
+        public async Task FailsIfStacktraceRequestedWhenNotPaused()
+        {
+            Skip.If(PsesStdioProcess.RunningInConstrainedLanguageMode,
+                "Breakpoints can't be set in Constrained Language Mode.");
+            string filePath = NewTestFile(GenerateScriptFromLoggingStatements(
+                "labelTestBreakpoint"
+            ));
+            // Set a breakpoint
+            await PsesDebugAdapterClient.SetBreakpoints(
+                new SetBreakpointsArguments
+                {
+                    Source = new Source { Name = Path.GetFileName(filePath), Path = filePath },
+                    Breakpoints = new SourceBreakpoint[] { new SourceBreakpoint { Line = 1 } },
+                    SourceModified = false,
+                }
+            );
+
+            // Signal to start the script
+            await PsesDebugAdapterClient.RequestConfigurationDone(new ConfigurationDoneArguments());
+            await PsesDebugAdapterClient.LaunchScript(filePath, Started);
+
+
+            // Get the stacktrace for the breakpoint
+            await Assert.ThrowsAsync<JsonRpcException>(() => PsesDebugAdapterClient.RequestStackTrace(
+                new StackTraceArguments { }
+            ));
+        }
+
+        [SkippableFact]
+        public async Task SendsInitialLabelBreakpointForPerformanceReasons()
+        {
+            Skip.If(PsesStdioProcess.RunningInConstrainedLanguageMode,
+                "Breakpoints can't be set in Constrained Language Mode.");
+            string filePath = NewTestFile(GenerateScriptFromLoggingStatements(
+                "before breakpoint",
+                "at breakpoint",
+                "after breakpoint"
+            ));
+
+            //TODO: This is technically wrong per the spec, configDone should be completed BEFORE launching, but this is how the vscode client does it today and we really need to fix that.
+            await PsesDebugAdapterClient.LaunchScript(filePath, Started);
+
+            // {"command":"setBreakpoints","arguments":{"source":{"name":"dfsdfg.ps1","path":"/Users/tyleonha/Code/PowerShell/Misc/foo/dfsdfg.ps1"},"lines":[2],"breakpoints":[{"line":2}],"sourceModified":false},"type":"request","seq":3}
+            SetBreakpointsResponse setBreakpointsResponse = await PsesDebugAdapterClient.SetBreakpoints(new SetBreakpointsArguments
+            {
+                Source = new Source { Name = Path.GetFileName(filePath), Path = filePath },
+                Breakpoints = new SourceBreakpoint[] { new SourceBreakpoint { Line = 2 } },
+                SourceModified = false,
+            });
+
+            Breakpoint breakpoint = setBreakpointsResponse.Breakpoints.First();
+            Assert.True(breakpoint.Verified);
+            Assert.Equal(filePath, breakpoint.Source.Path, ignoreCase: s_isWindows);
+            Assert.Equal(2, breakpoint.Line);
+
+            ConfigurationDoneResponse configDoneResponse = await PsesDebugAdapterClient.RequestConfigurationDone(new ConfigurationDoneArguments());
+
+            // FIXME: I think there is a race condition here. If you remove this, the following line Stack Trace fails because the breakpoint hasn't been hit yet. I think the whole getLog process just works long enough for ConfigurationDone to complete and for the breakpoint to be hit.
+
+            // I've tried to do this properly by waiting for a StoppedEvent, but that doesn't seem to work, I'm probably just not wiring it up right in the handler.
+            Assert.NotNull(configDoneResponse);
+            Assert.Collection(await GetLog(),
+                (i) => Assert.Equal("before breakpoint", i));
+            File.Delete(s_testOutputPath);
+
+            // Get the stacktrace for the breakpoint
+            StackTraceResponse stackTraceResponse = await PsesDebugAdapterClient.RequestStackTrace(
+                new StackTraceArguments { ThreadId = 1 }
+            );
+            DapStackFrame firstFrame = stackTraceResponse.StackFrames.First();
+            Assert.Equal(
+                firstFrame.PresentationHint,
+                StackFramePresentationHint.Label
+            );
+        }
+
         // This is a regression test for a bug where user code causes a new synchronization context
         // to be created, breaking the extension. It's most evident when debugging PowerShell
         // scripts that use System.Windows.Forms. It required fixing both Editor Services and