Skip to content

Commit

Permalink
support mongo db.instance and AggregateOperation test scenarios
Browse files Browse the repository at this point in the history
  • Loading branch information
youjie_li committed Nov 23, 2024
1 parent 5801ee2 commit 80b5ac8
Show file tree
Hide file tree
Showing 11 changed files with 217 additions and 74 deletions.
37 changes: 21 additions & 16 deletions ...in/src/main/java/org/apache/skywalking/apm/plugin/mongodb/v3/support/MongoSpanHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,17 +21,20 @@
import com.mongodb.MongoNamespace;
import org.apache.skywalking.apm.agent.core.context.ContextCarrier;
import org.apache.skywalking.apm.agent.core.context.ContextManager;
import org.apache.skywalking.apm.agent.core.context.tag.AbstractTag;
import org.apache.skywalking.apm.agent.core.context.tag.Tags;
import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer;
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;
import org.apache.skywalking.apm.plugin.mongodb.v3.MongoPluginConfig;
import org.apache.skywalking.apm.util.StringUtil;

import java.lang.reflect.Field;

public class MongoSpanHelper {

private static final AbstractTag<String> DB_COLLECTION_TAG = Tags.ofKey("db.collection");

private MongoSpanHelper() {
}

Expand All @@ -43,33 +46,35 @@ public static void createExitSpan(String executeMethod, String remotePeer, Objec
SpanLayer.asDB(span);

try {
Field namespaceField = operation.getClass().getDeclaredField("namespace");
Field.setAccessible(new Field[]{namespaceField}, true);
MongoNamespace namespace = (MongoNamespace) namespaceField.get(operation);
Tags.DB_INSTANCE.set(span, namespace.getFullName());
MongoNamespace namespace = tryToGetMongoNamespace(operation);
extractTagsFromNamespace(span, namespace);
} catch (Exception e) {
try {
Field wrappedField = operation.getClass().getDeclaredField("wrapped");
Field.setAccessible(new Field[]{wrappedField}, true);
Object wrappedOperation = wrappedField.get(operation);
Field wrappedNamespaceField = wrappedOperation.getClass().getDeclaredField("namespace");
Field.setAccessible(new Field[]{wrappedNamespaceField}, true);
MongoNamespace wrappedNamespace = (MongoNamespace) wrappedNamespaceField.get(wrappedOperation);
Tags.DB_INSTANCE.set(span, wrappedNamespace.getFullName());
MongoNamespace namespace = tryToGetMongoNamespace(wrappedOperation);
extractTagsFromNamespace(span, namespace);
} catch (Exception e2) {

}
}

if (operation instanceof EnhancedInstance) {
Object databaseName = ((EnhancedInstance) operation).getSkyWalkingDynamicField();
if (databaseName != null) {
Tags.DB_INSTANCE.set(span, (String) databaseName);
}
}

if (MongoPluginConfig.Plugin.MongoDB.TRACE_PARAM) {
Tags.DB_BIND_VARIABLES.set(span, MongoOperationHelper.getTraceParam(operation));
}
}

private static void extractTagsFromNamespace(AbstractSpan span, MongoNamespace namespace) {
Tags.DB_INSTANCE.set(span, namespace.getDatabaseName());
if (StringUtil.isNotEmpty(namespace.getCollectionName())) {
span.tag(DB_COLLECTION_TAG, namespace.getCollectionName());
}
}

private static MongoNamespace tryToGetMongoNamespace(Object operation) throws IllegalAccessException, NoSuchFieldException {
Field namespaceField = operation.getClass().getDeclaredField("namespace");
Field.setAccessible(new Field[]{namespaceField}, true);
return (MongoNamespace) namespaceField.get(operation);
}
}
41 changes: 38 additions & 3 deletions ...a/org/apache/skywalking/apm/plugin/mongodb/v3/interceptor/v30/MongoDBInterceptorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,10 @@
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

import com.mongodb.operation.AggregateOperation;
import org.apache.skywalking.apm.agent.core.context.trace.AbstractTracingSpan;
import org.apache.skywalking.apm.agent.core.context.trace.LogDataEntity;
import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer;
Expand Down Expand Up @@ -70,6 +73,7 @@ public class MongoDBInterceptorTest {
@Mock
private EnhancedInstance enhancedInstance;

private Decoder decoder;
private Object[] arguments;
private Class[] argumentTypes;

Expand All @@ -89,7 +93,7 @@ public void setUp() throws Exception {
BsonDocument document = new BsonDocument();
document.append("name", new BsonString("by"));
MongoNamespace mongoNamespace = new MongoNamespace("test.user");
Decoder decoder = mock(Decoder.class);
decoder = mock(Decoder.class);
FindOperation findOperation = new FindOperation(mongoNamespace, decoder);
findOperation.filter(document);

Expand All @@ -108,6 +112,24 @@ public void testIntercept() throws Throwable {
assertRedisSpan(spans.get(0));
}

@Test
public void testAggregateOperationIntercept() throws Throwable {
MongoNamespace mongoNamespace = new MongoNamespace("test.user");
BsonDocument matchStage = new BsonDocument("$match", new BsonDocument("name", new BsonString("by")));
List<BsonDocument> pipeline = Collections.singletonList(matchStage);
AggregateOperation<BsonDocument> aggregateOperation = new AggregateOperation(mongoNamespace, pipeline, decoder);
Object[] arguments = {aggregateOperation};
Class[] argumentTypes = {aggregateOperation.getClass()};

interceptor.beforeMethod(enhancedInstance, getExecuteMethod(), arguments, argumentTypes, null);
interceptor.afterMethod(enhancedInstance, getExecuteMethod(), arguments, argumentTypes, null);

MatcherAssert.assertThat(segmentStorage.getTraceSegments().size(), is(1));
TraceSegment traceSegment = segmentStorage.getTraceSegments().get(0);
List<AbstractTracingSpan> spans = SegmentHelper.getSpans(traceSegment);
assertMongoAggregateOperationSpan(spans.get(0));
}

@Test
public void testInterceptWithException() throws Throwable {
interceptor.beforeMethod(enhancedInstance, getExecuteMethod(), arguments, argumentTypes, null);
Expand All @@ -128,9 +150,22 @@ private void assertRedisSpan(AbstractTracingSpan span) {
assertThat(span.getOperationName(), is("MongoDB/FindOperation"));
assertThat(SpanHelper.getComponentId(span), is(42));
List<TagValuePair> tags = SpanHelper.getTags(span);
assertThat(tags.get(1).getValue(), is("test.user"));
assertThat(tags.get(2).getValue(), is("{\"name\": \"by\"}"));
assertThat(tags.get(0).getValue(), is("MongoDB"));
assertThat(tags.get(1).getValue(), is("test"));
assertThat(tags.get(2).getValue(), is("user"));
assertThat(tags.get(3).getValue(), is("{\"name\": \"by\"}"));
assertThat(span.isExit(), is(true));
assertThat(SpanHelper.getLayer(span), CoreMatchers.is(SpanLayer.DB));
}

private void assertMongoAggregateOperationSpan(AbstractTracingSpan span) {
assertThat(span.getOperationName(), is("MongoDB/AggregateOperation"));
assertThat(SpanHelper.getComponentId(span), is(42));
List<TagValuePair> tags = SpanHelper.getTags(span);
assertThat(tags.get(0).getValue(), is("MongoDB"));
assertThat(tags.get(1).getValue(), is("test"));
assertThat(tags.get(2).getValue(), is("user"));
assertThat(tags.get(3).getValue(), is("{\"$match\": {\"name\": \"by\"}},"));
assertThat(span.isExit(), is(true));
assertThat(SpanHelper.getLayer(span), CoreMatchers.is(SpanLayer.DB));
}
Expand Down
10 changes: 6 additions & 4 deletions ...alking/apm/plugin/mongodb/v3/interceptor/v37/MongoDBOperationExecutorInterceptorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -149,9 +149,10 @@ private void assertMongoAggregateOperationSpan(AbstractTracingSpan span) {
assertThat(span.getOperationName(), is("MongoDB/AggregateOperation"));
assertThat(SpanHelper.getComponentId(span), is(42));
List<TagValuePair> tags = SpanHelper.getTags(span);
assertThat(tags.get(1).getValue(), is("test.user"));
assertThat(tags.get(2).getValue(), is("{\"$match\": {\"name\": \"by\"}},"));
assertThat(tags.get(0).getValue(), is("MongoDB"));
assertThat(tags.get(1).getValue(), is("test"));
assertThat(tags.get(2).getValue(), is("user"));
assertThat(tags.get(3).getValue(), is("{\"$match\": {\"name\": \"by\"}},"));
assertThat(span.isExit(), is(true));
assertThat(SpanHelper.getLayer(span), CoreMatchers.is(SpanLayer.DB));
}
Expand All @@ -160,9 +161,10 @@ private void assertMongoFindOperationSpan(AbstractTracingSpan span) {
assertThat(span.getOperationName(), is("MongoDB/FindOperation"));
assertThat(SpanHelper.getComponentId(span), is(42));
List<TagValuePair> tags = SpanHelper.getTags(span);
assertThat(tags.get(1).getValue(), is("test.user"));
assertThat(tags.get(2).getValue(), is("{\"name\": \"by\"}"));
assertThat(tags.get(0).getValue(), is("MongoDB"));
assertThat(tags.get(1).getValue(), is("test"));
assertThat(tags.get(2).getValue(), is("user"));
assertThat(tags.get(3).getValue(), is("{\"name\": \"by\"}"));
assertThat(span.isExit(), is(true));
assertThat(SpanHelper.getLayer(span), CoreMatchers.is(SpanLayer.DB));
}
Expand Down
2 changes: 1 addition & 1 deletion ...che/skywalking/apm/plugin/mongodb/v4/interceptor/MongoDBOperationExecutorInterceptor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,7 @@ public class MongoDBOperationExecutorInterceptor implements InstanceMethodsAroun
@Override
public void beforeMethod(EnhancedInstance objInst, Method method, Object[] allArguments, Class<?>[] argumentsTypes,
MethodInterceptResult result) {
String executeMethod = allArguments[0].getClass().getSimpleName();
String executeMethod = argumentsTypes[0].getSimpleName();
// OperationExecutor has included th remotePeer
// See: MongoDBClientDelegateInterceptor.afterMethod
String remotePeer = (String) objInst.getSkyWalkingDynamicField();
Expand Down
3 changes: 1 addition & 2 deletions ...g/apm/plugin/mongodb/v4/interceptor/operation/OperationNamespaceConstructInterceptor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,7 @@ public class OperationNamespaceConstructInterceptor implements InstanceConstruct
@Override
public void onConstruct(EnhancedInstance objInst, Object[] allArguments) {
MongoNamespace mongoNamespace = (MongoNamespace) allArguments[0];
String databaseName = mongoNamespace.getDatabaseName();
objInst.setSkyWalkingDynamicField(databaseName);
objInst.setSkyWalkingDynamicField(mongoNamespace);
}

}
38 changes: 13 additions & 25 deletions ...in/src/main/java/org/apache/skywalking/apm/plugin/mongodb/v4/support/MongoSpanHelper.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,16 +21,18 @@
import com.mongodb.MongoNamespace;
import org.apache.skywalking.apm.agent.core.context.ContextCarrier;
import org.apache.skywalking.apm.agent.core.context.ContextManager;
import org.apache.skywalking.apm.agent.core.context.tag.AbstractTag;
import org.apache.skywalking.apm.agent.core.context.tag.Tags;
import org.apache.skywalking.apm.agent.core.context.trace.AbstractSpan;
import org.apache.skywalking.apm.agent.core.context.trace.SpanLayer;
import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.EnhancedInstance;
import org.apache.skywalking.apm.network.trace.component.ComponentsDefine;

import java.lang.reflect.Field;
import org.apache.skywalking.apm.util.StringUtil;

public class MongoSpanHelper {

private static final AbstractTag<String> DB_COLLECTION_TAG = Tags.ofKey("db.collection");

private MongoSpanHelper() {
}

Expand All @@ -42,34 +44,19 @@ private MongoSpanHelper() {
*/
public static void createExitSpan(String executeMethod, String remotePeer, Object operation) {
AbstractSpan span = ContextManager.createExitSpan(
MongoConstants.MONGO_DB_OP_PREFIX + executeMethod, new ContextCarrier(), remotePeer);
MongoConstants.MONGO_DB_OP_PREFIX + executeMethod, new ContextCarrier(), remotePeer);
span.setComponent(ComponentsDefine.MONGO_DRIVER);
Tags.DB_TYPE.set(span, MongoConstants.DB_TYPE);
SpanLayer.asDB(span);

try {
Field namespaceField = operation.getClass().getDeclaredField("namespace");
Field.setAccessible(new Field[]{namespaceField}, true);
MongoNamespace namespace = (MongoNamespace) namespaceField.get(operation);
Tags.DB_INSTANCE.set(span, namespace.getFullName());
} catch (Exception e) {
try {
Field wrappedField = operation.getClass().getDeclaredField("wrapped");
Field.setAccessible(new Field[]{wrappedField}, true);
Object wrappedOperation = wrappedField.get(operation);
Field wrappedNamespaceField = wrappedOperation.getClass().getDeclaredField("namespace");
Field.setAccessible(new Field[]{wrappedNamespaceField}, true);
MongoNamespace wrappedNamespace = (MongoNamespace) wrappedNamespaceField.get(wrappedOperation);
Tags.DB_INSTANCE.set(span, wrappedNamespace.getFullName());
} catch (Exception e2) {

}
}

if (operation instanceof EnhancedInstance) {
Object databaseName = ((EnhancedInstance) operation).getSkyWalkingDynamicField();
if (databaseName != null) {
Tags.DB_INSTANCE.set(span, (String) databaseName);
Object namespaceObj = ((EnhancedInstance) operation).getSkyWalkingDynamicField();
if (namespaceObj != null) {
MongoNamespace namespace = (MongoNamespace) namespaceObj;
Tags.DB_INSTANCE.set(span, namespace.getDatabaseName());
if (StringUtil.isNotEmpty(namespace.getCollectionName())) {
span.tag(DB_COLLECTION_TAG, namespace.getCollectionName());
}
}
}

Expand All @@ -78,3 +65,4 @@ public static void createExitSpan(String executeMethod, String remotePeer, Objec
}
}
}

Loading

0 comments on commit 80b5ac8

Please sign in to comment.