From 5f5b231aaeb8ea8baa2e461d982681e71933b98b Mon Sep 17 00:00:00 2001 From: Dongliang Xie Date: Sat, 20 Jun 2026 14:43:03 +0800 Subject: [PATCH] fix: complete methods for classes loaded by multiple ClassLoaders --- .../core/shell/cli/CompletionUtils.java | 38 ++-- .../core/shell/cli/CompletionUtilsTest.java | 213 ++++++++++++++++++ 2 files changed, 237 insertions(+), 14 deletions(-) create mode 100644 core/src/test/java/com/taobao/arthas/core/shell/cli/CompletionUtilsTest.java diff --git a/core/src/main/java/com/taobao/arthas/core/shell/cli/CompletionUtils.java b/core/src/main/java/com/taobao/arthas/core/shell/cli/CompletionUtils.java index 15d1f70cbdc..e04a4cd0da9 100644 --- a/core/src/main/java/com/taobao/arthas/core/shell/cli/CompletionUtils.java +++ b/core/src/main/java/com/taobao/arthas/core/shell/cli/CompletionUtils.java @@ -17,6 +17,7 @@ import java.util.Collection; import java.util.Collections; import java.util.HashSet; +import java.util.LinkedHashSet; import java.util.LinkedList; import java.util.List; import java.util.Set; @@ -220,28 +221,37 @@ public static boolean completeMethodName(Completion completion) { className = tokens.get(tokens.size() - 3).value(); } - Set> results = SearchUtils.searchClassOnly(completion.session().getInstrumentation(), className, 2); - if (results.size() != 1) { - // no class found or multiple class found + Set> results = SearchUtils.searchClassOnly(completion.session().getInstrumentation(), className, + false); + if (results.isEmpty()) { + // no class found completion.complete(Collections.emptyList()); return true; } - Class clazz = results.iterator().next(); - - List res = new ArrayList(); - - for (Method method : clazz.getDeclaredMethods()) { - if (StringUtils.isBlank(lastToken)) { - res.add(method.getName()); - } else if (method.getName().startsWith(lastToken)) { - res.add(method.getName()); + String matchedClassName = null; + Set res = new LinkedHashSet(); + for (Class clazz : results) { + if (matchedClassName == null) { + matchedClassName = clazz.getName(); + } else if (!matchedClassName.equals(clazz.getName())) { + completion.complete(Collections.emptyList()); + return true; } + for (Method method : clazz.getDeclaredMethods()) { + if (StringUtils.isBlank(lastToken)) { + res.add(method.getName()); + } else if (method.getName().startsWith(lastToken)) { + res.add(method.getName()); + } + } + } + if (StringUtils.isBlank(lastToken) || "".startsWith(lastToken)) { + res.add(""); } - res.add(""); if (res.size() == 1) { - completion.complete(res.get(0).substring(lastToken.length()), true); + completion.complete(res.iterator().next().substring(lastToken.length()), true); return true; } else { CompletionUtils.complete(completion, res); diff --git a/core/src/test/java/com/taobao/arthas/core/shell/cli/CompletionUtilsTest.java b/core/src/test/java/com/taobao/arthas/core/shell/cli/CompletionUtilsTest.java new file mode 100644 index 00000000000..928aaaa0ced --- /dev/null +++ b/core/src/test/java/com/taobao/arthas/core/shell/cli/CompletionUtilsTest.java @@ -0,0 +1,213 @@ +package com.taobao.arthas.core.shell.cli; + +import com.taobao.arthas.core.shell.cli.impl.CliTokenImpl; +import com.taobao.arthas.core.shell.session.Session; +import org.junit.Assert; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.mockito.Mockito; + +import javax.tools.JavaCompiler; +import javax.tools.ToolProvider; +import java.io.File; +import java.lang.instrument.Instrumentation; +import java.lang.reflect.Proxy; +import java.net.URL; +import java.net.URLClassLoader; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; + +public class CompletionUtilsTest { + + private static final String DUPLICATE_TARGET = "test.arthas.DuplicateTarget"; + + @Rule + public TemporaryFolder temporaryFolder = new TemporaryFolder(); + + @Test + public void shouldMergeMethodCandidatesFromClassesLoadedByDifferentClassLoaders() throws Exception { + Class first = compileDuplicateTarget("public void alpha() {}\npublic void common() {}"); + Class second = compileDuplicateTarget("public void beta() {}\npublic void common() {}"); + RecordingCompletion completion = completionFor(methodCompletionTokens(DUPLICATE_TARGET, ""), first, second); + + Assert.assertTrue(CompletionUtils.completeMethodName(completion)); + + Assert.assertNotNull(completion.candidates); + Assert.assertTrue(completion.candidates.contains("alpha")); + Assert.assertTrue(completion.candidates.contains("beta")); + Assert.assertTrue(completion.candidates.contains("common")); + Assert.assertTrue(completion.candidates.contains("")); + Assert.assertEquals(1, Collections.frequency(completion.candidates, "common")); + Assert.assertEquals(new HashSet(completion.candidates).size(), completion.candidates.size()); + } + + @Test + public void shouldCompletePartialMethodNameAcrossClassesLoadedByDifferentClassLoaders() throws Exception { + Class first = compileDuplicateTarget("public void alpha() {}\npublic void common() {}"); + Class second = compileDuplicateTarget("public void beta() {}\npublic void common() {}"); + RecordingCompletion completion = completionFor(methodCompletionTokens(DUPLICATE_TARGET, "al"), first, second); + + Assert.assertTrue(CompletionUtils.completeMethodName(completion)); + + Assert.assertEquals("pha", completion.value); + Assert.assertTrue(completion.terminal); + Assert.assertNull(completion.candidates); + } + + @Test + public void shouldCompleteConstructorPrefix() throws Exception { + Class first = compileDuplicateTarget("public void alpha() {}"); + Class second = compileDuplicateTarget("public void beta() {}"); + RecordingCompletion completion = completionFor(methodCompletionTokens(DUPLICATE_TARGET, "<"), first, second); + + Assert.assertTrue(CompletionUtils.completeMethodName(completion)); + + Assert.assertEquals("init>", completion.value); + Assert.assertTrue(completion.terminal); + Assert.assertNull(completion.candidates); + } + + @Test + public void shouldCompleteEmptyListWhenNoClassMatches() { + RecordingCompletion completion = completionFor(methodCompletionTokens(DUPLICATE_TARGET, "")); + + Assert.assertTrue(CompletionUtils.completeMethodName(completion)); + + Assert.assertEquals(Collections.emptyList(), completion.candidates); + Assert.assertNull(completion.value); + } + + @Test + public void shouldCompleteEmptyListWhenClassPatternMatchesDifferentClassNames() throws Exception { + Class first = compileTarget("test.arthas.FirstTarget", "public void alpha() {}"); + Class second = compileTarget("test.arthas.SecondTarget", "public void beta() {}"); + RecordingCompletion completion = completionFor(methodCompletionTokens("test.arthas.*Target", ""), first, second); + + Assert.assertTrue(CompletionUtils.completeMethodName(completion)); + + Assert.assertEquals(Collections.emptyList(), completion.candidates); + Assert.assertNull(completion.value); + } + + private Class compileDuplicateTarget(String methods) throws Exception { + return compileTarget(DUPLICATE_TARGET, methods); + } + + private Class compileTarget(String className, String methods) throws Exception { + JavaCompiler compiler = ToolProvider.getSystemJavaCompiler(); + Assert.assertNotNull("JDK compiler is required to compile test classes", compiler); + + File sourceRoot = temporaryFolder.newFolder(); + int packageEnd = className.lastIndexOf('.'); + String packageName = className.substring(0, packageEnd); + String simpleName = className.substring(packageEnd + 1); + File packageDir = new File(sourceRoot, packageName.replace('.', File.separatorChar)); + Assert.assertTrue(packageDir.mkdirs()); + File sourceFile = new File(packageDir, simpleName + ".java"); + String source = "package " + packageName + ";\npublic class " + simpleName + " {\n" + methods + "\n}\n"; + Files.write(sourceFile.toPath(), source.getBytes(StandardCharsets.UTF_8)); + + File outputRoot = temporaryFolder.newFolder(); + int exitCode = compiler.run(null, null, null, "-d", outputRoot.getAbsolutePath(), sourceFile.getAbsolutePath()); + Assert.assertEquals(0, exitCode); + + URLClassLoader classLoader = new URLClassLoader(new URL[] { outputRoot.toURI().toURL() }, null); + try { + return Class.forName(className, true, classLoader); + } finally { + classLoader.close(); + } + } + + private static RecordingCompletion completionFor(List tokens, Class... classes) { + Session session = Mockito.mock(Session.class); + Mockito.when(session.getInstrumentation()).thenReturn(instrumentationFor(classes)); + return new RecordingCompletion(session, tokens); + } + + private static List methodCompletionTokens(String className, String methodPrefix) { + if (methodPrefix.length() == 0) { + return Arrays.asList( + new CliTokenImpl(true, "watch"), + new CliTokenImpl(false, " "), + new CliTokenImpl(true, className), + new CliTokenImpl(false, " ")); + } + return Arrays.asList( + new CliTokenImpl(true, "watch"), + new CliTokenImpl(false, " "), + new CliTokenImpl(true, className), + new CliTokenImpl(false, " "), + new CliTokenImpl(true, methodPrefix)); + } + + private static Instrumentation instrumentationFor(final Class... classes) { + return (Instrumentation) Proxy.newProxyInstance( + CompletionUtilsTest.class.getClassLoader(), + new Class[] { Instrumentation.class }, + (proxy, method, args) -> { + if ("getAllLoadedClasses".equals(method.getName())) { + return classes; + } + Class returnType = method.getReturnType(); + if (returnType == boolean.class) { + return false; + } + if (returnType == long.class) { + return 0L; + } + if (returnType == int.class) { + return 0; + } + if (returnType == Class[].class) { + return new Class[0]; + } + return null; + }); + } + + private static class RecordingCompletion implements Completion { + + private final Session session; + private final List tokens; + private List candidates; + private String value; + private boolean terminal; + + private RecordingCompletion(Session session, List tokens) { + this.session = session; + this.tokens = tokens; + } + + @Override + public Session session() { + return session; + } + + @Override + public String rawLine() { + return ""; + } + + @Override + public List lineTokens() { + return tokens; + } + + @Override + public void complete(List candidates) { + this.candidates = candidates; + } + + @Override + public void complete(String value, boolean terminal) { + this.value = value; + this.terminal = terminal; + } + } +}