Addressed minor gaps in notebook cell execution

- Correctly printing the runtime errors stacktrace
- Handling null output stream callback
- Performing lineEnding normalization only where required in cell setContent.
- Minor performance improvements

Author: sid-srini

build.xml

@@ -66,9 +66,9 @@
             patches/dev-dependency-licenses.diff
             patches/nb-telemetry.diff
             patches/change-method-parameters-refactoring-qualified-names.diff
-	        patches/upgrade-lsp4j.diff
-	        patches/updated-show-input-params.diff
-	        patches/java-notebooks.diff
+            patches/upgrade-lsp4j.diff
+            patches/updated-show-input-params.diff
+            patches/java-notebooks.diff
         </string>
         <filterchain>
             <tokenfilter delimoutput=" ">

nbcode/notebooks/src/org/netbeans/modules/nbcode/java/notebook/CellState.java

@@ -83,7 +83,6 @@ public String getNotebookUri() {
     }
 
     public void setContent(String newContent, int newVersion) throws InterruptedException, ExecutionException {
-        String normalizedContent = NotebookUtils.normalizeLineEndings(newContent);
         VersionAwareContent currentContent = content.get();
 
         if (currentContent.getVersion() != newVersion - 1) {
@@ -100,13 +99,14 @@ public void setContent(String newContent, int newVersion) throws InterruptedExce
             int receivedVersion = newCellState.getVersion();
 
             if (receivedVersion > currentContent.getVersion()) {
-                VersionAwareContent newVersionContent = new VersionAwareContent(newCellState.getText(), receivedVersion);
+                VersionAwareContent newVersionContent = new VersionAwareContent(NotebookUtils.normalizeLineEndings(newCellState.getText()), receivedVersion);
                 content.updateAndGet(current -> current != currentContent && receivedVersion <= current.getVersion() ? current : newVersionContent);
             } else {
                 LOG.log(Level.WARNING, "Version mismatch: Received version to be greater than current version, received version:  {0}, current version: {1}", new Object[]{receivedVersion, currentContent.getVersion()});
             }
         } else {
-            VersionAwareContent newVersionContent = new VersionAwareContent(normalizedContent, newVersion);
+            // newContent is already normalized during applyChanges
+            VersionAwareContent newVersionContent = new VersionAwareContent(newContent, newVersion);
 
             if (!content.compareAndSet(currentContent, newVersionContent)) {
                 LOG.log(Level.WARNING, "Concurrent modification detected. Version expected: {0}, current: {1}", new Object[]{newVersion - 1, content.get().getVersion()});

nbcode/notebooks/src/org/netbeans/modules/nbcode/java/notebook/CodeEval.java

@@ -15,8 +15,8 @@
  */
 package org.netbeans.modules.nbcode.java.notebook;
 
-import jdk.jshell.JShell;
-import jdk.jshell.SnippetEvent;
+import java.io.PrintWriter;
+import java.io.Writer;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -29,8 +29,11 @@
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import jdk.jshell.Diag;
+import jdk.jshell.EvalException;
+import jdk.jshell.JShell;
 import jdk.jshell.JShellException;
 import jdk.jshell.SourceCodeAnalysis;
+import jdk.jshell.SnippetEvent;
 import org.netbeans.modules.java.lsp.server.notebook.CellExecutionResult;
 import org.netbeans.modules.java.lsp.server.notebook.NotebookCellExecutionProgressResultParams;
 import org.netbeans.modules.java.lsp.server.notebook.NotebookCellExecutionProgressResultParams.Builder;
@@ -228,17 +231,98 @@ private List<String> getRuntimeErrors(SnippetEvent event) {
         JShellException jshellException = event.exception();
         if (jshellException != null) {
             String msg = jshellException.getMessage();
+            boolean msgAdded = false;
             if (msg != null && !msg.isBlank()) {
                 runtimeErrors.add(msg);
+                msgAdded = true;
             }
-            if (!jshellException.fillInStackTrace().toString().isBlank()) {
-                runtimeErrors.add(jshellException.fillInStackTrace().toString());
+            // Getting the exception stacktrace/details:
+            // stacktrace for EvalException provides the exception that the snippet code generated
+            // stacktrace for non-EvalException is not helpful as it is only internal details
+            String stacktrace = jshellException instanceof EvalException
+                    ? getStackTrace((EvalException) jshellException)
+                    : msgAdded ? "" : jshellException.toString();
+            if (!stacktrace.isBlank()) {
+                runtimeErrors.add(stacktrace);
             }
         }
 
         return runtimeErrors;
     }
 
+    private String getStackTrace(EvalException exception) {
+        return printStackTrace(null, exception).toString();
+    }
+
+    private StringBuilder printStackTrace(StringBuilder output, EvalException exception) {
+        StringBuilder sb = printStackTrace(output, (Throwable) exception);
+
+        return correctExceptionName(sb, 0, exception);
+    }
+
+    private StringBuilder correctExceptionName(StringBuilder output, int startIndex, EvalException exception) {
+        // EvalException has the peculiarity that it replaces the actual cause,
+        // while retaining the name, stacktrace and subsequent causes.
+        // This is unhelpful since it hides the actual exception in the output.
+        // Note: jdk.internal.jshell.tool.JShellTool.displayEvalException() uses
+        // elaborate code to perform the user-friendly printing on console.
+        String actualName = exception.getExceptionClassName();
+        String wrapperName = exception.getClass().getName();
+        if (actualName != null && !wrapperName.equals(actualName)) {
+            int foundAt = output.indexOf(wrapperName, startIndex);
+            if (foundAt >= 0) {
+                output.replace(foundAt, foundAt + wrapperName.length(), actualName);
+                foundAt += actualName.length();
+                if (foundAt < output.length()) {
+                    Throwable cause = exception;
+                    Throwable cycleDetector = cause;
+                    do {
+                        cause = cause.getCause();
+
+                        if (cycleDetector != null) {
+                            // Check for loops in cause using a tortoise-hare detector.
+                            cycleDetector = cycleDetector.getCause();
+                            if (cycleDetector != null) {
+                                cycleDetector = cycleDetector.getCause();
+                                if (cycleDetector == cause) {
+                                    cause = null;   // Cycle has been detected; break
+                                }
+                            }
+                        }
+                    } while (cause != null && !(cause instanceof EvalException));
+                    if (cause != null) {
+                        correctExceptionName(output, foundAt, (EvalException) cause);
+                    }
+                }
+            }
+        }
+        return output;
+    }
+    
+    private StringBuilder printStackTrace(StringBuilder output, Throwable exception) {
+        if (exception == null) {
+            return output != null ? output : new StringBuilder(0);
+        }
+        StringBuilder sb = output != null ? output : new StringBuilder();
+        PrintWriter stackWriter = new PrintWriter(new Writer() {
+            @Override
+            public void write(char[] cbuf, int off, int len) {
+                sb.append(cbuf, off, len);
+            }
+            
+            @Override
+            public void flush() {
+            }
+            
+            @Override
+            public void close() {
+            }
+        });
+        exception.printStackTrace(stackWriter);
+        
+        return sb;
+    }
+
     private List<String> getSnippetValue(SnippetEvent event) {
         return event.value() != null ? List.of(event.value()) : Collections.emptyList();
     }
@@ -250,11 +334,11 @@ private void flushStreams(String notebookId) {
         }
     }
 
-    // This method is directly taken from JShell tool implementation in jdk with some minor modifications
+    // Note: This method is taken from jdk.internal.jshell.tool.JShellTool with some simplifications
     private List<String> displayableDiagnostic(String source, Diag diag) {
         List<String> toDisplay = new ArrayList<>();
 
-        for (String line : diag.getMessage(null).split("\\r?\\n")) {
+        for (String line : diag.getMessage(null).split("\\R")) {
             if (!line.trim().startsWith("location:")) {
                 toDisplay.add(line);
             }

nbcode/notebooks/src/org/netbeans/modules/nbcode/java/notebook/CustomInputStream.java

@@ -47,13 +47,12 @@ public CustomInputStream(NbCodeLanguageClient client) {
     @Override
     public synchronized int read(byte[] b, int off, int len) throws IOException {
         try {
-            NbCodeLanguageClient client = this.client.get();
-            if (client == null) {
-                LOG.log(Level.WARNING, "client is null");
-                return -1;
-            }
-
             if (currentStream == null || currentStream.available() == 0) {
+                NbCodeLanguageClient client = this.client.get();
+                if (client == null) {
+                    LOG.log(Level.WARNING, "client is null");
+                    return -1;
+                }
                 CompletableFuture<String> future = client.showInputBox(new ShowInputBoxParams(USER_PROMPT_REQUEST, "", true));
                 String userInput = future.get();
 

nbcode/notebooks/src/org/netbeans/modules/nbcode/java/notebook/JshellStreamsHandler.java

@@ -61,14 +61,6 @@ private Consumer<byte[]> createCallback(BiConsumer<String, byte[]> callback) {
         return callback != null ? output -> callback.accept(notebookId, output) : null;
     }
 
-    public void setOutStreamCallback(BiConsumer<String, byte[]> callback) {
-        outStream.setCallback(createCallback(callback));
-    }
-
-    public void setErrStreamCallback(BiConsumer<String, byte[]> callback) {
-        errStream.setCallback(createCallback(callback));
-    }
-
     public PrintStream getPrintOutStream() {
         return printOutStream;
     }
@@ -89,37 +81,41 @@ public void flushOutputStreams() {
         try {
             outStream.flush();
         } catch (IOException exception) {
-            LOG.log(Level.WARNING, "IOException occurred while flushing out stream: {0}", exception.getMessage());
+            LOG.log(Level.WARNING, "IOException occurred while flushing out stream: {0}", exception.toString());
         }
         try {
             errStream.flush();
         } catch (IOException exception) {
-            LOG.log(Level.WARNING, "IOException occurred while flushing error stream: {0}", exception.getMessage());
+            LOG.log(Level.WARNING, "IOException occurred while flushing error stream: {0}", exception.toString());
         }
     }
 
     @Override
     public void close() {
         try {
             printOutStream.close();
+        } catch (Exception exception) {
+            LOG.log(Level.WARNING, "Exception occurred while closing print out stream: {0}", exception.toString());
+        }
+        try {
             printErrStream.close();
         } catch (Exception exception) {
-            LOG.log(Level.WARNING, "Exception occurred while closing print streams: {0}", exception.getMessage());
+            LOG.log(Level.WARNING, "Exception occurred while closing print err stream: {0}", exception.toString());
         }
         try {
             outStream.close();
         } catch (IOException exception) {
-            LOG.log(Level.WARNING, "IOException occurred while closing out stream: {0}", exception.getMessage());
+            LOG.log(Level.WARNING, "IOException occurred while closing out stream: {0}", exception.toString());
         }
         try {
             errStream.close();
         } catch (IOException exception) {
-            LOG.log(Level.WARNING, "IOException occurred while closing error stream: {0}", exception.getMessage());
+            LOG.log(Level.WARNING, "IOException occurred while closing error stream: {0}", exception.toString());
         }
         try {
             inputStream.close();
         } catch (IOException exception) {
-            LOG.log(Level.WARNING, "IOException occurred while closing input stream: {0}", exception.getMessage());
+            LOG.log(Level.WARNING, "IOException occurred while closing input stream: {0}", exception.toString());
         }
     }
 }

nbcode/notebooks/src/org/netbeans/modules/nbcode/java/notebook/NotebookDocumentServiceHandlerImpl.java

@@ -35,9 +35,6 @@
 import org.eclipse.lsp4j.CompletionList;
 import org.eclipse.lsp4j.CompletionParams;
 import org.eclipse.lsp4j.DefinitionParams;
-import org.eclipse.lsp4j.services.LanguageClient;
-import org.netbeans.modules.java.lsp.server.protocol.NbCodeLanguageClient;
-import org.netbeans.modules.java.lsp.server.notebook.NotebookDocumentServiceHandler;
 import org.eclipse.lsp4j.DidChangeNotebookDocumentParams;
 import org.eclipse.lsp4j.DidCloseNotebookDocumentParams;
 import org.eclipse.lsp4j.DidOpenNotebookDocumentParams;
@@ -71,6 +68,9 @@
 import org.eclipse.lsp4j.WillSaveTextDocumentParams;
 import org.eclipse.lsp4j.jsonrpc.messages.Either;
 import org.eclipse.lsp4j.jsonrpc.messages.Either3;
+import org.eclipse.lsp4j.services.LanguageClient;
+import org.netbeans.modules.java.lsp.server.notebook.NotebookDocumentServiceHandler;
+import org.netbeans.modules.java.lsp.server.protocol.NbCodeLanguageClient;
 import org.netbeans.modules.java.lsp.server.protocol.ShowStatusMessageParams;
 import org.openide.util.NbBundle;
 import org.openide.util.lookup.ServiceProvider;
@@ -133,8 +133,12 @@ public void didSave(DidSaveNotebookDocumentParams params) {
     public void didClose(DidCloseNotebookDocumentParams params) {
         String notebookUri = params.getNotebookDocument().getUri();
         NotebookSessionManager.getInstance().closeSession(notebookUri);
-        notebookStateMap.remove(notebookUri);
-        notebookCellMap.entrySet().removeIf(entry -> notebookUri.equals(entry.getValue()));
+        NotebookDocumentStateManager state = notebookStateMap.remove(notebookUri);
+        if (state != null) {
+            state.getCellsMap().keySet().forEach(notebookCellMap::remove);
+        } else {
+            notebookCellMap.values().removeIf(notebookUri::equals);
+        }
     }
 
     @Override

nbcode/notebooks/src/org/netbeans/modules/nbcode/java/notebook/NotebookDocumentStateManager.java

@@ -121,13 +121,14 @@ private void updateNotebookCellStructure(NotebookDocumentChangeEventCellStructur
         List<NotebookCell> cellsDetail = updatedStructure.getArray().getCells();
 
         if (cellsItem != null && cellsDetail != null) {
-            for (int i = 0; i < cellsDetail.size(); i++) {
-                TextDocumentItem cellItem = cellsItem.get(i);
+            Iterator<NotebookCell> details = cellsDetail.iterator();
+            for (TextDocumentItem cellItem: cellsItem) {
                 String uri = cellItem.getUri();
                 openedCellUris.add(uri);
-
-                addNewCellState(cellsDetail.get(i), cellItem);
                 cellsNotebookMap.put(uri, notebookDoc.getUri());
+                if (details.hasNext()) {
+                    addNewCellState(details.next(), cellItem);
+                }
             }
         }
 

nbcode/notebooks/src/org/netbeans/modules/nbcode/java/notebook/StreamingOutputStream.java

@@ -20,7 +20,6 @@
 import java.io.OutputStream;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
-import org.eclipse.lsp4j.jsonrpc.validation.NonNull;
 import org.openide.util.RequestProcessor;
 import org.openide.util.RequestProcessor.Task;
 
@@ -31,9 +30,10 @@
 public class StreamingOutputStream extends OutputStream {
 
     private final ByteArrayOutputStream buffer = new ByteArrayOutputStream();
-    private volatile Consumer<byte[]> callback;
+    private final Consumer<byte[]> callback;
     private static final int MAX_BUFFER_SIZE = 1024;
     private final AtomicBoolean isPeriodicFlushOutputStream;
+    private final boolean noop;
 
     static RequestProcessor getRequestProcessor() {
         return RPSingleton.instance;
@@ -50,19 +50,22 @@ private static final class RPSingleton {
     }
 
     public StreamingOutputStream(Consumer<byte[]> callback) {
+        this.noop = callback == null;
         this.callback = callback;
-        this.isPeriodicFlushOutputStream = new AtomicBoolean(true);
+        this.isPeriodicFlushOutputStream = new AtomicBoolean(!noop);
         createAndScheduleTask();
     }
 
     @Override
     public synchronized void write(int b) throws IOException {
+        if (noop) return;
         buffer.write(b);
         ifBufferOverflowFlush();
     }
 
     @Override
     public synchronized void write(byte[] b, int off, int len) throws IOException {
+        if (noop) return;
         if (len >= MAX_BUFFER_SIZE) {
             flushToCallback();
             byte[] chunk = new byte[len];
@@ -76,33 +79,34 @@ public synchronized void write(byte[] b, int off, int len) throws IOException {
 
     @Override
     public synchronized void flush() throws IOException {
+        if (noop) return;
         flushToCallback();
     }
 
     @Override
     public synchronized void write(byte[] b) throws IOException {
+        if (noop) return;
         write(b, 0, b.length);
     }
 
     @Override
     public synchronized void close() throws IOException {
-        flushToCallback();
-        isPeriodicFlushOutputStream.set(false);
+        if (!noop) {
+            flushToCallback();
+            isPeriodicFlushOutputStream.set(false);
+        }
         super.close();
     }
 
-    @NonNull
-    public void setCallback(Consumer<byte[]> cb) {
-        this.callback = cb;
-    }
-
     private void ifBufferOverflowFlush() {
+        if (noop) return;
         if (buffer.size() > MAX_BUFFER_SIZE) {
             flushToCallback();
         }
     }
 
     private synchronized void flushToCallback() {
+        if (noop) return;
         if (buffer.size() > 0) {
             byte[] output = buffer.toByteArray();
             buffer.reset();