Implement JVM HotSwap via Java Agent

jjava-distro/pom.xml

@@ -132,6 +132,8 @@
                         <transformer implementation="org.apache.maven.plugins.shade.resource.ManifestResourceTransformer">
                             <manifestEntries>
                                 <Main-Class>org.dflib.jjava.distro.JJava</Main-Class>
+                                <Launcher-Agent-Class>org.dflib.jjava.kernel.execution.HotswapAgent</Launcher-Agent-Class>
+                                <Can-Redefine-Classes>true</Can-Redefine-Classes>
                             </manifestEntries>
                         </transformer>
                         <transformer implementation="org.apache.maven.plugins.shade.resource.ServicesResourceTransformer" />

jjava-distro/src/test/java/org/dflib/jjava/distro/KernelExecutionIT.java

@@ -22,4 +22,23 @@ public void variableSurvivesLaterImports() throws Exception {
         assertEquals(0, snippetResult.getExitCode(), snippetResult.getStdout());
         assertThat(snippetResult.getStdout(), not(containsString("|")));
     }
+
+    /**
+     * @see <a href="https://github.com/dflib/jjava/issues/119#issuecomment-4754144547">#119</a>
+     */
+    @Test
+    public void variableRedeclarationUsesLatestValue() throws Exception {
+        String snippet = String.join("\n",
+                "var v = \"first\";",
+                "v",
+                "var v = \"second\";",
+                "v"
+        );
+
+        Container.ExecResult snippetResult = executeInKernel(snippet);
+
+        assertEquals(0, snippetResult.getExitCode(), snippetResult.getStdout());
+        assertThat(snippetResult.getStdout(), containsString("Out[2]: first"));
+        assertThat(snippetResult.getStdout(), containsString("Out[4]: second"));
+    }
 }

jjava-kernel/src/main/java/org/dflib/jjava/kernel/execution/CodeEvaluator.java

@@ -169,9 +169,11 @@ private void dropSnippet(JShell shell, Snippet snippet) {
         // snippet.classFullName() returns name of a wrapper class created for a snippet
         String className = snippetClassName(snippet);
         // check that this class is not used by other snippets
-        if (shell.snippets()
+        boolean isUnused = shell.snippets()
+                .filter(s -> shell.status(s).isActive())
                 .map(this::snippetClassName)
-                .noneMatch(className::equals)) {
+                .noneMatch(className::equals);
+        if (isUnused) {
             execControl.unloadClass(className);
         }
     }

jjava-kernel/src/main/java/org/dflib/jjava/kernel/execution/HotswapAgent.java

@@ -0,0 +1,32 @@
+package org.dflib.jjava.kernel.execution;
+
+import java.lang.instrument.Instrumentation;
+
+/**
+ * Captures the JVM {@link Instrumentation} instance so that {@link JJavaLoaderDelegate} can perform
+ * in-place (HotSwap) redefinition of JShell snippet classes.
+ */
+public final class HotswapAgent {
+
+    private static volatile Instrumentation instrumentation;
+
+    private HotswapAgent() {
+    }
+
+    public static void premain(String agentArgs, Instrumentation inst) {
+        instrumentation = inst;
+    }
+
+    public static void agentmain(String agentArgs, Instrumentation inst) {
+        instrumentation = inst;
+    }
+
+    public static Instrumentation instrumentation() {
+        return instrumentation;
+    }
+
+    public static boolean canRedefine() {
+        Instrumentation inst = instrumentation;
+        return inst != null && inst.isRedefineClassesSupported();
+    }
+}

jjava-kernel/src/main/java/org/dflib/jjava/kernel/execution/JJavaLoaderDelegate.java

@@ -4,11 +4,16 @@
 import jdk.jshell.spi.ExecutionControl;
 import org.dflib.jjava.jupyter.kernel.util.PathsHandler;
 
+import java.lang.instrument.ClassDefinition;
+import java.lang.instrument.Instrumentation;
+import java.lang.instrument.UnmodifiableClassException;
 import java.net.MalformedURLException;
 import java.net.URL;
 import java.net.URLClassLoader;
 import java.nio.file.Path;
 import java.security.CodeSource;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -53,8 +58,41 @@ public void load(ExecutionControl.ClassBytecodes[] cbcs) throws ExecutionControl
 
     @Override
     public void classesRedefined(ExecutionControl.ClassBytecodes[] cbcs) {
+        List<ClassDefinition> redefinitions = new ArrayList<>(cbcs.length);
         for (ExecutionControl.ClassBytecodes cbc : cbcs) {
             declaredClasses.put(cbc.name(), cbc.bytecodes());
+
+            Class<?> loaded = loadedClasses.get(cbc.name());
+            if (loaded != null) {
+                redefinitions.add(new ClassDefinition(loaded, cbc.bytecodes()));
+            }
+        }
+        redefineLoadedClasses(redefinitions);
+    }
+
+   /**
+    * Perform in-place (HotSwap) redefinition of the given classes. This mirrors how the {@code jshell} tool
+    * redefines classes in its remote JVM via JDI, preserving {@link Class} identity and the values of
+    * static fields (i.e. JShell variable values).
+    */
+    private void redefineLoadedClasses(List<ClassDefinition> redefinitions) {
+        if (redefinitions.isEmpty() || !HotswapAgent.canRedefine()) {
+            return;
+        }
+
+        Instrumentation instrumentation = HotswapAgent.instrumentation();
+        try {
+            // Redefine all classes at once to handle interdependent changes to more than one class quickly.
+            instrumentation.redefineClasses(redefinitions.toArray(new ClassDefinition[0]));
+        } catch (UnsupportedOperationException | UnmodifiableClassException | ClassNotFoundException e) {
+            // Schema-changing redefinition in some class: try to redefine each class separately.
+            for (ClassDefinition redefinition : redefinitions) {
+                try {
+                    instrumentation.redefineClasses(redefinition);
+                } catch (UnsupportedOperationException | UnmodifiableClassException | ClassNotFoundException ee) {
+                    // Schema-changing redefinition: keep the new bytecode to declaredClasses.
+                }
+            }
         }
     }
 
@@ -120,7 +158,9 @@ protected Class<?> findClass(String name) throws ClassNotFoundException {
                 return super.findClass(name);
             }
             try {
-                return super.defineClass(name, data, 0, data.length, (CodeSource) null);
+                Class<?> klass = super.defineClass(name, data, 0, data.length, (CodeSource) null);
+                loadedClasses.put(name, klass);
+                return klass;
             } catch (LinkageError er) {
                 // rethrow as ClassNotFoundException to let the caller properly handle this case
                 // this error could be thrown in some cases (like static method signature change)