diff --git a/boot/src/main/java/com/taobao/arthas/boot/Bootstrap.java b/boot/src/main/java/com/taobao/arthas/boot/Bootstrap.java index 5f10d0340..a17c1ca6d 100644 --- a/boot/src/main/java/com/taobao/arthas/boot/Bootstrap.java +++ b/boot/src/main/java/com/taobao/arthas/boot/Bootstrap.java @@ -106,6 +106,9 @@ public class Bootstrap { private String command; private String batchFile; + private String tunnelServer; + private String agentId; + static { ARTHAS_LIB_DIR = new File( System.getProperty("user.home") + File.separator + ".arthas" + File.separator + "lib"); @@ -230,6 +233,18 @@ public class Bootstrap { this.verbose = verbose; } + @Option(longName = "tunnel-server") + @Description("The tunnel server url") + public void setTunnelServer(String tunnelServer) { + this.tunnelServer = tunnelServer; + } + + @Option(longName = "agent-id") + @Description("The agent id register to tunnel server") + public void setAgentId(String agentId) { + this.agentId = agentId; + } + public static void main(String[] args) throws ParserConfigurationException, SAXException, IOException, ClassNotFoundException, NoSuchMethodException, SecurityException, IllegalAccessException, IllegalArgumentException, InvocationTargetException { @@ -458,6 +473,15 @@ public class Bootstrap { attachArgs.add("" + bootstrap.getSessionTimeout()); } + if (bootstrap.getTunnelServer() != null) { + attachArgs.add("-tunnel-server"); + attachArgs.add(bootstrap.getTunnelServer()); + } + if (bootstrap.getAgentId() != null) { + attachArgs.add("-agent-id"); + attachArgs.add(bootstrap.getAgentId()); + } + AnsiLog.info("Try to attach process " + pid); AnsiLog.debug("Start arthas-core.jar args: " + attachArgs); ProcessUtils.startArthasCore(pid, attachArgs); @@ -634,4 +658,12 @@ public class Bootstrap { public Integer getWidth() { return width; } + + public String getTunnelServer() { + return tunnelServer; + } + + public String getAgentId() { + return agentId; + } } diff --git a/core/pom.xml b/core/pom.xml index 3ac4e0d77..0f608030f 100644 --- a/core/pom.xml +++ b/core/pom.xml @@ -67,6 +67,11 @@ arthas-memorycompiler ${project.version} + + com.taobao.arthas + arthas-tunnel-client + ${project.version} + org.ow2.asm asm diff --git a/core/src/main/java/com/taobao/arthas/core/Arthas.java b/core/src/main/java/com/taobao/arthas/core/Arthas.java index 049e84302..a739a20e2 100644 --- a/core/src/main/java/com/taobao/arthas/core/Arthas.java +++ b/core/src/main/java/com/taobao/arthas/core/Arthas.java @@ -39,8 +39,12 @@ public class Arthas { .setShortName("http-port").setDefaultValue(DEFAULT_HTTP_PORT); Option sessionTimeout = new TypedOption().setType(Integer.class) .setShortName("session-timeout").setDefaultValue("" + Configure.DEFAULT_SESSION_TIMEOUT_SECONDS); + + Option tunnelServer = new TypedOption().setType(String.class).setShortName("tunnel-server"); + Option agentId = new TypedOption().setType(String.class).setShortName("agent-id"); + CLI cli = CLIs.create("arthas").addOption(pid).addOption(core).addOption(agent).addOption(target) - .addOption(telnetPort).addOption(httpPort).addOption(sessionTimeout); + .addOption(telnetPort).addOption(httpPort).addOption(sessionTimeout).addOption(tunnelServer).addOption(agentId); CommandLine commandLine = cli.parse(Arrays.asList(args)); Configure configure = new Configure(); @@ -56,6 +60,9 @@ public class Arthas { configure.setIp((String) commandLine.getOptionValue("target-ip")); configure.setTelnetPort((Integer) commandLine.getOptionValue("telnet-port")); configure.setHttpPort((Integer) commandLine.getOptionValue("http-port")); + + configure.setTunnelServer((String) commandLine.getOptionValue("tunnel-server")); + configure.setAgentId((String) commandLine.getOptionValue("agent-id")); return configure; } diff --git a/core/src/main/java/com/taobao/arthas/core/config/Configure.java b/core/src/main/java/com/taobao/arthas/core/config/Configure.java index f0b41d037..8db9c9b7c 100644 --- a/core/src/main/java/com/taobao/arthas/core/config/Configure.java +++ b/core/src/main/java/com/taobao/arthas/core/config/Configure.java @@ -24,6 +24,9 @@ public class Configure { private String arthasCore; private String arthasAgent; + private String tunnelServer; + private String agentId; + /** * session timeout seconds */ @@ -85,6 +88,22 @@ public class Configure { this.sessionTimeout = sessionTimeout; } + public String getTunnelServer() { + return tunnelServer; + } + + public void setTunnelServer(String tunnelServer) { + this.tunnelServer = tunnelServer; + } + + public String getAgentId() { + return agentId; + } + + public void setAgentId(String agentId) { + this.agentId = agentId; + } + // 对象的编码解码器 private final static FeatureCodec codec = new FeatureCodec(';', '='); @@ -106,7 +125,10 @@ public class Configure { // 非静态的才需要纳入非序列化过程 try { - map.put(field.getName(), String.valueOf(ArthasReflectUtils.getFieldValueByField(this, field))); + Object fieldValue = ArthasReflectUtils.getFieldValueByField(this, field); + if (fieldValue != null) { + map.put(field.getName(), String.valueOf(fieldValue)); + } } catch (Throwable t) { // } diff --git a/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java b/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java index da122ca3b..3c6270755 100644 --- a/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java +++ b/core/src/main/java/com/taobao/arthas/core/server/ArthasBootstrap.java @@ -1,6 +1,7 @@ package com.taobao.arthas.core.server; import com.taobao.arthas.core.config.Configure; +import com.alibaba.arthas.tunnel.client.TunnelClient; import com.taobao.arthas.core.command.BuiltinCommandPack; import com.taobao.arthas.core.shell.ShellServer; import com.taobao.arthas.core.shell.ShellServerOptions; @@ -9,19 +10,26 @@ import com.taobao.arthas.core.shell.handlers.BindHandler; import com.taobao.arthas.core.shell.impl.ShellServerImpl; import com.taobao.arthas.core.shell.term.impl.HttpTermServer; import com.taobao.arthas.core.shell.term.impl.TelnetTermServer; +import com.taobao.arthas.core.util.ArthasBanner; import com.taobao.arthas.core.util.Constants; import com.taobao.arthas.core.util.LogUtil; import com.taobao.arthas.core.util.UserStatUtil; import com.taobao.middleware.logger.Logger; +import io.netty.channel.ChannelFuture; + import java.io.IOException; import java.lang.instrument.Instrumentation; import java.lang.reflect.Method; +import java.net.URI; import java.util.ArrayList; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.ThreadFactory; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; @@ -39,6 +47,7 @@ public class ArthasBootstrap { private Thread shutdown; private ShellServer shellServer; private ExecutorService executorService; + private TunnelClient tunnelClient; private ArthasBootstrap(int pid, Instrumentation instrumentation) { this.pid = pid; @@ -78,11 +87,40 @@ public class ArthasBootstrap { throw new IllegalStateException("already bind"); } + String agentId = null; + try { + if (configure.getHttpPort() > 0) { + tunnelClient = new TunnelClient(); + tunnelClient.setId(configure.getAgentId()); + tunnelClient.setTunnelServerUrl(configure.getTunnelServer()); + // ws://127.0.0.1:8563/ws + String host = "127.0.0.1"; + if(configure.getIp() != null) { + host = configure.getIp(); + } + URI uri = new URI("ws", null, host, configure.getHttpPort(), "/ws", null, null); + tunnelClient.setLocalServerUrl(uri.toString()); + ChannelFuture channelFuture = tunnelClient.start(); + channelFuture.await(10, TimeUnit.SECONDS); + if(channelFuture.isSuccess()) { + agentId = tunnelClient.getId(); + } + } + } catch (Throwable t) { + logger.error("arthas", "start tunnel client error", t); + } + try { ShellServerOptions options = new ShellServerOptions() .setInstrumentation(instrumentation) .setPid(pid) .setSessionTimeout(configure.getSessionTimeout() * 1000); + + if (agentId != null) { + Map welcomeInfos = new HashMap(); + welcomeInfos.put("id", agentId); + options.setWelcomeMessage(ArthasBanner.welcome(welcomeInfos)); + } shellServer = new ShellServerImpl(options, this); BuiltinCommandPack builtinCommands = new BuiltinCommandPack(); List resolvers = new ArrayList(); @@ -132,6 +170,13 @@ public class ArthasBootstrap { } public void destroy() { + if (this.tunnelClient != null) { + try { + tunnelClient.stop(); + } catch (Throwable e) { + logger.error("arthas", "stop tunnel client error", e); + } + } executorService.shutdownNow(); UserStatUtil.destroy(); // clear the reference in Spy class. diff --git a/core/src/main/java/com/taobao/arthas/core/util/ArthasBanner.java b/core/src/main/java/com/taobao/arthas/core/util/ArthasBanner.java index 4205315e3..a906c9116 100644 --- a/core/src/main/java/com/taobao/arthas/core/util/ArthasBanner.java +++ b/core/src/main/java/com/taobao/arthas/core/util/ArthasBanner.java @@ -9,6 +9,9 @@ import com.taobao.text.ui.TableElement; import com.taobao.text.util.RenderUtil; import java.io.InputStream; +import java.util.Collections; +import java.util.Map; +import java.util.Map.Entry; import static com.taobao.text.ui.Element.label; @@ -93,6 +96,10 @@ public class ArthasBanner { } public static String welcome() { + return welcome(Collections.emptyMap()); + } + + public static String welcome(Map infos) { logger.info("arthas version: " + version()); TableElement table = new TableElement().rightCellPadding(1) .row("wiki", wiki()) @@ -100,6 +107,9 @@ public class ArthasBanner { .row("version", version()) .row("pid", PidUtils.currentPid()) .row("time", DateUtils.getCurrentDate()); + for (Entry entry : infos.entrySet()) { + table.row(entry.getKey(), entry.getValue()); + } return logo() + "\n" + RenderUtil.render(table); } diff --git a/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClient.java b/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClient.java index 2835d946c..00ee63292 100644 --- a/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClient.java +++ b/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClient.java @@ -52,11 +52,11 @@ public class TunnelClient { // agent id, generated by tunnel server. if reconnect, reuse the id volatile private String id; - public void start() throws IOException, InterruptedException, URISyntaxException { - connect(false); + public ChannelFuture start() throws IOException, InterruptedException, URISyntaxException { + return connect(false); } - public void connect(boolean reconnect) throws SSLException, URISyntaxException, InterruptedException { + public ChannelFuture connect(boolean reconnect) throws SSLException, URISyntaxException, InterruptedException { QueryStringEncoder queryEncoder = new QueryStringEncoder(this.tunnelServerUrl); queryEncoder.addParam("method", "agentRegister"); if (id != null) { @@ -94,7 +94,13 @@ public class TunnelClient { sslCtx = null; } + WebSocketClientHandshaker newHandshaker = WebSocketClientHandshakerFactory.newHandshaker(agentRegisterURI, + WebSocketVersion.V13, null, true, new DefaultHttpHeaders()); + final WebSocketClientProtocolHandler websocketClientHandler = new WebSocketClientProtocolHandler(newHandshaker); + final TunnelClientSocketClientHandler handler = new TunnelClientSocketClientHandler(TunnelClient.this); + Bootstrap bs = new Bootstrap(); + bs.group(eventLoopGroup).channel(NioSocketChannel.class).remoteAddress(agentRegisterURI.getHost(), port) .handler(new ChannelInitializer() { @Override @@ -103,12 +109,7 @@ public class TunnelClient { if (sslCtx != null) { p.addLast(sslCtx.newHandler(ch.alloc(), host, port)); } - WebSocketClientHandshaker newHandshaker = WebSocketClientHandshakerFactory.newHandshaker( - agentRegisterURI, WebSocketVersion.V13, null, true, new DefaultHttpHeaders()); - WebSocketClientProtocolHandler websocketClientHandler = new WebSocketClientProtocolHandler( - newHandshaker); - final TunnelClientSocketClientHandler handler = new TunnelClientSocketClientHandler( - TunnelClient.this); + p.addLast(new HttpClientCodec(), new HttpObjectAggregator(8192), websocketClientHandler, handler); } @@ -126,6 +127,8 @@ public class TunnelClient { }); } channel = connectFuture.sync().channel(); + + return handler.registerFuture(); } public void stop() { diff --git a/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClientSocketClientHandler.java b/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClientSocketClientHandler.java index 75bc8dda5..e02cdcfd5 100644 --- a/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClientSocketClientHandler.java +++ b/tunnel-client/src/main/java/com/alibaba/arthas/tunnel/client/TunnelClientSocketClientHandler.java @@ -9,7 +9,9 @@ import java.util.concurrent.TimeUnit; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import io.netty.channel.ChannelFuture; import io.netty.channel.ChannelHandlerContext; +import io.netty.channel.ChannelPromise; import io.netty.channel.SimpleChannelInboundHandler; import io.netty.handler.codec.http.QueryStringDecoder; import io.netty.handler.codec.http.QueryStringEncoder; @@ -25,11 +27,21 @@ public class TunnelClientSocketClientHandler extends SimpleChannelInboundHandler private final static Logger logger = LoggerFactory.getLogger(TunnelClientSocketClientHandler.class); private TunnelClient tunnelClient; + private ChannelPromise registerPromise; public TunnelClientSocketClientHandler(TunnelClient tunnelClient) { this.tunnelClient = tunnelClient; } + public ChannelFuture registerFuture() { + return registerPromise; + } + + @Override + public void handlerAdded(ChannelHandlerContext ctx) { + registerPromise = ctx.newPromise(); + } + @Override public void channelRead0(ChannelHandlerContext ctx, WebSocketFrame frame) throws Exception { if (frame instanceof TextWebSocketFrame) { @@ -51,6 +63,7 @@ public class TunnelClientSocketClientHandler extends SimpleChannelInboundHandler if (idList != null && !idList.isEmpty()) { this.tunnelClient.setId(idList.get(0)); } + registerPromise.setSuccess(); } if ("startTunnel".equals(method)) { @@ -84,4 +97,12 @@ public class TunnelClientSocketClientHandler extends SimpleChannelInboundHandler } }, tunnelClient.getReconnectDelay(), TimeUnit.SECONDS); } + + @Override + public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) { + if (!registerPromise.isDone()) { + registerPromise.setFailure(cause); + } + ctx.fireExceptionCaught(cause); + } }