Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow executing commands on BTrace server #559

Merged
merged 1 commit into from
Jul 9, 2022
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
90 changes: 33 additions & 57 deletions btrace-agent/src/main/java/org/openjdk/btrace/agent/Client.java
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
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
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