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 24, 2024
1 parent 80b5ac8 commit 3e64803
Show file tree
Hide file tree
Showing 5 changed files with 138 additions and 38 deletions.
72 changes: 55 additions & 17 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 @@ -19,6 +19,7 @@
package org.apache.skywalking.apm.plugin.mongodb.v3.support;

import com.mongodb.MongoNamespace;
import lombok.SneakyThrows;
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;
Expand All @@ -38,26 +39,18 @@ public class MongoSpanHelper {
private MongoSpanHelper() {
}

@SneakyThrows
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 {
MongoNamespace namespace = tryToGetMongoNamespace(operation);
Field[] declaredFields = operation.getClass().getDeclaredFields();
MongoNamespace namespace = tryToGetMongoNamespace(operation, declaredFields);
if (namespace != null) {
extractTagsFromNamespace(span, namespace);
} catch (Exception e) {
try {
Field wrappedField = operation.getClass().getDeclaredField("wrapped");
Field.setAccessible(new Field[]{wrappedField}, true);
Object wrappedOperation = wrappedField.get(operation);
MongoNamespace namespace = tryToGetMongoNamespace(wrappedOperation);
extractTagsFromNamespace(span, namespace);
} catch (Exception e2) {

}
}

if (MongoPluginConfig.Plugin.MongoDB.TRACE_PARAM) {
Expand All @@ -72,9 +65,54 @@ private static void extractTagsFromNamespace(AbstractSpan span, MongoNamespace n
}
}

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);
private static MongoNamespace tryToGetMongoNamespace(Object operation, Field[] declaredFields) throws IllegalAccessException {
Field namespaceField = null;
Field wrappedField = null;
Field databaseField = null;
Field collectionField = null;
for (Field field : declaredFields) {
if ("namespace".equals(field.getName())) {
namespaceField = field;
Field.setAccessible(new Field[]{field}, true);
}
if ("wrapped".equals(field.getName())) {
wrappedField = field;
Field.setAccessible(new Field[]{field}, true);
}
if ("databaseName".equals(field.getName())) {
databaseField = field;
Field.setAccessible(new Field[]{field}, true);
}
if ("collectionName".equals(field.getName())) {
collectionField = field;
Field.setAccessible(new Field[]{field}, true);
}
}
if (namespaceField != null) {
return (MongoNamespace) namespaceField.get(operation);
}
String database = null;
String collection = null;
if (databaseField != null) {
database = (String) databaseField.get(operation);
}
if (collectionField != null) {
collection = (String) collectionField.get(operation);
}
if (database != null && collection == null) {
return new MongoNamespace(database);
}
if (database != null && collection != null) {
return new MongoNamespace(database, collection);
}
if (wrappedField != null) {
Object wrapped = wrappedField.get(operation);
if (wrapped != null && wrapped != operation) {
Field[] declaredFieldsInWrapped = wrapped.getClass().getDeclaredFields();
return tryToGetMongoNamespace(wrapped, declaredFieldsInWrapped);
}
}
return null;

}
}
27 changes: 27 additions & 0 deletions ...alking/apm/plugin/mongodb/v3/interceptor/v37/MongoDBOperationExecutorInterceptorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import com.mongodb.ReadConcern;
import com.mongodb.client.internal.OperationExecutor;
import com.mongodb.operation.AggregateOperation;
import com.mongodb.operation.CreateCollectionOperation;
import com.mongodb.operation.FindOperation;
import com.mongodb.operation.WriteOperation;
import org.apache.skywalking.apm.agent.core.context.trace.AbstractTracingSpan;
Expand Down Expand Up @@ -112,6 +113,20 @@ public void testIntercept() throws Throwable {
assertMongoFindOperationSpan(spans.get(0));
}

@Test
public void testCreateCollectionOperationIntercept() throws Throwable {
CreateCollectionOperation createCollectionOperation = new CreateCollectionOperation("test", "user");
Object[] arguments = {createCollectionOperation};
Class[] argumentTypes = {createCollectionOperation.getClass()};
interceptor.beforeMethod(enhancedInstance, getMethod(), arguments, argumentTypes, null);
interceptor.afterMethod(enhancedInstance, getMethod(), arguments, argumentTypes, null);

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

@Test
public void testAggregateOperationIntercept() throws Throwable {
MongoNamespace mongoNamespace = new MongoNamespace("test.user");
Expand Down Expand Up @@ -145,6 +160,18 @@ public void testInterceptFindOperationWithException() throws Throwable {
SpanAssert.assertException(logDataEntities.get(0), RuntimeException.class);
}

private void assertMongoCreateCollectionOperationSpan(AbstractTracingSpan span) {
assertThat(span.getOperationName(), is("MongoDB/CreateCollectionOperation"));
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("user"));
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));
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 = argumentsTypes[0].getSimpleName();
String executeMethod = allArguments[0].getClass().getSimpleName();
// OperationExecutor has included th remotePeer
// See: MongoDBClientDelegateInterceptor.afterMethod
String remotePeer = (String) objInst.getSkyWalkingDynamicField();
Expand Down
13 changes: 8 additions & 5 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 @@ -38,9 +38,10 @@ private MongoSpanHelper() {

/**
* createExitSpan
*
* @param executeMethod executeMethod
* @param remotePeer remotePeer
* @param operation operation
* @param remotePeer remotePeer
* @param operation operation
*/
public static void createExitSpan(String executeMethod, String remotePeer, Object operation) {
AbstractSpan span = ContextManager.createExitSpan(
Expand All @@ -50,13 +51,15 @@ public static void createExitSpan(String executeMethod, String remotePeer, Objec
SpanLayer.asDB(span);

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

Expand Down
62 changes: 47 additions & 15 deletions .../org/apache/skywalking/apm/plugin/mongodb/v4/MongoDBOperationExecutorInterceptorTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,15 +18,13 @@

package org.apache.skywalking.apm.plugin.mongodb.v4;

import static org.hamcrest.CoreMatchers.is;
import static org.junit.Assert.assertThat;
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.MongoNamespace;
import com.mongodb.ReadConcern;
import com.mongodb.client.internal.OperationExecutor;
import com.mongodb.internal.operation.AggregateOperation;
import com.mongodb.internal.operation.CreateCollectionOperation;
import com.mongodb.internal.operation.FindOperation;
import com.mongodb.internal.operation.WriteOperation;
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 @@ -58,11 +56,16 @@
import org.mockito.Spy;
import org.mockito.junit.MockitoJUnit;
import org.mockito.junit.MockitoRule;
import com.mongodb.MongoNamespace;
import com.mongodb.ReadConcern;
import com.mongodb.client.internal.OperationExecutor;
import com.mongodb.internal.operation.FindOperation;
import com.mongodb.internal.operation.WriteOperation;

import java.lang.reflect.Method;
import java.util.Collections;
import java.util.List;

import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.CoreMatchers.startsWith;
import static org.junit.Assert.assertThat;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

@RunWith(TracingSegmentRunner.class)
public class MongoDBOperationExecutorInterceptorTest {
Expand Down Expand Up @@ -136,6 +139,24 @@ public void testConstructIntercept() throws Throwable {
MatcherAssert.assertThat(enhancedObjInstance.getSkyWalkingDynamicField(), Is.is(mongoNamespace));
}

@Test
public void testCreateCollectionOperationIntercept() throws Throwable {
CreateCollectionOperation createCollectionOperation = new CreateCollectionOperation("test", "user");
CreateCollectionOperation enhancedInstanceForCreateCollectionOperation = mock(CreateCollectionOperation.class, Mockito.withSettings().extraInterfaces(EnhancedInstance.class));
when(((EnhancedInstance) enhancedInstanceForCreateCollectionOperation).getSkyWalkingDynamicField()).thenReturn("test");
when(enhancedInstanceForCreateCollectionOperation.getCollectionName()).thenReturn("user");

Object[] arguments = {enhancedInstanceForCreateCollectionOperation};
Class[] argumentTypes = {createCollectionOperation.getClass()};
interceptor.beforeMethod(enhancedInstance, getMethod(), arguments, argumentTypes, null);
interceptor.afterMethod(enhancedInstance, getMethod(), arguments, argumentTypes, null);

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

@Test
public void testIntercept() throws Throwable {
interceptor.beforeMethod(enhancedInstance, getMethod(), arguments, argumentTypes, null);
Expand Down Expand Up @@ -185,7 +206,7 @@ public void testInterceptWithException() throws Throwable {
}

private void assertMongoFindOperationSpan(AbstractTracingSpan span) {
assertThat(span.getOperationName(), is("MongoDB/FindOperation"));
assertThat(span.getOperationName(), startsWith("MongoDB/FindOperation"));
assertThat(SpanHelper.getComponentId(span), is(42));
List<TagValuePair> tags = SpanHelper.getTags(span);
assertThat(tags.get(0).getValue(), is("MongoDB"));
Expand All @@ -196,8 +217,19 @@ private void assertMongoFindOperationSpan(AbstractTracingSpan span) {
assertThat(SpanHelper.getLayer(span), CoreMatchers.is(SpanLayer.DB));
}

private void assertMongoCreateCollectionOperationSpan(AbstractTracingSpan span) {
assertThat(span.getOperationName(), startsWith("MongoDB/CreateCollectionOperation"));
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(span.isExit(), is(true));
assertThat(SpanHelper.getLayer(span), CoreMatchers.is(SpanLayer.DB));
}

private void assertMongoAggregateOperationSpan(AbstractTracingSpan span) {
assertThat(span.getOperationName(), is("MongoDB/AggregateOperation"));
assertThat(span.getOperationName(), startsWith("MongoDB/AggregateOperation"));
assertThat(SpanHelper.getComponentId(span), is(42));
List<TagValuePair> tags = SpanHelper.getTags(span);
assertThat(tags.get(0).getValue(), is("MongoDB"));
Expand Down

0 comments on commit 3e64803

Please sign in to comment.