diff --git a/btrace-agent/src/main/java/io/btrace/agent/Client.java b/btrace-agent/src/main/java/io/btrace/agent/Client.java index 27f4668cc..aa123806e 100644 --- a/btrace-agent/src/main/java/io/btrace/agent/Client.java +++ b/btrace-agent/src/main/java/io/btrace/agent/Client.java @@ -288,6 +288,11 @@ static Collection listProbes() { return probes; } + /** Returns {@code true} once {@link #onExit(int)} has started removing this probe. */ + final boolean isShuttingDown() { + return shuttingDown; + } + synchronized void onExit(int exitCode) { if (!shuttingDown) { shuttingDown = true; diff --git a/btrace-agent/src/main/java/io/btrace/agent/RemoteClient.java b/btrace-agent/src/main/java/io/btrace/agent/RemoteClient.java index c4d56c057..78cf8a259 100644 --- a/btrace-agent/src/main/java/io/btrace/agent/RemoteClient.java +++ b/btrace-agent/src/main/java/io/btrace/agent/RemoteClient.java @@ -103,6 +103,8 @@ public Boolean apply(Command value) { AtomicReferenceFieldUpdater.newUpdater(RemoteClient.class, Socket.class, "sock"); private final AtomicReferenceFieldUpdater protocolUpdater = AtomicReferenceFieldUpdater.newUpdater(RemoteClient.class, WireProtocol.class, "protocol"); + // Keeps sock and protocol a consistent pair across reconnect() and releaseTransport(). + private final Object transportLock = new Object(); private final CircularBuffer delayedCommands = new CircularBuffer<>(5000); private final AtomicReference terminalHandshake = new AtomicReference<>(); @@ -279,13 +281,20 @@ private void initClient() { () -> { try { BTraceRuntime.enter(); - while (true) { + while (!isShuttingDown()) { + WireProtocol input; + Socket inputSocket; + synchronized (transportLock) { + input = protocol; + inputSocket = sock; + } try { - if (protocol == null) { + if (input == null) { + // Detached: wait for reconnect() to install a new transport. LockSupport.parkNanos(500_000_000L); // sleep 500ms continue; } - Command cmd = protocol.read(); + Command cmd = input.read(); switch (cmd.getType()) { case Command.EXIT: { @@ -341,6 +350,12 @@ private void initClient() { if (log.isDebugEnabled()) { log.debug("client command stream closed: {}", exp.toString()); } + if (terminalHandshake.get() == null) { + // The probe outlives its client: keep it reconnectable and keep reading + // once a client reconnects. + releaseTransport(input, inputSocket); + continue; + } } else { log.debug("Error while processing BTrace command", exp); } @@ -548,10 +563,38 @@ protected void closeAll() throws IOException { } } + /** + * Drops a client transport whose stream has ended, leaving the probe running and reconnectable. A + * transport that reconnect() has already replaced is closed without touching the new one. + */ + private void releaseTransport(WireProtocol input, Socket inputSocket) { + synchronized (transportLock) { + if (protocol == input) { + protocol = null; + sock = null; + disconnected = true; + } + } + try { + input.close(); + } catch (IOException ignore) { + // best effort + } + if (inputSocket != null) { + try { + inputSocket.close(); + } catch (IOException ignore) { + // best effort + } + } + } + void reconnect(WireProtocol protocol, Socket socket) throws IOException { - this.sock = socket; - this.protocol = protocol; - this.disconnected = false; + synchronized (transportLock) { + this.sock = socket; + this.protocol = protocol; + this.disconnected = false; + } onCommand(Command.NULL); } } diff --git a/btrace-agent/src/test/java/io/btrace/agent/RemoteClientReconnectTest.java b/btrace-agent/src/test/java/io/btrace/agent/RemoteClientReconnectTest.java new file mode 100644 index 000000000..91a929377 --- /dev/null +++ b/btrace-agent/src/test/java/io/btrace/agent/RemoteClientReconnectTest.java @@ -0,0 +1,127 @@ +/* + * Copyright (c) 2026, Jaroslav Bachorik . + * All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.btrace.agent; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import io.btrace.core.ArgsMap; +import io.btrace.core.BTraceRuntime; +import io.btrace.core.SharedSettings; +import io.btrace.core.comm.BinaryWireProtocol; +import io.btrace.core.comm.ExitCommand; +import io.btrace.core.comm.WireProtocol; +import java.lang.reflect.Proxy; +import java.net.InetAddress; +import java.net.ServerSocket; +import java.net.Socket; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; + +class RemoteClientReconnectTest { + private static final int EXIT_CODE = 41; + + @Test + void reconnectedClientCommandsReachTheProbeAfterThePeerWentAway() throws Exception { + CountDownLatch exitRequested = new CountDownLatch(1); + AtomicInteger requestedExitCode = new AtomicInteger(Integer.MIN_VALUE); + BTraceRuntime.Impl runtime = + runtimeProxy( + exitCode -> { + requestedExitCode.set(exitCode); + exitRequested.countDown(); + }); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (ServerSocket server = new ServerSocket(0, 1, InetAddress.getLoopbackAddress())) { + Future accepted = + executor.submit( + () -> { + Socket socket = server.accept(); + ClientContext context = + new ClientContext(null, null, new ArgsMap(), new SharedSettings()); + return RemoteClient.createForTerminalTest( + context, protocol(socket), socket, runtime, exitCode -> {}); + }); + RemoteClient remote; + try (Socket first = new Socket(server.getInetAddress(), server.getLocalPort())) { + remote = accepted.get(5, TimeUnit.SECONDS); + } + // The client side closed its socket, as a detaching client does. + awaitTrue(remote::isDisconnected, "agent did not notice the client going away"); + + Future reconnecting = executor.submit(server::accept); + try (Socket second = new Socket(server.getInetAddress(), server.getLocalPort())) { + Socket agentSide = reconnecting.get(5, TimeUnit.SECONDS); + remote.reconnect(protocol(agentSide), agentSide); + assertFalse(remote.isDisconnected()); + + WireProtocol client = protocol(second); + client.write(new ExitCommand(EXIT_CODE)); + client.flush(); + + assertTrue( + exitRequested.await(5, TimeUnit.SECONDS), + "EXIT from the reconnected client never reached the probe"); + assertEquals(EXIT_CODE, requestedExitCode.get()); + } + } finally { + executor.shutdownNow(); + } + } + + private static WireProtocol protocol(Socket socket) throws java.io.IOException { + return new BinaryWireProtocol(socket.getInputStream(), socket.getOutputStream()); + } + + private static void awaitTrue(java.util.function.BooleanSupplier condition, String message) + throws InterruptedException { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (!condition.getAsBoolean()) { + if (System.nanoTime() > deadline) { + throw new AssertionError(message); + } + Thread.sleep(20); + } + } + + private static BTraceRuntime.Impl runtimeProxy(final ExitRecorder recorder) { + return (BTraceRuntime.Impl) + Proxy.newProxyInstance( + RemoteClientReconnectTest.class.getClassLoader(), + new Class[] {BTraceRuntime.Impl.class}, + (proxy, method, args) -> { + if ("handleExit".equals(method.getName())) { + recorder.record(((Integer) args[0]).intValue()); + } + Class type = method.getReturnType(); + if (type == Boolean.TYPE) return Boolean.FALSE; + if (type == Integer.TYPE) return Integer.valueOf(0); + if (type == Long.TYPE) return Long.valueOf(0L); + return null; + }); + } + + private interface ExitRecorder { + void record(int exitCode); + } +} diff --git a/btrace-client/src/main/java/io/btrace/client/Client.java b/btrace-client/src/main/java/io/btrace/client/Client.java index 2cf8cd022..c6caa71be 100644 --- a/btrace-client/src/main/java/io/btrace/client/Client.java +++ b/btrace-client/src/main/java/io/btrace/client/Client.java @@ -301,9 +301,8 @@ public byte[] compile(String fileName, String classPath) { * Scans the extensions directory and returns a classpath string with all API JARs. * *

Looks in the following locations (first match wins): - System property 'btrace.libs' - * (assumed to point to BTrace libs directory); scans sibling 'extensions/' - Derived from - * java.class.path entry containing 'btrace-client' (for dist layouts); scans sibling - * 'extensions/' + * (assumed to point to BTrace libs directory); scans sibling 'extensions/' - Derived from a + * btrace-client JAR on java.class.path (for legacy dist layouts); scans sibling 'extensions/' */ private String getExtensionApiClasspath() { try { @@ -318,19 +317,12 @@ private String getExtensionApiClasspath() { } } - // 2) Fall back to locating btrace-client in the classpath - String classPath = System.getProperty("java.class.path"); - for (String entry : classPath.split(File.pathSeparator)) { - if (entry.contains("btrace-client")) { - File clientPath = new File(entry); - File libsDir = clientPath.getParentFile(); - if (libsDir != null) { - File btraceHome = libsDir.getParentFile(); - String cp = scanExtensionsDir(new File(btraceHome, "extensions")); - if (!cp.isEmpty()) { - return cp; - } - } + // 2) Fall back to locating the btrace-client JAR in the classpath + File btraceHome = btraceHomeFromClassPath(System.getProperty("java.class.path")); + if (btraceHome != null) { + String cp = scanExtensionsDir(new File(btraceHome, "extensions")); + if (!cp.isEmpty()) { + return cp; } } // 3) Fall back to extensions embedded in the agent JAR itself. A Maven or jbang user has no @@ -348,6 +340,31 @@ private String getExtensionApiClasspath() { return ""; } + /** + * Returns the BTrace home of a {@code /libs/btrace-client[-].jar} entry on {@code + * classPath}. Only the entry's file name is matched, so a directory that merely contains + * "btrace-client" in its path, such as a checkout or worktree, is never taken for the client. + * + * @return the home directory, or {@code null} when no such JAR is on the classpath + */ + static File btraceHomeFromClassPath(String classPath) { + if (classPath == null) { + return null; + } + for (String entry : classPath.split(File.pathSeparator)) { + File clientJar = new File(entry); + String name = clientJar.getName(); + boolean isClientJar = + name.endsWith(".jar") + && (name.equals("btrace-client.jar") || name.startsWith("btrace-client-")); + File libsDir = clientJar.getAbsoluteFile().getParentFile(); + if (isClientJar && libsDir != null && libsDir.getParentFile() != null) { + return libsDir.getParentFile(); + } + } + return null; + } + /** * Returns the agent JAR as a compile classpath entry when it declares embedded extensions. * @@ -1130,6 +1147,22 @@ void connectAndListFailedExtensions(String host, CommandListener listener) throw } } + /** + * Reconnects to a probe that a client detached from and stops it, as {@code -r exit} + * does. Returns once the agent confirms the exit; the connection is always closed. + * + * @param host the agent host + * @param probeId a probe id as reported by {@link #connectAndListProbes} + * @throws IOException if the agent is unavailable or does not know {@code probeId} + */ + public void connectAndExitProbe(String host, String probeId) throws IOException { + try { + reconnect(host, probeId, cmd -> {}, new String[] {"exit", null}); + } finally { + closeQuietly(); + } + } + void reconnect(String host, String resumeProbe, CommandListener listener, String[] command) throws IOException { if (sock != null) { diff --git a/btrace-client/src/test/java/io/btrace/client/ClientTest.java b/btrace-client/src/test/java/io/btrace/client/ClientTest.java index c564796cf..0aaf962f4 100644 --- a/btrace-client/src/test/java/io/btrace/client/ClientTest.java +++ b/btrace-client/src/test/java/io/btrace/client/ClientTest.java @@ -20,6 +20,7 @@ import com.sun.tools.attach.AgentLoadException; import io.btrace.core.comm.Command; +import io.btrace.core.comm.ExitCommand; import io.btrace.core.comm.JavaSerializationProtocol; import io.btrace.core.comm.ListFailedExtensionsCommand; import io.btrace.core.comm.ListProbesCommand; @@ -539,6 +540,90 @@ void reconnectThrowsForUnknownProbe() throws Exception { } } + @Test + void connectAndExitProbeStopsTheProbeAndClosesSocket() throws Exception { + forceV1Protocol(); + ExecutorService executor = Executors.newSingleThreadExecutor(); + try (ServerSocket server = new ServerSocket(0, 1, InetAddress.getLoopbackAddress())) { + // The agent side of a reconnect: acknowledge the probe, then echo the client's EXIT. + Future> received = + executor.submit( + () -> { + List commands = new ArrayList<>(); + try (Socket socket = server.accept()) { + socket.setSoTimeout(10_000); + JavaSerializationProtocol protocol = + new JavaSerializationProtocol( + socket.getInputStream(), socket.getOutputStream()); + commands.add(protocol.read()); + protocol.write(new StatusCommand(ReconnectCommand.STATUS_FLAG)); + protocol.flush(); + Command exit = protocol.read(); + commands.add(exit); + protocol.write(exit); + protocol.flush(); + try { + protocol.read(); + } catch (IOException e) { + commands.add(null); // EOF: the client closed the connection + } + } + return commands; + }); + Properties properties = new Properties(); + properties.setProperty("btrace.port", String.valueOf(server.getLocalPort())); + properties.setProperty("btrace.address", "127.0.0.1"); + Client client = new Client(0); + client.configureConnection(properties); + + client.connectAndExitProbe("localhost", "probe-1"); + + List commands = received.get(10, TimeUnit.SECONDS); + assertEquals(3, commands.size(), "socket must be closed after the EXIT echo"); + assertEquals("probe-1", ((ReconnectCommand) commands.get(0)).getProbeId()); + assertEquals(Command.EXIT, commands.get(1).getType()); + assertEquals(0, ((ExitCommand) commands.get(1)).getExitCode()); + } finally { + executor.shutdownNow(); + clearProtocolProperties(); + } + } + + @Test + void connectAndExitProbeThrowsForUnknownProbe() throws Exception { + forceV1Protocol(); + try (FakeAgent agent = new FakeAgent(1, new StatusCommand(-ReconnectCommand.STATUS_FLAG))) { + Client client = agent.client(); + + IOException failure = + assertThrows( + IOException.class, () -> client.connectAndExitProbe("localhost", "no-such-probe")); + + assertTrue(failure.getMessage().contains("no-such-probe")); + assertTrue(agent.clientClosed.get(10, TimeUnit.SECONDS), "socket must be closed"); + } finally { + clearProtocolProperties(); + } + } + + @Test + void btraceHomeIsDerivedOnlyFromABtraceClientJar() { + String sep = File.pathSeparator; + String home = new File("/opt/btrace").getAbsolutePath(); + // A directory merely named like the client -- a checkout or worktree -- is not the client. + assertNull( + Client.btraceHomeFromClassPath( + "/work/btrace-client-feature/build/classes/java/main" + + sep + + "/work/btrace-client-feature/lib/other.jar")); + assertNull(Client.btraceHomeFromClassPath("/opt/btrace/libs/btrace-clientx.jar")); + assertEquals( + new File(home), + Client.btraceHomeFromClassPath("/x/a.jar" + sep + home + "/libs/btrace-client.jar")); + assertEquals( + new File(home), Client.btraceHomeFromClassPath(home + "/libs/btrace-client-3.0.0.jar")); + } + private static Object readField(Client client, String name) throws Exception { Field field = Client.class.getDeclaredField(name); field.setAccessible(true); diff --git a/btrace-core/src/main/java/io/btrace/core/comm/v2/BinaryProtocol.java b/btrace-core/src/main/java/io/btrace/core/comm/v2/BinaryProtocol.java index 80f4bd9fe..b163a23c2 100644 --- a/btrace-core/src/main/java/io/btrace/core/comm/v2/BinaryProtocol.java +++ b/btrace-core/src/main/java/io/btrace/core/comm/v2/BinaryProtocol.java @@ -16,6 +16,7 @@ */ package io.btrace.core.comm.v2; +import java.io.EOFException; import java.io.IOException; import java.io.InputStream; import java.io.OutputStream; @@ -37,7 +38,7 @@ public class BinaryProtocol { public static byte readByte(InputStream is) throws IOException { int value = is.read(); if (value == -1) { - throw new IOException("End of stream reached"); + throw new EOFException("End of stream reached"); } return (byte) value; } @@ -181,7 +182,7 @@ public static void readFully(InputStream is, byte[] buffer, int offset, int leng while (totalBytesRead < length) { int bytesRead = is.read(buffer, offset + totalBytesRead, length - totalBytesRead); if (bytesRead == -1) { - throw new IOException("End of stream reached"); + throw new EOFException("End of stream reached"); } totalBytesRead += bytesRead; } diff --git a/integration-tests/src/test/java/tests/BTraceFunctionalTests.java b/integration-tests/src/test/java/tests/BTraceFunctionalTests.java index 920f6d5e6..7d853b114 100644 --- a/integration-tests/src/test/java/tests/BTraceFunctionalTests.java +++ b/integration-tests/src/test/java/tests/BTraceFunctionalTests.java @@ -47,6 +47,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.ValueSource; import tests.harness.Completion; @@ -533,10 +534,51 @@ public void validate(String stdout, String stderr, int retcode, String jfrFile) @Test public void testOnMethodUnattended() throws Exception { TestApp testApp = launchTestApp("resources.Main"); + try { + String host = "localhost"; + String probeId = deployAndDetachOnMethodProbe(testApp, host); + + List probes = listProbesWithProtocol(host); + long matches = + probes.stream().filter(p -> extractProbeClassName(p).endsWith("OnMethodTest")).count(); + assertEquals(1, matches, "expected exactly one OnMethodTest probe listed by -lp"); + assertTrue( + probes.stream().anyMatch(p -> p.startsWith(probeId + " ")), + "probe id not present in -lp output"); + } finally { + // The target's agent holds the BTrace port; a later test attaching a new target needs it. + testApp.stop(); + } + } + + @Test + @Timeout(value = 2, unit = TimeUnit.MINUTES) + public void testStopDetachedProbe() throws Exception { + TestApp testApp = launchTestApp("resources.Main"); + try { + String host = "localhost"; + String probeId = deployAndDetachOnMethodProbe(testApp, host); + assertTrue( + listProbesWithProtocol(host).stream().anyMatch(p -> p.startsWith(probeId + " ")), + "detached probe not listed"); + + Client stopper = createClientForTests(locateTrace("btrace/OnMethodTest.java").getParent()); + stopper.attach(String.valueOf(testApp.getPid()), null, getEventsClassPath()); + stopper.connectAndExitProbe(host, probeId); + + assertFalse( + listProbesWithProtocol(host).stream().anyMatch(p -> p.startsWith(probeId + " ")), + "stopped probe is still listed"); + } finally { + testApp.stop(); + } + } + + /** Deploys OnMethodTest into {@code testApp}, detaches from it and returns its probe id. */ + private String deployAndDetachOnMethodProbe(TestApp testApp, String host) throws Exception { File traceFile = locateTrace("btrace/OnMethodTest.java"); String pid = String.valueOf(testApp.getPid()); - String host = "localhost"; Client client = createClientForTests(traceFile.getParentFile().getAbsolutePath()); client.attach(pid, null, getEventsClassPath()); byte[] code = client.compile(traceFile.getAbsolutePath(), getEventsClassPath()); @@ -587,14 +629,7 @@ public void testOnMethodUnattended() throws Exception { } finally { executor.shutdownNow(); } - - List probes = listProbesWithProtocol(host); - long matches = - probes.stream().filter(p -> extractProbeClassName(p).endsWith("OnMethodTest")).count(); - assertEquals(1, matches, "expected exactly one OnMethodTest probe listed by -lp"); - assertTrue( - probes.stream().anyMatch(p -> p.startsWith(probeId + " ")), - "probe id not present in -lp output"); + return probeId; } private static String extractProbeClassName(String probeEntry) {