Skip to content

Commit

Permalink
Allow executing commands on BTrace server (#559)
Browse files Browse the repository at this point in the history
  • Loading branch information
@jbachorik
jbachorik committed Jul 9, 2022
1 parent e5e93b0 commit 6c059c3
Show file tree
Hide file tree
Showing 9 changed files with 552 additions and 101 deletions.
90 changes: 33 additions & 57 deletions btrace-agent/src/main/java/org/openjdk/btrace/agent/Client.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,14 +52,7 @@
import org.openjdk.btrace.core.BTraceRuntime;
import org.openjdk.btrace.core.DebugSupport;
import org.openjdk.btrace.core.SharedSettings;
import org.openjdk.btrace.core.comm.CommandListener;
import org.openjdk.btrace.core.comm.ErrorCommand;
import org.openjdk.btrace.core.comm.ExitCommand;
import org.openjdk.btrace.core.comm.InstrumentCommand;
import org.openjdk.btrace.core.comm.MessageCommand;
import org.openjdk.btrace.core.comm.RenameCommand;
import org.openjdk.btrace.core.comm.RetransformationStartNotification;
import org.openjdk.btrace.core.comm.StatusCommand;
import org.openjdk.btrace.core.comm.*;
import org.openjdk.btrace.instr.BTraceProbe;
import org.openjdk.btrace.instr.BTraceProbeFactory;
import org.openjdk.btrace.instr.BTraceProbePersisted;
Expand Down Expand Up @@ -338,13 +331,6 @@ final Class<?> loadClass(InstrumentCommand instr, boolean canLoadPack) throws IO
errorExit(th);
return null;
}

if (probe.isClassRenamed()) {
if (isDebug()) {
debugPrint("class renamed to " + probe.getClassName());
}
onCommand(new RenameCommand(probe.getClassName()));
}
if (isDebug()) {
debugPrint("creating BTraceRuntime instance for " + probe.getClassName());
}
Expand All @@ -357,12 +343,18 @@ final Class<?> loadClass(InstrumentCommand instr, boolean canLoadPack) throws IO
runtime.handleExit(0);
}
}));
if (probe.isClassRenamed()) {
if (isDebug()) {
debugPrint("class renamed to " + probe.getClassName());
}
sendCommand(new RenameCommand(probe.getClassName()));
}
if (isDebug()) {
debugPrint("created BTraceRuntime instance for " + probe.getClassName());
debugPrint("sending Okay command");
}

onCommand(new StatusCommand());
sendCommand(new StatusCommand());

boolean entered = false;
try {
Expand Down Expand Up @@ -391,9 +383,9 @@ protected void closeAll() throws IOException {

private void errorExit(Throwable th) throws IOException {
debugPrint("sending error command");
onCommand(new ErrorCommand(th));
sendCommand(new ErrorCommand(th));
debugPrint("sending exit command");
onCommand(new ExitCommand(1));
sendCommand(new ExitCommand(1));
closeAll();
}

Expand Down Expand Up @@ -445,23 +437,15 @@ private final boolean isCandidate(Class<?> c) {
}

private final void startRetransformClasses(int numClasses) {
try {
onCommand(new RetransformationStartNotification(numClasses));
if (isDebug()) {
debugPrint("calling retransformClasses (" + numClasses + " classes to be retransformed)");
}
} catch (IOException e) {
debugPrint(e);
sendCommand(new RetransformationStartNotification(numClasses));
if (isDebug()) {
debugPrint("calling retransformClasses (" + numClasses + " classes to be retransformed)");
}
}

final void endRetransformClasses() {
try {
onCommand(new StatusCommand());
if (isDebug()) debugPrint("finished retransformClasses");
} catch (IOException e) {
debugPrint(e);
}
sendCommand(new StatusCommand());
if (isDebug()) debugPrint("finished retransformClasses");
}

// Internals only below this point
Expand Down Expand Up @@ -514,19 +498,13 @@ boolean retransformLoaded() throws UnmodifiableClassException {
} catch (ClassFormatError | VerifyError e) {
debugPrint("Class '" + c.getName() + "' verification failed");
debugPrint(e);
try {
onCommand(
new MessageCommand(
"[BTRACE WARN] Class verification failed: "
+ c.getName()
+ " ("
+ e.getMessage()
+ ")"));
} catch (IOException ioException) {
if (isDebug()) {
debug.debug(ioException);
}
}
sendCommand(
new MessageCommand(
"[BTRACE WARN] Class verification failed: "
+ c.getName()
+ " ("
+ e.getMessage()
+ ")"));
}
}
} else {
Expand All @@ -544,19 +522,13 @@ boolean retransformLoaded() throws UnmodifiableClassException {
} catch (ClassFormatError | VerifyError e1) {
debugPrint("Class '" + c.getName() + "' verification failed");
debugPrint(e1);
try {
onCommand(
new MessageCommand(
"[BTRACE WARN] Class verification failed: "
+ c.getName()
+ " ("
+ e.getMessage()
+ ")"));
} catch (IOException ioException) {
if (isDebug()) {
debug.debug(ioException);
}
}
sendCommand(
new MessageCommand(
"[BTRACE WARN] Class verification failed: "
+ c.getName()
+ " ("
+ e.getMessage()
+ ")"));
}
}
}
Expand All @@ -566,6 +538,10 @@ boolean retransformLoaded() throws UnmodifiableClassException {
return true;
}

protected void sendCommand(Command command) {
runtime.send(command);
}

static Client findClient(String uuid) {
try {
UUID id = UUID.fromString(uuid);
Expand Down
81 changes: 50 additions & 31 deletions btrace-agent/src/main/java/org/openjdk/btrace/agent/RemoteClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
import java.net.SocketException;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
import java.util.concurrent.locks.LockSupport;
import org.openjdk.btrace.core.*;
import org.openjdk.btrace.core.comm.Command;
Expand Down Expand Up @@ -70,6 +71,13 @@ public Boolean apply(Command value) {
private volatile ObjectInputStream ois;
private volatile ObjectOutputStream oos;

private final AtomicReferenceFieldUpdater<RemoteClient, Socket> sockUpdater =
AtomicReferenceFieldUpdater.newUpdater(RemoteClient.class, Socket.class, "sock");
private final AtomicReferenceFieldUpdater<RemoteClient, ObjectInputStream> oisUpdater =
AtomicReferenceFieldUpdater.newUpdater(RemoteClient.class, ObjectInputStream.class, "ois");
private final AtomicReferenceFieldUpdater<RemoteClient, ObjectOutputStream> oosUpdater =
AtomicReferenceFieldUpdater.newUpdater(RemoteClient.class, ObjectOutputStream.class, "oos");

private final CircularBuffer<Command> delayedCommands = new CircularBuffer<>(5000);

static Client getClient(ClientContext ctx, Socket sock, Function<Client, Future<?>> initCallback)
Expand All @@ -93,7 +101,7 @@ static Client getClient(ClientContext ctx, Socket sock, Function<Client, Future<
try {
Client client = new RemoteClient(ctx, ois, oos, sock, (InstrumentCommand) cmd);
initCallback.apply(client).get();
WireIO.write(oos, new StatusCommand(InstrumentCommand.STATUS_FLAG));
client.sendCommand(new StatusCommand(InstrumentCommand.STATUS_FLAG));
return client;
} catch (ExecutionException | InterruptedException e) {
WireIO.write(oos, new StatusCommand(-1 * InstrumentCommand.STATUS_FLAG));
Expand All @@ -102,10 +110,11 @@ static Client getClient(ClientContext ctx, Socket sock, Function<Client, Future<
}
case Command.RECONNECT:
{
Client client = Client.findClient(((ReconnectCommand) cmd).getProbeId());
String probeId = ((ReconnectCommand) cmd).getProbeId();
Client client = Client.findClient(probeId);
if (client instanceof RemoteClient) {
((RemoteClient) client).reconnect(ois, oos, sock);
WireIO.write(oos, new StatusCommand(ReconnectCommand.STATUS_FLAG));
client.sendCommand(new StatusCommand(ReconnectCommand.STATUS_FLAG));
return client;
}
WireIO.write(oos, new StatusCommand(-1 * ReconnectCommand.STATUS_FLAG));
Expand Down Expand Up @@ -212,7 +221,8 @@ private void initClient() {
@SuppressWarnings("RedundantThrows")
@Override
public void onCommand(Command cmd) throws IOException {
if (oos == null) {
ObjectOutputStream output = oos;
if (output == null) {
if (!cmd.isUrgent()) {
delayedCommands.add(cmd);
}
Expand All @@ -224,9 +234,7 @@ public void onCommand(Command cmd) throws IOException {
try {
boolean isConnected = true;
try {
synchronized (oos) {
oos.reset();
}
output.reset();
} catch (SocketException e) {
isConnected = false;
}
Expand All @@ -248,35 +256,45 @@ private boolean dispatchCommand(Command cmd, boolean isConnected) {
if (cmd == Command.NULL) {
return true; // do not dispatch the NULL command
}
ObjectOutputStream output = oos;
ObjectInputStream input = ois;
Socket socket = sock;
if (output == null) {
return false;
}
try {
switch (cmd.getType()) {
case Command.EXIT:
if (isConnected) {
WireIO.write(oos, cmd);
WireIO.write(output, cmd);
}
onExit(((ExitCommand) cmd).getExitCode());
break;
case Command.LIST_PROBES:
{
if (isConnected) {
((ListProbesCommand) cmd).setProbes(listProbes());
WireIO.write(oos, cmd);
WireIO.write(output, cmd);
}
break;
}
case Command.DISCONNECT:
{
((DisconnectCommand) cmd).setProbeId(id.toString());
synchronized (oos) {
WireIO.write(oos, cmd);
oos.flush();
ois.close();
oos.close();
if (output != null) {
WireIO.write(output, cmd);
output.flush();
output.close();
oosUpdater.compareAndSet(this, output, null);
}
if (input != null) {
input.close();
oisUpdater.compareAndSet(this, input, null);
}
if (socket != null) {
socket.close();
sockUpdater.compareAndSet(this, socket, null);
}
sock.close();
ois = null;
oos = null;
sock = null;
break;
}
default:
Expand All @@ -301,27 +319,28 @@ public boolean isDisconnected() {
}

@Override
protected synchronized void closeAll() throws IOException {
protected void closeAll() throws IOException {
super.closeAll();

if (oos != null) {
synchronized (oos) {
oos.close();
}
oos = null;
ObjectOutputStream output = oos;
if (output != null) {
output.close();
oosUpdater.compareAndSet(this, output, null);
}
if (ois != null) {
ois.close();
ois = null;
ObjectInputStream input = ois;
if (input != null) {
input.close();
oisUpdater.compareAndSet(this, input, null);
}
if (sock != null) {
sock.close();
sock = null;
Socket socket = sock;
if (socket != null) {
socket.close();
sockUpdater.compareAndSet(this, socket, null);
}
}

void reconnect(ObjectInputStream ois, ObjectOutputStream oos, Socket socket) throws IOException {
sock = socket;
this.sock = socket;
this.ois = ois;
this.oos = oos;
onCommand(Command.NULL);
Expand Down
33 changes: 32 additions & 1 deletion btrace-client/src/main/java/org/openjdk/btrace/client/Client.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -419,7 +419,8 @@ void connectAndListProbes(String host, CommandListener listener) throws IOExcept
}
}

void reconnect(String host, String resumeProbe, CommandListener listener) throws IOException {
void reconnect(String host, String resumeProbe, CommandListener listener, String[] command)
throws IOException {
if (sock != null) {
throw new IllegalStateException();
}
Expand Down Expand Up @@ -467,6 +468,36 @@ public void onCommand(Command cmd) throws IOException {
if (statusCommand.getFlag() == ReconnectCommand.STATUS_FLAG) {
if (statusCommand.isSuccess()) {
DebugSupport.info("Reconnected to an active probe: " + resumeProbe);
String probeCommand = command[0];
String probeCommandArg = command[1];
if (probeCommand != null) {
debugPrint(
"Executing remote command '"
+ probeCommand
+ "'"
+ (probeCommandArg != null ? "(" + probeCommandArg + ")" : ""));
switch (probeCommand) {
case "exit":
{
sendExit();
break;
}
case "event":
{
if (probeCommandArg == null || probeCommandArg.equals("*")) {
sendEvent();
} else {
sendEvent(probeCommandArg);
}
sendDisconnect();
break;
}
default:
{
DebugSupport.warning("Unrecognized BTrace command " + probeCommand);
}
}
}
} else {
DebugSupport.warning("Unable to reconnect to an active probe: " + resumeProbe);
System.exit(1);
Expand Down
Loading

0 comments on commit 6c059c3

Please sign in to comment.