Skip to content

Commit e329aec

Browse files
AhmadMasryAhmadMasry
authored
Opensearch support (#1)
Add support for Amazon Opensearch instrumentation.
1 parent e725829 commit e329aec

13 files changed

Lines changed: 603 additions & 0 deletions

File tree

‎dd-java-agent/instrumentation/opensearch/build.gradle‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
apply from: "$rootDir/gradle/java.gradle"
2+
3+
dependencies {
4+
compileOnly group: 'org.opensearch.client', name: 'opensearch-rest-client', version: '1.0.0'
5+
compileOnly group: 'org.opensearch', name: 'opensearch', version: '1.0.0'
6+
}

‎dd-java-agent/instrumentation/opensearch/rest/build.gradle‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
2+
muzzle {
3+
pass {
4+
group = "org.opensearch.client"
5+
module = "opensearch-rest-client"
6+
versions = "[1,)"
7+
assertInverse = true
8+
}
9+
}
10+
11+
apply from: "$rootDir/gradle/java.gradle"
12+
13+
addTestSuiteForDir('latestDepTest', 'test')
14+
15+
dependencies {
16+
compileOnly group: 'org.opensearch.client', name: 'opensearch-rest-client', version: '1.0.0'
17+
18+
implementation project(':dd-java-agent:instrumentation:opensearch')
19+
20+
testImplementation project(':dd-java-agent:instrumentation:apache-httpclient-4')
21+
testImplementation project(':dd-java-agent:instrumentation:apache-httpasyncclient-4')
22+
// Netty is used, but it adds complexity to the tests since we're using embedded ES.
23+
//testImplementation project(':dd-java-agent:instrumentation:netty-4.1')
24+
25+
testImplementation group: 'org.apache.logging.log4j', name: 'log4j-core', version: '2.11.0'
26+
testImplementation group: 'org.apache.logging.log4j', name: 'log4j-api', version: '2.11.0'
27+
28+
testImplementation group: 'org.opensearch.client', name: 'opensearch-rest-client', version: '1.0.0'
29+
testImplementation group: 'org.opensearch', name: 'opensearch', version: '1.0.0'
30+
testImplementation group: 'org.opensearch.plugin', name: 'transport-netty4-client', version: '1.0.0'
31+
32+
latestDepTestImplementation group: 'org.opensearch.client', name: 'opensearch-rest-client', version: '1.0.+'
33+
latestDepTestImplementation group: 'org.opensearch.client', name: 'transport', version: '1.0.+'
34+
latestDepTestImplementation group: 'org.opensearch', name: 'opensearch', version: '1.0.+'
35+
latestDepTestImplementation group: 'org.opensearch.plugin', name: 'transport-netty4-client', version: '1.0.+'
36+
}

‎dd-java-agent/instrumentation/opensearch/rest/src/main/java/datadog/trace/instrumentation/opensearch/OpensearchRestClientInstrumentation.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
package datadog.trace.instrumentation.opensearch;
2+
3+
import static datadog.trace.agent.tooling.bytebuddy.matcher.ClassLoaderMatchers.hasClassNamed;
4+
import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
5+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
6+
import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
7+
import static datadog.trace.instrumentation.opensearch.OpensearchRestClientDecorator.DECORATE;
8+
import static datadog.trace.instrumentation.opensearch.OpensearchRestClientDecorator.OPERATION_NAME;
9+
import static net.bytebuddy.matcher.ElementMatchers.isMethod;
10+
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
11+
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
12+
13+
import com.google.auto.service.AutoService;
14+
import datadog.trace.agent.tooling.Instrumenter;
15+
import datadog.trace.bootstrap.instrumentation.api.AgentScope;
16+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
17+
import net.bytebuddy.asm.Advice;
18+
import net.bytebuddy.implementation.bytecode.assign.Assigner;
19+
import net.bytebuddy.matcher.ElementMatcher;
20+
import org.opensearch.client.Request;
21+
import org.opensearch.client.Response;
22+
import org.opensearch.client.ResponseListener;
23+
24+
@AutoService(Instrumenter.class)
25+
public class OpensearchRestClientInstrumentation extends Instrumenter.Tracing
26+
implements Instrumenter.ForSingleType {
27+
28+
public OpensearchRestClientInstrumentation() {
29+
super("opensearch", "opensearch-rest");
30+
}
31+
32+
@Override
33+
public ElementMatcher<ClassLoader> classLoaderMatcher() {
34+
// Avoid matching pre-ES7 releases which have their own instrumentations.
35+
return hasClassNamed("org.opensearch.client.RestClient$InternalRequest");
36+
}
37+
38+
@Override
39+
public String[] helperClassNames() {
40+
return new String[] {
41+
"datadog.trace.instrumentation.opensearch.OpensearchRestClientDecorator",
42+
packageName + ".RestResponseListener",
43+
};
44+
}
45+
46+
@Override
47+
public String instrumentedType() {
48+
return "org.opensearch.client.RestClient";
49+
}
50+
51+
@Override
52+
public void adviceTransformations(AdviceTransformation transformation) {
53+
transformation.applyAdvice(
54+
isMethod()
55+
.and(named("performRequest"))
56+
.and(takesArguments(1))
57+
.and(takesArgument(0, named("org.opensearch.client.Request"))),
58+
OpensearchRestClientInstrumentation.class.getName() + "$OpensearchRestClientAdvice");
59+
transformation.applyAdvice(
60+
isMethod()
61+
.and(named("performRequestAsync"))
62+
.and(takesArguments(2))
63+
.and(takesArgument(0, named("org.opensearch.client.Request")))
64+
.and(takesArgument(1, named("org.opensearch.client.ResponseListener"))),
65+
OpensearchRestClientInstrumentation.class.getName() + "$OpensearchRestClientAdvice");
66+
}
67+
68+
public static class OpensearchRestClientAdvice {
69+
70+
@Advice.OnMethodEnter(suppress = Throwable.class)
71+
public static AgentScope onEnter(
72+
@Advice.Argument(0) final Request request,
73+
@Advice.Argument(value = 1, readOnly = false, optional = true)
74+
ResponseListener responseListener) {
75+
76+
final AgentSpan span = startSpan(OPERATION_NAME);
77+
DECORATE.afterStart(span);
78+
DECORATE.onRequest(span, request.getMethod(), request.getEndpoint());
79+
80+
if (responseListener != null) {
81+
responseListener = new RestResponseListener(responseListener, span);
82+
}
83+
84+
return activateSpan(span);
85+
}
86+
87+
@Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
88+
public static void stopSpan(
89+
@Advice.Enter final AgentScope scope,
90+
@Advice.Thrown final Throwable throwable,
91+
@Advice.Return(typing = Assigner.Typing.DYNAMIC) final Object result) {
92+
if (throwable != null) {
93+
final AgentSpan span = scope.span();
94+
DECORATE.onError(span, throwable);
95+
DECORATE.beforeFinish(span);
96+
span.finish();
97+
} else if (result instanceof Response) {
98+
final AgentSpan span = scope.span();
99+
if (((Response) result).getHost() != null) {
100+
DECORATE.onResponse(span, ((Response) result));
101+
}
102+
DECORATE.beforeFinish(span);
103+
span.finish();
104+
} else {
105+
// async call, span finished by RestResponseListener
106+
}
107+
scope.close();
108+
}
109+
}
110+
}

‎dd-java-agent/instrumentation/opensearch/rest/src/main/java/datadog/trace/instrumentation/opensearch/RestResponseListener.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
package datadog.trace.instrumentation.opensearch;
2+
3+
import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
4+
import org.opensearch.client.Response;
5+
import org.opensearch.client.ResponseListener;
6+
7+
/** This class is identical to version 6's instrumentation. */
8+
public class RestResponseListener implements ResponseListener {
9+
10+
private final ResponseListener listener;
11+
private final AgentSpan span;
12+
13+
public RestResponseListener(final ResponseListener listener, final AgentSpan span) {
14+
this.listener = listener;
15+
this.span = span;
16+
}
17+
18+
@Override
19+
public void onSuccess(final Response response) {
20+
if (response.getHost() != null) {
21+
OpensearchRestClientDecorator.DECORATE.onResponse(span, response);
22+
}
23+
24+
try {
25+
listener.onSuccess(response);
26+
} finally {
27+
OpensearchRestClientDecorator.DECORATE.beforeFinish(span);
28+
span.finish();
29+
}
30+
}
31+
32+
@Override
33+
public void onFailure(final Exception e) {
34+
OpensearchRestClientDecorator.DECORATE.onError(span, e);
35+
36+
try {
37+
listener.onFailure(e);
38+
} finally {
39+
OpensearchRestClientDecorator.DECORATE.beforeFinish(span);
40+
span.finish();
41+
}
42+
}
43+
}

‎dd-java-agent/instrumentation/opensearch/rest/src/test/groovy/OpensearchRestClientTest.groovy‎

Lines changed: 120 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,120 @@
1+
import datadog.trace.agent.test.AgentTestRunner
2+
import datadog.trace.api.DDSpanTypes
3+
import datadog.trace.bootstrap.instrumentation.api.Tags
4+
import groovy.json.JsonSlurper
5+
import org.apache.http.HttpHost
6+
import org.apache.http.client.config.RequestConfig
7+
import org.apache.http.util.EntityUtils
8+
import org.opensearch.client.Request
9+
import org.opensearch.client.Response
10+
import org.opensearch.client.RestClient
11+
import org.opensearch.client.RestClientBuilder
12+
import org.opensearch.common.io.FileSystemUtils
13+
import org.opensearch.common.settings.Settings
14+
import org.opensearch.common.transport.TransportAddress
15+
import org.opensearch.http.HttpServerTransport
16+
import org.opensearch.node.InternalSettingsPreparer
17+
import org.opensearch.node.Node
18+
import org.opensearch.transport.Netty4Plugin
19+
import spock.lang.Shared
20+
21+
class OpensearchRestClientTest extends AgentTestRunner {
22+
@Shared
23+
TransportAddress httpTransportAddress
24+
@Shared
25+
Node testNode
26+
@Shared
27+
File esWorkingDir
28+
@Shared
29+
String clusterName = UUID.randomUUID().toString()
30+
31+
@Shared
32+
RestClient client
33+
34+
def setupSpec() {
35+
36+
esWorkingDir = File.createTempDir("test-es-working-dir-", "")
37+
esWorkingDir.deleteOnExit()
38+
println "ES work dir: $esWorkingDir"
39+
40+
def settings = Settings.builder()
41+
.put("path.home", esWorkingDir.path)
42+
.put("cluster.name", clusterName)
43+
.put("node.name", "test-node")
44+
.put("transport.type", "netty4")
45+
.build()
46+
testNode = new Node(InternalSettingsPreparer.prepareEnvironment(
47+
settings, [:], null, null), [Netty4Plugin], false) {}
48+
testNode.start()
49+
httpTransportAddress = testNode.injector().getInstance(HttpServerTransport).boundAddress().publishAddress()
50+
51+
client = RestClient.builder(new HttpHost(httpTransportAddress.address, httpTransportAddress.port))
52+
.setRequestConfigCallback(new RestClientBuilder.RequestConfigCallback() {
53+
@Override
54+
RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder builder) {
55+
return builder.setConnectTimeout(Integer.MAX_VALUE).setSocketTimeout(Integer.MAX_VALUE)
56+
}
57+
})
58+
.build()
59+
}
60+
61+
def cleanupSpec() {
62+
testNode?.close()
63+
if (esWorkingDir != null) {
64+
FileSystemUtils.deleteSubDirectories(esWorkingDir.toPath())
65+
esWorkingDir.delete()
66+
}
67+
}
68+
69+
def "test opensearch status #nr"() {
70+
setup:
71+
Request request = new Request("GET", "_cluster/health")
72+
Response response = client.performRequest(request)
73+
74+
Map result = new JsonSlurper().parseText(EntityUtils.toString(response.entity))
75+
76+
expect:
77+
result.status == "green"
78+
79+
assertTraces(1) {
80+
sortSpansByStart()
81+
trace(2) {
82+
span {
83+
serviceName "opensearch"
84+
resourceName "GET _cluster/health"
85+
operationName "opensearch.rest.query"
86+
spanType DDSpanTypes.OPENSEARCH
87+
parent()
88+
tags {
89+
"$Tags.COMPONENT" "opensearch-java"
90+
"$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
91+
"$Tags.PEER_HOSTNAME" httpTransportAddress.address
92+
"$Tags.PEER_PORT" httpTransportAddress.port
93+
"$Tags.HTTP_URL" "_cluster/health"
94+
"$Tags.HTTP_METHOD" "GET"
95+
"$Tags.DB_TYPE" "opensearch"
96+
defaultTags()
97+
}
98+
}
99+
span {
100+
serviceName "opensearch"
101+
resourceName "GET _cluster/health"
102+
operationName "http.request"
103+
spanType DDSpanTypes.HTTP_CLIENT
104+
childOf span(0)
105+
tags {
106+
"$Tags.COMPONENT" "apache-httpasyncclient"
107+
"$Tags.SPAN_KIND" Tags.SPAN_KIND_CLIENT
108+
"$Tags.HTTP_URL" "_cluster/health"
109+
"$Tags.HTTP_METHOD" "GET"
110+
"$Tags.HTTP_STATUS" 200
111+
defaultTags()
112+
}
113+
}
114+
}
115+
}
116+
117+
where:
118+
nr << (1..101)
119+
}
120+
}

0 commit comments

Comments
 (0)