diff --git a/apm-application-toolkit/apm-toolkit-opentracing/src/main/java/org/skywalking/apm/toolkit/opentracing/SkywalkingSpan.java b/apm-application-toolkit/apm-toolkit-opentracing/src/main/java/org/skywalking/apm/toolkit/opentracing/SkywalkingSpan.java
index 4b82d6021ccc81596b2f3a636415fde27feac975..ba982ec81c7e0fbf51d24d672c7db1f9be44db9d 100644
--- a/apm-application-toolkit/apm-toolkit-opentracing/src/main/java/org/skywalking/apm/toolkit/opentracing/SkywalkingSpan.java
+++ b/apm-application-toolkit/apm-toolkit-opentracing/src/main/java/org/skywalking/apm/toolkit/opentracing/SkywalkingSpan.java
@@ -17,12 +17,13 @@ public class SkywalkingSpan implements Span {
/**
* Create a shell span for {@link SkywalkingTracer#activeSpan()}
+ *
* @param tracer
*/
@NeedSnifferActivation(
"1. set the span reference to the dynamic field of enhanced SkywalkingSpan"
)
- public SkywalkingSpan(SkywalkingTracer tracer){
+ public SkywalkingSpan(SkywalkingTracer tracer) {
}
diff --git a/apm-sniffer/apm-agent-core/pom.xml b/apm-sniffer/apm-agent-core/pom.xml
index f94d04f9668231161327d8b330d01b48de1294cb..b0316d1272f3a8f55c84593d54c400f121bde89c 100644
--- a/apm-sniffer/apm-agent-core/pom.xml
+++ b/apm-sniffer/apm-agent-core/pom.xml
@@ -77,6 +77,24 @@
${jetty.version}
test
+
+ com.github.tomakehurst
+ wiremock
+ 2.6.0
+ test
+
+
+ io.grpc
+ grpc-testing
+ 1.4.0
+
+
+ mockito-core
+ org.mockito
+
+
+ test
+
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/conf/Config.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/conf/Config.java
index df95a68ccac3448ef31529f5231bd4c1d1c0cce5..5e25a95f2285d3df62ee333c890f57d7ff3b4139 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/conf/Config.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/conf/Config.java
@@ -34,6 +34,7 @@ public class Config {
}
public static class Collector {
+ public static long DISCOVERY_CHECK_INTERVAL = 60 * 1000;
/**
* Collector REST-Service address.
* e.g.
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/ContextManager.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/ContextManager.java
index d97490dc2de7322c5b6ff9501a80ed379e015b16..fbe727ab05489d84e9e0a49cd9d24894061e9eb9 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/ContextManager.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/ContextManager.java
@@ -5,7 +5,6 @@ import org.skywalking.apm.agent.core.boot.ServiceManager;
import org.skywalking.apm.agent.core.conf.Config;
import org.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
import org.skywalking.apm.agent.core.context.trace.AbstractSpan;
-import org.skywalking.apm.agent.core.context.trace.AbstractTracingSpan;
import org.skywalking.apm.agent.core.context.trace.TraceSegment;
import org.skywalking.apm.agent.core.dictionary.DictionaryUtil;
import org.skywalking.apm.agent.core.sampling.SamplingService;
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/trace/AbstractTracingSpan.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/trace/AbstractTracingSpan.java
index e1d1b80627cad534e011e7e8426bc9e0ec23adcf..28a32138140901d62340cdbdc490e52a94c675c5 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/trace/AbstractTracingSpan.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/trace/AbstractTracingSpan.java
@@ -214,14 +214,14 @@ public abstract class AbstractTracingSpan implements AbstractSpan {
spanBuilder.setParentSpanId(parentSpanId);
spanBuilder.setStartTime(startTime);
spanBuilder.setEndTime(endTime);
- if (operationId == DictionaryUtil.nullValue()) {
+ if (operationId != DictionaryUtil.nullValue()) {
spanBuilder.setOperationNameId(operationId);
} else {
spanBuilder.setOperationName(operationName);
}
spanBuilder.setSpanType(SpanType.Entry);
spanBuilder.setSpanLayerValue(this.layer.getCode());
- if (componentId == DictionaryUtil.nullValue()) {
+ if (componentId != DictionaryUtil.nullValue()) {
spanBuilder.setComponentId(componentId);
} else {
spanBuilder.setComponent(componentName);
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/trace/ExitSpan.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/trace/ExitSpan.java
index 524aeceeb643201e0b8a4af08112d52aa06b01e4..a60486531fa78fafaf11d0828ff6ebae50bfa744 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/trace/ExitSpan.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/context/trace/ExitSpan.java
@@ -2,6 +2,7 @@ package org.skywalking.apm.agent.core.context.trace;
import org.skywalking.apm.agent.core.dictionary.DictionaryUtil;
import org.skywalking.apm.network.proto.SpanObject;
+import org.skywalking.apm.network.proto.SpanType;
import org.skywalking.apm.network.trace.component.Component;
/**
@@ -113,6 +114,7 @@ public class ExitSpan extends AbstractTracingSpan {
} else {
spanBuilder.setPeer(peer);
}
+ spanBuilder = spanBuilder.setSpanType(SpanType.Exit);
return spanBuilder;
}
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/CollectorDiscoveryService.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/CollectorDiscoveryService.java
index 2b940077c2e4e0ce29239820c69b2d4d856f3296..f464abe24477c065d016ff50fea716b4df168fa2 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/CollectorDiscoveryService.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/CollectorDiscoveryService.java
@@ -1,6 +1,9 @@
package org.skywalking.apm.agent.core.remote;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
import org.skywalking.apm.agent.core.boot.BootService;
+import org.skywalking.apm.agent.core.conf.Config;
/**
* The CollectorDiscoveryService
is responsible for start {@link DiscoveryRestServiceClient}.
@@ -15,9 +18,9 @@ public class CollectorDiscoveryService implements BootService {
@Override
public void boot() throws Throwable {
- Thread collectorClientThread = new Thread(new DiscoveryRestServiceClient(), "collectorClientThread");
- collectorClientThread.setDaemon(true);
- collectorClientThread.start();
+ Executors.newSingleThreadScheduledExecutor()
+ .scheduleAtFixedRate(new DiscoveryRestServiceClient(), 0,
+ Config.Collector.DISCOVERY_CHECK_INTERVAL, TimeUnit.MILLISECONDS);
}
@Override
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/DiscoveryRestServiceClient.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/DiscoveryRestServiceClient.java
index 3d7f0fa70a5d78a23b16a30b67b3a85dc056fc10..843147354ac040acb4aec53c89e4eea9414e5a10 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/DiscoveryRestServiceClient.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/DiscoveryRestServiceClient.java
@@ -31,22 +31,25 @@ public class DiscoveryRestServiceClient implements Runnable {
private volatile int selectedServer = -1;
public DiscoveryRestServiceClient() {
+ if (Config.Collector.SERVERS == null || Config.Collector.SERVERS.trim().length() == 0) {
+ logger.warn("Collector server not configured.");
+ return;
+ }
+
serverList = Config.Collector.SERVERS.split(",");
Random r = new Random();
if (serverList.length > 0) {
selectedServer = r.nextInt(serverList.length);
}
+
}
@Override
public void run() {
- while (true) {
- try {
- try2Sleep(60 * 1000);
- findServerList();
- } catch (Throwable t) {
- logger.error(t, "Find server list fail.");
- }
+ try {
+ findServerList();
+ } catch (Throwable t) {
+ logger.error(t, "Find server list fail.");
}
}
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/GRPCChannelManager.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/GRPCChannelManager.java
index 068809d425f19c07225dd379791860a017d9e02a..f3034efd29fd857311a3fbe44334341a1be14fdb 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/GRPCChannelManager.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/GRPCChannelManager.java
@@ -62,11 +62,16 @@ public class GRPCChannelManager implements BootService, Runnable {
.maxInboundMessageSize(1024 * 1024 * 50)
.usePlaintext(true);
managedChannel = channelBuilder.build();
- reconnect = false;
- notify(GRPCChannelStatus.CONNECTED);
+ if (!managedChannel.isShutdown() && !managedChannel.isTerminated()) {
+ reconnect = false;
+ notify(GRPCChannelStatus.CONNECTED);
+ } else {
+ notify(GRPCChannelStatus.DISCONNECT);
+ }
return;
} catch (Throwable t) {
logger.error(t, "Create channel to {} fail.", server);
+ notify(GRPCChannelStatus.DISCONNECT);
}
}
diff --git a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/TraceSegmentServiceClient.java b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/TraceSegmentServiceClient.java
index 71c8af3142323ca60e0e19159bda1947def91ad2..e6637b7fcc9ea4c05f1483bfa1209afa344a25fa 100644
--- a/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/TraceSegmentServiceClient.java
+++ b/apm-sniffer/apm-agent-core/src/main/java/org/skywalking/apm/agent/core/remote/TraceSegmentServiceClient.java
@@ -26,6 +26,7 @@ import static org.skywalking.apm.agent.core.remote.GRPCChannelStatus.CONNECTED;
*/
public class TraceSegmentServiceClient implements BootService, IConsumer, TracingContextListener, GRPCChannelListener {
private static final ILog logger = LogManager.getLogger(TraceSegmentServiceClient.class);
+ private static final int TIMEOUT = 30 * 1000;
private volatile DataCarrier carrier;
private volatile TraceSegmentServiceGrpc.TraceSegmentServiceStub serviceStub;
@@ -88,7 +89,7 @@ public class TraceSegmentServiceClient implements BootService, IConsumer registryService = getFieldValue(ServiceManager.INSTANCE, "bootedServices");
- assertThat(registryService.size(), is(6));
+ assertThat(registryService.size(), is(7));
assertTraceSegmentServiceClient(ServiceManager.INSTANCE.findService(TraceSegmentServiceClient.class));
assertContextManager(ServiceManager.INSTANCE.findService(ContextManager.class));
@@ -48,7 +55,7 @@ public class ServiceManagerTest {
private void assertTracingContextListener() throws Exception {
List LISTENERS = getFieldValue(TracingContext.ListenerManager.class, "LISTENERS");
- assertThat(LISTENERS.size(), is(2));
+ assertThat(LISTENERS.size(), is(3));
assertThat(LISTENERS.contains(ServiceManager.INSTANCE.findService(ContextManager.class)), is(true));
assertThat(LISTENERS.contains(ServiceManager.INSTANCE.findService(TraceSegmentServiceClient.class)), is(true));
@@ -62,9 +69,7 @@ public class ServiceManagerTest {
assertNotNull(service);
List listeners = getFieldValue(service, "listeners");
- assertEquals(listeners.size(), 1);
- assertThat(listeners.get(0), is((GRPCChannelListener)ServiceManager.INSTANCE.
- findService(TraceSegmentServiceClient.class)));
+ assertEquals(listeners.size(), 3);
}
private void assertSamplingService(SamplingService service) {
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/ContextManagerTest.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/ContextManagerTest.java
index c527656ed3cb358b1707c992dfb94de308b4b4ce..1fe50131962b5d286034a432c8338c59414abe92 100644
--- a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/ContextManagerTest.java
+++ b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/ContextManagerTest.java
@@ -1,9 +1,11 @@
package org.skywalking.apm.agent.core.context;
+import com.google.protobuf.InvalidProtocolBufferException;
import java.util.List;
import org.junit.After;
import org.junit.Before;
import org.junit.BeforeClass;
+import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.skywalking.apm.agent.core.boot.ServiceManager;
@@ -17,11 +19,19 @@ import org.skywalking.apm.agent.core.context.trace.TraceSegment;
import org.skywalking.apm.agent.core.context.trace.TraceSegmentRef;
import org.skywalking.apm.agent.core.context.util.AbstractTracingSpanHelper;
import org.skywalking.apm.agent.core.context.util.SegmentHelper;
-import org.skywalking.apm.agent.core.context.util.SegmentStorage;
-import org.skywalking.apm.agent.core.context.util.SegmentStoragePoint;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
+import org.skywalking.apm.agent.core.test.tools.SegmentStorage;
+import org.skywalking.apm.agent.core.test.tools.SegmentStoragePoint;
import org.skywalking.apm.agent.core.context.util.TraceSegmentRefHelper;
-import org.skywalking.apm.agent.core.context.util.TracingSegmentRunner;
+import org.skywalking.apm.agent.core.test.tools.TracingSegmentRunner;
import org.skywalking.apm.agent.core.dictionary.DictionaryUtil;
+import org.skywalking.apm.network.proto.KeyWithStringValue;
+import org.skywalking.apm.network.proto.LogMessage;
+import org.skywalking.apm.network.proto.SpanObject;
+import org.skywalking.apm.network.proto.SpanType;
+import org.skywalking.apm.network.proto.TraceSegmentObject;
+import org.skywalking.apm.network.proto.TraceSegmentReference;
+import org.skywalking.apm.network.proto.UpstreamSegment;
import org.skywalking.apm.network.trace.component.ComponentsDefine;
import static org.hamcrest.CoreMatchers.is;
@@ -35,10 +45,8 @@ public class ContextManagerTest {
@SegmentStoragePoint
private SegmentStorage tracingData;
- @BeforeClass
- public static void setUpBeforeClass() {
- ServiceManager.INSTANCE.boot();
- }
+ @Rule
+ public AgentServiceRule agentServiceRule = new AgentServiceRule();
@Before
public void setUp() throws Exception {
@@ -186,4 +194,75 @@ public class ContextManagerTest {
RemoteDownstreamConfig.Agent.APPLICATION_INSTANCE_ID = DictionaryUtil.nullValue();
}
+ @Test
+ public void testTransform() throws InvalidProtocolBufferException {
+ ContextCarrier contextCarrier = new ContextCarrier().deserialize("S.1499176688384.581928182.80935.69.1|3|1|#192.168.1.8 :18002|#/portal/|T.1499176688386.581928182.80935.69.2");
+ assertTrue(contextCarrier.isValid());
+
+ AbstractSpan firstEntrySpan = ContextManager.createEntrySpan("/testFirstEntry", contextCarrier);
+ firstEntrySpan.setComponent(ComponentsDefine.TOMCAT);
+ Tags.HTTP.METHOD.set(firstEntrySpan, "GET");
+ Tags.URL.set(firstEntrySpan, "127.0.0.1:8080");
+ SpanLayer.asHttp(firstEntrySpan);
+
+ AbstractSpan secondEntrySpan = ContextManager.createEntrySpan("/testSecondEntry", contextCarrier);
+ secondEntrySpan.setComponent(ComponentsDefine.DUBBO);
+ Tags.URL.set(firstEntrySpan, "dubbo://127.0.0.1:8080");
+ SpanLayer.asRPCFramework(secondEntrySpan);
+
+ ContextCarrier injectContextCarrier = new ContextCarrier();
+ AbstractSpan exitSpan = ContextManager.createExitSpan("/textExitSpan", injectContextCarrier, "127.0.0.1:12800");
+ exitSpan.errorOccurred();
+ exitSpan.log(new RuntimeException("exception"));
+ exitSpan.setComponent(ComponentsDefine.HTTPCLIENT);
+ SpanLayer.asHttp(exitSpan);
+
+ ContextManager.stopSpan();
+ ContextManager.stopSpan();
+ ContextManager.stopSpan();
+
+ TraceSegment actualSegment = tracingData.getTraceSegments().get(0);
+
+ UpstreamSegment upstreamSegment = actualSegment.transform();
+ assertThat(upstreamSegment.getGlobalTraceIdsCount(), is(1));
+ TraceSegmentObject traceSegmentObject = TraceSegmentObject.parseFrom(upstreamSegment.getSegment());
+ TraceSegmentReference reference = traceSegmentObject.getRefs(0);
+
+ assertThat(reference.getEntryServiceName(), is("/portal/"));
+ assertThat(reference.getNetworkAddress(), is("192.168.1.8 :18002"));
+ assertThat(reference.getParentSpanId(), is(3));
+
+ assertThat(traceSegmentObject.getApplicationId(), is(1));
+ assertThat(traceSegmentObject.getRefsCount(), is(1));
+
+ assertThat(traceSegmentObject.getSpansCount(), is(2));
+
+ SpanObject actualSpan = traceSegmentObject.getSpans(1);
+ assertThat(actualSpan.getComponentId(), is(3));
+ assertThat(actualSpan.getComponent(), is(""));
+
+ assertThat(actualSpan.getOperationName(), is("/testSecondEntry"));
+ assertThat(actualSpan.getParentSpanId(), is(-1));
+ assertThat(actualSpan.getSpanId(), is(0));
+ assertThat(actualSpan.getSpanType(), is(SpanType.Entry));
+
+ SpanObject exitSpanObject = traceSegmentObject.getSpans(0);
+ assertThat(exitSpanObject.getComponentId(), is(2));
+ assertThat(exitSpanObject.getComponent(), is(""));
+ assertThat(exitSpanObject.getSpanType(), is(SpanType.Exit));
+
+ assertThat(exitSpanObject.getOperationName(), is("/textExitSpan"));
+ assertThat(exitSpanObject.getParentSpanId(), is(0));
+ assertThat(exitSpanObject.getSpanId(), is(1));
+
+ assertThat(exitSpanObject.getLogsCount(), is(1));
+ LogMessage logMessage = exitSpanObject.getLogs(0);
+ assertThat(logMessage.getDataCount(), is(4));
+ List values = logMessage.getDataList();
+
+ assertThat(values.get(0).getValue(), is("error"));
+ assertThat(values.get(1).getValue(), is(RuntimeException.class.getName()));
+ assertThat(values.get(2).getValue(), is("exception"));
+ assertTrue(values.get(2).getValue().length() <= 4000);
+ }
}
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/IgnoredTracerContextTest.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/IgnoredTracerContextTest.java
index 81fc04c375791d536362caf617e81f7186090f02..ed9f233f8155095039548f5c28e2f6fce3097f6c 100644
--- a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/IgnoredTracerContextTest.java
+++ b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/IgnoredTracerContextTest.java
@@ -3,6 +3,7 @@ package org.skywalking.apm.agent.core.context;
import java.util.LinkedList;
import org.junit.After;
import org.junit.Before;
+import org.junit.Rule;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.skywalking.apm.agent.core.boot.ServiceManager;
@@ -10,9 +11,10 @@ import org.skywalking.apm.agent.core.conf.Config;
import org.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
import org.skywalking.apm.agent.core.context.trace.AbstractSpan;
import org.skywalking.apm.agent.core.context.trace.NoopSpan;
-import org.skywalking.apm.agent.core.context.util.SegmentStorage;
-import org.skywalking.apm.agent.core.context.util.SegmentStoragePoint;
-import org.skywalking.apm.agent.core.context.util.TracingSegmentRunner;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
+import org.skywalking.apm.agent.core.test.tools.SegmentStorage;
+import org.skywalking.apm.agent.core.test.tools.SegmentStoragePoint;
+import org.skywalking.apm.agent.core.test.tools.TracingSegmentRunner;
import static junit.framework.TestCase.assertNull;
import static org.hamcrest.CoreMatchers.is;
@@ -24,6 +26,9 @@ public class IgnoredTracerContextTest {
@SegmentStoragePoint
private SegmentStorage storage;
+ @Rule
+ public AgentServiceRule agentServiceRule = new AgentServiceRule();
+
@Before
public void setUp() throws Exception {
RemoteDownstreamConfig.Agent.APPLICATION_ID = 1;
@@ -57,7 +62,6 @@ public class IgnoredTracerContextTest {
@Test
public void ignoredTraceContextWithExcludeOperationName() {
- ServiceManager.INSTANCE.boot();
AbstractSpan abstractSpan = ContextManager.createEntrySpan("test.js", null);
ContextManager.stopSpan();
@@ -68,7 +72,6 @@ public class IgnoredTracerContextTest {
@Test
public void ignoredTraceContextWithEmptyOperationName() {
- ServiceManager.INSTANCE.boot();
ContextCarrier contextCarrier = new ContextCarrier();
AbstractSpan abstractSpan = ContextManager.createExitSpan("", contextCarrier, "127.0.0.1:2181");
ContextManager.stopSpan();
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/TracingSegmentRunner.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/TracingSegmentRunner.java
deleted file mode 100644
index 94d4ea639fc5cfff69a0a6350aadf3bc5c0ebad5..0000000000000000000000000000000000000000
--- a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/TracingSegmentRunner.java
+++ /dev/null
@@ -1,86 +0,0 @@
-package org.skywalking.apm.agent.core.context.util;
-
-import java.lang.reflect.Field;
-import org.junit.runner.Description;
-import org.junit.runner.notification.Failure;
-import org.junit.runner.notification.RunNotifier;
-import org.junit.runners.BlockJUnit4ClassRunner;
-import org.junit.runners.model.InitializationError;
-import org.skywalking.apm.agent.core.context.IgnoreTracerContextListener;
-import org.skywalking.apm.agent.core.context.IgnoredTracerContext;
-import org.skywalking.apm.agent.core.context.TracingContext;
-import org.skywalking.apm.agent.core.context.TracingContextListener;
-import org.skywalking.apm.agent.core.context.trace.TraceSegment;
-
-public class TracingSegmentRunner extends BlockJUnit4ClassRunner {
- private TracingContextListener tracingContextListener;
- private IgnoreTracerContextListener ignoreTracerContextListener;
- private Field field;
- private Object targetObject;
- private SegmentStorage tracingData;
-
- public TracingSegmentRunner(Class> klass) throws InitializationError {
- super(klass);
- for (Field field : klass.getDeclaredFields()) {
- if (field.isAnnotationPresent(SegmentStoragePoint.class) && field.getType().equals(SegmentStorage.class)) {
- this.field = field;
- this.field.setAccessible(true);
- break;
- }
- }
- }
-
- @Override
- protected Object createTest() throws Exception {
- targetObject = super.createTest();
- return targetObject;
- }
-
- @Override
- public void run(RunNotifier notifier) {
- notifier.addListener(new RunListener());
- super.run(notifier);
- }
-
- class RunListener extends org.junit.runner.notification.RunListener {
- @Override
- public void testStarted(Description description) throws Exception {
- if (field != null) {
- try {
- tracingData = new SegmentStorage();
- field.set(targetObject, tracingData);
- } catch (IllegalAccessException e) {
- }
- }
- tracingContextListener = new TracingContextListener() {
- @Override
- public void afterFinished(TraceSegment traceSegment) {
- tracingData.addTraceSegment(traceSegment);
- }
- };
-
- ignoreTracerContextListener = new IgnoreTracerContextListener() {
- @Override
- public void afterFinished(IgnoredTracerContext tracerContext) {
- tracingData.addIgnoreTraceContext(tracerContext);
- }
- };
- TracingContext.ListenerManager.add(tracingContextListener);
- IgnoredTracerContext.ListenerManager.add(ignoreTracerContextListener);
- super.testStarted(description);
- }
-
- @Override
- public void testFinished(Description description) throws Exception {
- super.testFinished(description);
- TracingContext.ListenerManager.remove(tracingContextListener);
- IgnoredTracerContext.ListenerManager.remove(ignoreTracerContextListener);
- }
-
- @Override
- public void testFailure(Failure failure) throws Exception {
- super.testFailure(failure);
- TracingContext.ListenerManager.remove(tracingContextListener);
- }
- }
-}
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/DiscoveryRestServiceClientTest.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/DiscoveryRestServiceClientTest.java
new file mode 100644
index 0000000000000000000000000000000000000000..d7b5d990616c6ee1c9df285079ced825deac97b2
--- /dev/null
+++ b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/DiscoveryRestServiceClientTest.java
@@ -0,0 +1,108 @@
+package org.skywalking.apm.agent.core.remote;
+
+import com.github.tomakehurst.wiremock.junit.WireMockRule;
+import java.io.IOException;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.skywalking.apm.agent.core.conf.Config;
+import org.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
+
+import static com.github.tomakehurst.wiremock.client.WireMock.aResponse;
+import static com.github.tomakehurst.wiremock.client.WireMock.get;
+import static com.github.tomakehurst.wiremock.client.WireMock.stubFor;
+import static com.github.tomakehurst.wiremock.client.WireMock.urlEqualTo;
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+
+public class DiscoveryRestServiceClientTest {
+
+ @Rule
+ public AgentServiceRule agentServiceRule = new AgentServiceRule();
+
+ private DiscoveryRestServiceClient client;
+ @Rule
+ public WireMockRule wireMockRule = new WireMockRule(8089);
+
+ @Before
+ public void setUpBeforeClass() {
+ Config.Collector.DISCOVERY_CHECK_INTERVAL = 1;
+ stubFor(get(urlEqualTo("/withoutResult"))
+ .willReturn(aResponse()
+ .withStatus(200)
+ .withHeader("Content-Type", "application/json")
+ .withBody("[]")));
+ stubFor(get(urlEqualTo("/withResult"))
+ .willReturn(aResponse()
+ .withStatus(200)
+ .withHeader("Content-Type", "application/json")
+ .withBody("['127.0.0.1:8080','127.0.0.1:8090']")));
+ stubFor(get(urlEqualTo("/withSameResult"))
+ .willReturn(aResponse()
+ .withStatus(200)
+ .withHeader("Content-Type", "application/json")
+ .withBody("['127.0.0.1:8090','127.0.0.1:8080']")));
+ stubFor(get(urlEqualTo("/withDifferenceResult"))
+ .willReturn(aResponse()
+ .withStatus(200)
+ .withHeader("Content-Type", "application/json")
+ .withBody("['127.0.0.1:9090','127.0.0.1:18090']")));
+ stubFor(get(urlEqualTo("/with404"))
+ .willReturn(aResponse()
+ .withStatus(400)));
+ }
+
+ @Test
+ public void testWithoutCollectorServer() throws RESTResponseStatusError, IOException {
+ client = new DiscoveryRestServiceClient();
+ client.run();
+ assertThat(RemoteDownstreamConfig.Collector.GRPC_SERVERS.size(), is(0));
+ }
+
+ @Test
+ public void testWithGRPCAddress() throws RESTResponseStatusError, IOException {
+ Config.Collector.SERVERS = "127.0.0.1:8089";
+ Config.Collector.DISCOVERY_SERVICE_NAME = "/withResult";
+ client = new DiscoveryRestServiceClient();
+ client.run();
+
+ assertThat(RemoteDownstreamConfig.Collector.GRPC_SERVERS.size(), is(2));
+ assertThat(RemoteDownstreamConfig.Collector.GRPC_SERVERS.contains("127.0.0.1:8080"), is(true));
+ assertThat(RemoteDownstreamConfig.Collector.GRPC_SERVERS.contains("127.0.0.1:8090"), is(true));
+ }
+
+ @Test
+ public void testWithoutGRPCAddress() throws RESTResponseStatusError, IOException {
+ Config.Collector.SERVERS = "127.0.0.1:8089";
+ Config.Collector.DISCOVERY_SERVICE_NAME = "/withoutResult";
+ client = new DiscoveryRestServiceClient();
+ client.run();
+
+ assertThat(RemoteDownstreamConfig.Collector.GRPC_SERVERS.size(), is(0));
+ }
+
+ @Test
+ public void testChangeGrpcAddress() throws RESTResponseStatusError, IOException {
+ Config.Collector.SERVERS = "127.0.0.1:8089";
+ Config.Collector.DISCOVERY_SERVICE_NAME = "/withResult";
+ client = new DiscoveryRestServiceClient();
+ client.run();
+
+ Config.Collector.DISCOVERY_SERVICE_NAME = "/withDifferenceResult";
+ client.run();
+
+ assertThat(RemoteDownstreamConfig.Collector.GRPC_SERVERS.size(), is(2));
+ assertThat(RemoteDownstreamConfig.Collector.GRPC_SERVERS.contains("127.0.0.1:9090"), is(true));
+ assertThat(RemoteDownstreamConfig.Collector.GRPC_SERVERS.contains("127.0.0.1:18090"), is(true));
+ }
+
+ @After
+ public void tearDown() {
+ Config.Collector.SERVERS = "";
+ Config.Collector.DISCOVERY_SERVICE_NAME = "/grpc/address";
+ RemoteDownstreamConfig.Collector.GRPC_SERVERS.clear();
+ }
+
+}
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/GRPCChannelManagerTest.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/GRPCChannelManagerTest.java
new file mode 100644
index 0000000000000000000000000000000000000000..40e795e8a35018b505739c4671e113d9d631b4c4
--- /dev/null
+++ b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/GRPCChannelManagerTest.java
@@ -0,0 +1,104 @@
+package org.skywalking.apm.agent.core.remote;
+
+import io.grpc.NameResolver;
+import io.grpc.Status;
+import io.grpc.StatusRuntimeException;
+import io.grpc.internal.DnsNameResolverProvider;
+import io.grpc.netty.NettyChannelBuilder;
+import io.grpc.testing.GrpcServerRule;
+import java.util.ArrayList;
+import java.util.List;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Spy;
+import org.powermock.core.classloader.annotations.PrepareForTest;
+import org.powermock.modules.junit4.PowerMockRunner;
+import org.powermock.reflect.Whitebox;
+import org.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.anyString;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.powermock.api.mockito.PowerMockito.mockStatic;
+import static org.powermock.api.mockito.PowerMockito.when;
+
+@RunWith(PowerMockRunner.class)
+@PrepareForTest({GRPCChannelManager.class, NettyChannelBuilder.class})
+public class GRPCChannelManagerTest {
+
+ @Rule
+ private GrpcServerRule grpcServerRule = new GrpcServerRule().directExecutor();
+
+ @Spy
+ private GRPCChannelManager grpcChannelManager = new GRPCChannelManager();
+
+ @Mock
+ private NettyChannelBuilder mock;
+
+ @Spy
+ private MockGRPCChannelListener listener = new MockGRPCChannelListener();
+
+ @Before
+ public void setUp() throws Throwable {
+ List grpcServers = new ArrayList();
+ grpcServers.add("127.0.0.1:2181");
+ RemoteDownstreamConfig.Collector.GRPC_SERVERS = grpcServers;
+ Whitebox.setInternalState(grpcChannelManager, "retryCycle", 1);
+
+ mockStatic(NettyChannelBuilder.class);
+ when(NettyChannelBuilder.forAddress(anyString(), anyInt())).thenReturn(mock);
+ when(mock.nameResolverFactory(any(NameResolver.Factory.class))).thenReturn(mock);
+ when(mock.maxInboundMessageSize(anyInt())).thenReturn(mock);
+ when(mock.usePlaintext(true)).thenReturn(mock);
+ when(mock.build()).thenReturn(grpcServerRule.getChannel());
+
+ grpcChannelManager.addChannelListener(listener);
+ }
+
+ @Test
+ public void changeStatusToConnectedWithReportError() throws Throwable {
+ grpcChannelManager.reportError(new StatusRuntimeException(Status.ABORTED));
+ grpcChannelManager.run();
+
+ verify(listener, times(1)).statusChanged(GRPCChannelStatus.CONNECTED);
+ assertThat(listener.status, is(GRPCChannelStatus.CONNECTED));
+ }
+
+ @Test
+ public void changeStatusToDisConnectedWithReportError() throws Throwable {
+ doThrow(new RuntimeException()).when(mock).nameResolverFactory(any(NameResolver.Factory.class));
+ grpcChannelManager.run();
+
+ verify(listener, times(1)).statusChanged(GRPCChannelStatus.DISCONNECT);
+ assertThat(listener.status, is(GRPCChannelStatus.DISCONNECT));
+ }
+
+ @Test
+ public void reportErrorWithoutChangeStatus() throws Throwable {
+ grpcChannelManager.run();
+ grpcChannelManager.reportError(new RuntimeException());
+ grpcChannelManager.run();
+
+ verify(listener, times(1)).statusChanged(GRPCChannelStatus.CONNECTED);
+ assertThat(listener.status, is(GRPCChannelStatus.CONNECTED));
+ }
+
+ private class MockGRPCChannelListener implements GRPCChannelListener {
+ private GRPCChannelStatus status;
+
+ @Override
+ public void statusChanged(GRPCChannelStatus status) {
+ this.status = status;
+ }
+ }
+
+}
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/TraceSegmentServiceClientTest.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/TraceSegmentServiceClientTest.java
new file mode 100644
index 0000000000000000000000000000000000000000..a180eb1fa3377552935e9233a66cd84993689413
--- /dev/null
+++ b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/remote/TraceSegmentServiceClientTest.java
@@ -0,0 +1,136 @@
+package org.skywalking.apm.agent.core.remote;
+
+import com.google.protobuf.InvalidProtocolBufferException;
+import io.grpc.stub.StreamObserver;
+import io.grpc.testing.GrpcServerRule;
+import java.util.ArrayList;
+import java.util.List;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.powermock.reflect.Whitebox;
+import org.skywalking.apm.agent.core.boot.ServiceManager;
+import org.skywalking.apm.agent.core.conf.RemoteDownstreamConfig;
+import org.skywalking.apm.agent.core.context.ContextManager;
+import org.skywalking.apm.agent.core.context.tag.Tags;
+import org.skywalking.apm.agent.core.context.trace.AbstractSpan;
+import org.skywalking.apm.agent.core.context.trace.SpanLayer;
+import org.skywalking.apm.agent.core.test.tools.AgentServiceRule;
+import org.skywalking.apm.agent.core.test.tools.SegmentStorage;
+import org.skywalking.apm.agent.core.test.tools.SegmentStoragePoint;
+import org.skywalking.apm.agent.core.test.tools.TracingSegmentRunner;
+import org.skywalking.apm.network.proto.Downstream;
+import org.skywalking.apm.network.proto.SpanObject;
+import org.skywalking.apm.network.proto.SpanType;
+import org.skywalking.apm.network.proto.TraceSegmentObject;
+import org.skywalking.apm.network.proto.TraceSegmentServiceGrpc;
+import org.skywalking.apm.network.proto.UpstreamSegment;
+import org.skywalking.apm.network.trace.component.ComponentsDefine;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.mockito.Mockito.spy;
+
+@RunWith(TracingSegmentRunner.class)
+public class TraceSegmentServiceClientTest {
+
+ @Rule
+ public AgentServiceRule agentServiceRule = new AgentServiceRule();
+
+ @Rule
+ public GrpcServerRule grpcServerRule = new GrpcServerRule().directExecutor();
+
+ @SegmentStoragePoint
+ private SegmentStorage storage;
+
+ private TraceSegmentServiceClient serviceClient = new TraceSegmentServiceClient();
+ private List upstreamSegments;
+
+ private TraceSegmentServiceGrpc.TraceSegmentServiceImplBase serviceImplBase = new TraceSegmentServiceGrpc.TraceSegmentServiceImplBase() {
+ @Override
+ public StreamObserver collect(final StreamObserver responseObserver) {
+ return new StreamObserver() {
+ @Override
+ public void onNext(UpstreamSegment value) {
+ upstreamSegments.add(value);
+ }
+
+ @Override
+ public void onError(Throwable t) {
+ }
+
+ @Override
+ public void onCompleted() {
+ responseObserver.onNext(Downstream.getDefaultInstance());
+ responseObserver.onCompleted();
+ }
+ };
+ }
+ };
+
+ @BeforeClass
+ public static void setUpBeforeClass() {
+ RemoteDownstreamConfig.Agent.APPLICATION_ID = 1;
+ RemoteDownstreamConfig.Agent.APPLICATION_INSTANCE_ID = 1;
+ }
+
+ @Before
+ public void setUp() throws Throwable {
+ Whitebox.setInternalState(ServiceManager.INSTANCE.findService(GRPCChannelManager.class), "reconnect", false);
+ spy(serviceClient);
+
+ Whitebox.setInternalState(serviceClient, "serviceStub",
+ TraceSegmentServiceGrpc.newStub(grpcServerRule.getChannel()));
+ Whitebox.setInternalState(serviceClient, "status", GRPCChannelStatus.CONNECTED);
+
+ upstreamSegments = new ArrayList();
+ }
+
+ @Test
+ public void testSendTraceSegmentWithoutException() throws InvalidProtocolBufferException {
+ grpcServerRule.getServiceRegistry().addService(serviceImplBase);
+
+ AbstractSpan firstEntrySpan = ContextManager.createEntrySpan("/testFirstEntry", null);
+ firstEntrySpan.setComponent(ComponentsDefine.TOMCAT);
+ Tags.HTTP.METHOD.set(firstEntrySpan, "GET");
+ Tags.URL.set(firstEntrySpan, "127.0.0.1:8080");
+ SpanLayer.asHttp(firstEntrySpan);
+ ContextManager.stopSpan();
+
+ serviceClient.consume(storage.getTraceSegments());
+
+ assertThat(upstreamSegments.size(), is(1));
+ UpstreamSegment upstreamSegment = upstreamSegments.get(0);
+ assertThat(upstreamSegment.getGlobalTraceIdsCount(), is(1));
+ TraceSegmentObject traceSegmentObject = TraceSegmentObject.parseFrom(upstreamSegment.getSegment());
+ assertThat(traceSegmentObject.getRefsCount(), is(0));
+ assertThat(traceSegmentObject.getSpansCount(), is(1));
+
+ SpanObject spanObject = traceSegmentObject.getSpans(0);
+ assertThat(spanObject.getSpanType(), is(SpanType.Entry));
+ assertThat(spanObject.getSpanId(), is(0));
+ assertThat(spanObject.getParentSpanId(), is(-1));
+ }
+
+ @Test
+ public void testSendTraceSegmentWithException() throws InvalidProtocolBufferException {
+ grpcServerRule.getServiceRegistry().addService(serviceImplBase);
+
+ AbstractSpan firstEntrySpan = ContextManager.createEntrySpan("/testFirstEntry", null);
+ firstEntrySpan.setComponent(ComponentsDefine.TOMCAT);
+ Tags.HTTP.METHOD.set(firstEntrySpan, "GET");
+ Tags.URL.set(firstEntrySpan, "127.0.0.1:8080");
+ SpanLayer.asHttp(firstEntrySpan);
+ ContextManager.stopSpan();
+ grpcServerRule.getServer().shutdownNow();
+ serviceClient.consume(storage.getTraceSegments());
+
+ assertThat(upstreamSegments.size(), is(0));
+
+ boolean reconnect = Whitebox.getInternalState(ServiceManager.INSTANCE.findService(GRPCChannelManager.class), "reconnect");
+ assertThat(reconnect, is(true));
+
+ }
+}
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/AgentServiceRule.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/AgentServiceRule.java
new file mode 100644
index 0000000000000000000000000000000000000000..ea1b48cce5ca167668b75edc5b629b3847d03afe
--- /dev/null
+++ b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/AgentServiceRule.java
@@ -0,0 +1,28 @@
+package org.skywalking.apm.agent.core.test.tools;
+
+import java.util.HashMap;
+import java.util.LinkedList;
+import org.junit.rules.ExternalResource;
+import org.powermock.reflect.Whitebox;
+import org.skywalking.apm.agent.core.boot.BootService;
+import org.skywalking.apm.agent.core.boot.ServiceManager;
+import org.skywalking.apm.agent.core.context.IgnoredTracerContext;
+import org.skywalking.apm.agent.core.context.TracingContext;
+import org.skywalking.apm.agent.core.context.TracingContextListener;
+
+public class AgentServiceRule extends ExternalResource {
+
+ @Override
+ protected void after() {
+ super.after();
+ Whitebox.setInternalState(ServiceManager.INSTANCE, "bootedServices", new HashMap());
+ Whitebox.setInternalState(TracingContext.ListenerManager.class, "LISTENERS", new LinkedList() );
+ Whitebox.setInternalState(IgnoredTracerContext.ListenerManager.class, "LISTENERS", new LinkedList() );
+ }
+
+ @Override
+ protected void before() throws Throwable {
+ super.before();
+ ServiceManager.INSTANCE.boot();
+ }
+}
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/SegmentStorage.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/SegmentStorage.java
similarity index 91%
rename from apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/SegmentStorage.java
rename to apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/SegmentStorage.java
index ea310bba99eac0ebd360825afcde7c909b409dec..de4742cb257bb551c257268de86e9b2b74fe8d8d 100644
--- a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/SegmentStorage.java
+++ b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/SegmentStorage.java
@@ -1,4 +1,4 @@
-package org.skywalking.apm.agent.core.context.util;
+package org.skywalking.apm.agent.core.test.tools;
import java.util.LinkedList;
import java.util.List;
@@ -9,7 +9,7 @@ public class SegmentStorage {
private LinkedList traceSegments;
private LinkedList ignoredTracerContexts;
- SegmentStorage() {
+ public SegmentStorage() {
traceSegments = new LinkedList();
ignoredTracerContexts = new LinkedList();
}
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/SegmentStoragePoint.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/SegmentStoragePoint.java
similarity index 83%
rename from apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/SegmentStoragePoint.java
rename to apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/SegmentStoragePoint.java
index 21a450c5c14a2c7d864fed913e38b2c0af7d50ce..2bafc0b9479d6a937cfd2a8ca380ec2ff1b73c5c 100644
--- a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/context/util/SegmentStoragePoint.java
+++ b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/SegmentStoragePoint.java
@@ -1,4 +1,4 @@
-package org.skywalking.apm.agent.core.context.util;
+package org.skywalking.apm.agent.core.test.tools;
import java.lang.annotation.ElementType;
import java.lang.annotation.Retention;
diff --git a/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/TracingSegmentRunner.java b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/TracingSegmentRunner.java
new file mode 100644
index 0000000000000000000000000000000000000000..ea77dad8b820393b472b08c7025ef22f185759e3
--- /dev/null
+++ b/apm-sniffer/apm-agent-core/src/test/java/org/skywalking/apm/agent/core/test/tools/TracingSegmentRunner.java
@@ -0,0 +1,73 @@
+package org.skywalking.apm.agent.core.test.tools;
+
+import java.lang.reflect.Field;
+import org.junit.runners.BlockJUnit4ClassRunner;
+import org.junit.runners.model.FrameworkMethod;
+import org.junit.runners.model.InitializationError;
+import org.junit.runners.model.Statement;
+import org.skywalking.apm.agent.core.context.IgnoreTracerContextListener;
+import org.skywalking.apm.agent.core.context.IgnoredTracerContext;
+import org.skywalking.apm.agent.core.context.TracingContext;
+import org.skywalking.apm.agent.core.context.TracingContextListener;
+import org.skywalking.apm.agent.core.context.trace.TraceSegment;
+
+public class TracingSegmentRunner extends BlockJUnit4ClassRunner {
+ private TracingContextListener tracingContextListener;
+ private IgnoreTracerContextListener ignoreTracerContextListener;
+ private Field field;
+ private Object targetObject;
+ private SegmentStorage tracingData;
+
+ public TracingSegmentRunner(Class> klass) throws InitializationError {
+ super(klass);
+ for (Field field : klass.getDeclaredFields()) {
+ if (field.isAnnotationPresent(SegmentStoragePoint.class) && field.getType().equals(SegmentStorage.class)) {
+ this.field = field;
+ this.field.setAccessible(true);
+ break;
+ }
+ }
+ }
+
+ @Override
+ protected Object createTest() throws Exception {
+ targetObject = super.createTest();
+ return targetObject;
+ }
+
+ @Override protected Statement withAfters(FrameworkMethod method, Object target, final Statement statement) {
+ return new Statement() {
+ @Override public void evaluate() throws Throwable {
+ if (field != null) {
+ try {
+ tracingData = new SegmentStorage();
+ field.set(targetObject, tracingData);
+ } catch (IllegalAccessException e) {
+ }
+ }
+ tracingContextListener = new TracingContextListener() {
+ @Override
+ public void afterFinished(TraceSegment traceSegment) {
+ tracingData.addTraceSegment(traceSegment);
+ }
+ };
+
+ ignoreTracerContextListener = new IgnoreTracerContextListener() {
+ @Override
+ public void afterFinished(IgnoredTracerContext tracerContext) {
+ tracingData.addIgnoreTraceContext(tracerContext);
+ }
+ };
+
+ TracingContext.ListenerManager.add(tracingContextListener);
+ IgnoredTracerContext.ListenerManager.add(ignoreTracerContextListener);
+ try {
+ statement.evaluate();
+ } finally {
+ TracingContext.ListenerManager.remove(tracingContextListener);
+ IgnoredTracerContext.ListenerManager.remove(ignoreTracerContextListener);
+ }
+ }
+ };
+ }
+}