Skip to content

Commit b54e8b3

Browse files
pavolloffaypavolloffay
authored
Java streams instrumentation tests (#188)
Signed-off-by: Pavol Loffay <[email protected]>
1 parent fd1270d commit b54e8b3

File tree

8 files changed

+353
-42
lines changed

8 files changed

+353
-42
lines changed

‎instrumentation/java-streams/build.gradle.kts‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,14 @@
11
plugins {
22
`java-library`
33
id("net.bytebuddy.byte-buddy")
4+
id("io.opentelemetry.instrumentation.auto-instrumentation")
5+
muzzle
6+
}
7+
8+
muzzle {
9+
pass {
10+
coreJdk()
11+
}
412
}
513

614
afterEvaluate{
@@ -10,3 +18,7 @@ afterEvaluate{
1018
project(":javaagent-tooling").configurations["instrumentationMuzzle"] + configurations.runtimeClasspath
1119
).configure()
1220
}
21+
22+
dependencies {
23+
testImplementation(project(":testing-common"))
24+
}

‎instrumentation/java-streams/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/java/inputstream/InputStreamInstrumentationModule.java‎

Lines changed: 57 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
import net.bytebuddy.description.type.TypeDescription;
3838
import net.bytebuddy.matcher.ElementMatcher;
3939
import org.hypertrace.agent.core.GlobalObjectRegistry;
40+
import org.hypertrace.agent.core.GlobalObjectRegistry.SpanAndBuffer;
4041

4142
/**
4243
* {@link InputStream} instrumentation. The type matcher applies to all implementations. However
@@ -67,7 +68,7 @@ static class InputStreamInstrumentation implements TypeInstrumentation {
6768

6869
@Override
6970
public ElementMatcher<? super TypeDescription> typeMatcher() {
70-
return extendsClass(named("java.io.InputStream"));
71+
return extendsClass(named(InputStream.class.getName()));
7172
}
7273

7374
@Override
@@ -93,7 +94,7 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
9394
+ "$InputStream_ReadByteArrayOffsetAdvice");
9495
transformers.put(
9596
named("readAllBytes").and(takesArguments(0)).and(isPublic()),
96-
InputStream_ReadAllBytes.class.getName());
97+
InputStreamInstrumentationModule.class.getName() + "$InputStream_ReadAllBytes");
9798
transformers.put(
9899
named("readNBytes")
99100
.and(takesArguments(0))
@@ -107,50 +108,93 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
107108
}
108109

109110
public static class InputStream_ReadNoArgsAdvice {
111+
@Advice.OnMethodEnter(suppress = Throwable.class)
112+
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
113+
return InputStreamUtils.check(thizz);
114+
}
115+
110116
@Advice.OnMethodExit(suppress = Throwable.class)
111-
public static void exit(@Advice.This InputStream thizz, @Advice.Return int read) {
112-
InputStreamUtils.read(thizz, read);
117+
public static void exit(
118+
@Advice.This InputStream thizz,
119+
@Advice.Return int read,
120+
@Advice.Enter SpanAndBuffer spanAndBuffer) {
121+
if (spanAndBuffer != null) {
122+
InputStreamUtils.read(thizz, spanAndBuffer, read);
123+
}
113124
}
114125
}
115126

116127
public static class InputStream_ReadByteArrayAdvice {
128+
@Advice.OnMethodEnter(suppress = Throwable.class)
129+
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
130+
return InputStreamUtils.check(thizz);
131+
}
132+
117133
@Advice.OnMethodExit(suppress = Throwable.class)
118134
public static void exit(
119-
@Advice.This InputStream thizz, @Advice.Return int read, @Advice.Argument(0) byte b[]) {
120-
InputStreamUtils.read(thizz, read, b);
135+
@Advice.This InputStream thizz,
136+
@Advice.Return int read,
137+
@Advice.Argument(0) byte b[],
138+
@Advice.Enter SpanAndBuffer spanAndBuffer) {
139+
if (spanAndBuffer != null) {
140+
InputStreamUtils.read(thizz, spanAndBuffer, read, b);
141+
}
121142
}
122143
}
123144

124145
public static class InputStream_ReadByteArrayOffsetAdvice {
146+
@Advice.OnMethodEnter(suppress = Throwable.class)
147+
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
148+
return InputStreamUtils.check(thizz);
149+
}
150+
125151
@Advice.OnMethodExit(suppress = Throwable.class)
126152
public static void exit(
127153
@Advice.This InputStream thizz,
128154
@Advice.Return int read,
129155
@Advice.Argument(0) byte b[],
130156
@Advice.Argument(1) int off,
131-
@Advice.Argument(2) int len) {
132-
InputStreamUtils.read(thizz, read, b, off, len);
157+
@Advice.Argument(2) int len,
158+
@Advice.Enter SpanAndBuffer spanAndBuffer) {
159+
if (spanAndBuffer != null) {
160+
InputStreamUtils.read(thizz, spanAndBuffer, read, b, off, len);
161+
}
133162
}
134163
}
135164

136165
public static class InputStream_ReadAllBytes {
166+
@Advice.OnMethodEnter(suppress = Throwable.class)
167+
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
168+
return InputStreamUtils.check(thizz);
169+
}
170+
137171
@Advice.OnMethodExit(suppress = Throwable.class)
138-
public static void exit(@Advice.This InputStream thizz, @Advice.Return byte[] b)
172+
public static void exit(
173+
@Advice.This InputStream thizz,
174+
@Advice.Return byte[] b,
175+
@Advice.Enter SpanAndBuffer spanAndBuffer)
139176
throws IOException {
140-
InputStreamUtils.readAll(thizz, b);
177+
if (spanAndBuffer != null) {
178+
InputStreamUtils.readAll(thizz, spanAndBuffer, b);
179+
}
141180
}
142181
}
143182

144183
public static class InputStream_ReadNBytes {
184+
@Advice.OnMethodEnter(suppress = Throwable.class)
185+
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
186+
return InputStreamUtils.check(thizz);
187+
}
188+
145189
@Advice.OnMethodExit(suppress = Throwable.class)
146190
public static void exit(
147191
@Advice.This InputStream thizz,
148192
@Advice.Return int read,
149193
@Advice.Argument(0) byte[] b,
150194
@Advice.Argument(1) int off,
151-
@Advice.Argument(2) int len)
152-
throws IOException {
153-
InputStreamUtils.readNBytes(thizz, read, b, off, len);
195+
@Advice.Argument(2) int len,
196+
@Advice.Enter SpanAndBuffer spanAndBuffer) {
197+
InputStreamUtils.readNBytes(thizz, spanAndBuffer, read, b, off, len);
154198
}
155199
}
156200
}

‎instrumentation/java-streams/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/java/inputstream/InputStreamUtils.java‎

Lines changed: 25 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import io.opentelemetry.api.trace.Span;
2222
import io.opentelemetry.api.trace.Tracer;
2323
import io.opentelemetry.context.Context;
24+
import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
2425
import java.io.ByteArrayOutputStream;
2526
import java.io.IOException;
2627
import java.io.InputStream;
@@ -69,12 +70,21 @@ public static void addBody(
6970
}
7071
}
7172

72-
public static void read(InputStream inputStream, int read) {
73+
public static SpanAndBuffer check(InputStream inputStream) {
7374
SpanAndBuffer spanAndBuffer =
7475
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
7576
if (spanAndBuffer == null) {
76-
return;
77+
return null;
7778
}
79+
80+
int callDepth = CallDepthThreadLocalMap.incrementCallDepth(InputStream.class);
81+
if (callDepth > 0) {
82+
return null;
83+
}
84+
return spanAndBuffer;
85+
}
86+
87+
public static void read(InputStream inputStream, SpanAndBuffer spanAndBuffer, int read) {
7888
if (read != -1) {
7989
spanAndBuffer.byteArrayBuffer.write((byte) read);
8090
} else if (read == -1) {
@@ -85,14 +95,11 @@ public static void read(InputStream inputStream, int read) {
8595
spanAndBuffer.charset);
8696
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
8797
}
98+
CallDepthThreadLocalMap.reset(InputStream.class);
8899
}
89100

90-
public static void read(InputStream inputStream, int read, byte[] b) {
91-
SpanAndBuffer spanAndBuffer =
92-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
93-
if (spanAndBuffer == null) {
94-
return;
95-
}
101+
public static void read(
102+
InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b) {
96103
if (read > 0) {
97104
spanAndBuffer.byteArrayBuffer.write(b, 0, read);
98105
} else if (read == -1) {
@@ -103,14 +110,11 @@ public static void read(InputStream inputStream, int read, byte[] b) {
103110
spanAndBuffer.charset);
104111
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
105112
}
113+
CallDepthThreadLocalMap.reset(InputStream.class);
106114
}
107115

108-
public static void read(InputStream inputStream, int read, byte[] b, int off, int len) {
109-
SpanAndBuffer spanAndBuffer =
110-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
111-
if (spanAndBuffer == null) {
112-
return;
113-
}
116+
public static void read(
117+
InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b, int off, int len) {
114118
if (read > 0) {
115119
spanAndBuffer.byteArrayBuffer.write(b, off, read);
116120
} else if (read == -1) {
@@ -121,24 +125,18 @@ public static void read(InputStream inputStream, int read, byte[] b, int off, in
121125
spanAndBuffer.charset);
122126
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
123127
}
128+
CallDepthThreadLocalMap.reset(InputStream.class);
124129
}
125130

126-
public static void readAll(InputStream inputStream, byte[] b) throws IOException {
127-
SpanAndBuffer spanAndBuffer =
128-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
129-
if (spanAndBuffer == null) {
130-
return;
131-
}
131+
public static void readAll(InputStream inputStream, SpanAndBuffer spanAndBuffer, byte[] b)
132+
throws IOException {
132133
spanAndBuffer.byteArrayBuffer.write(b);
133134
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
135+
CallDepthThreadLocalMap.reset(InputStream.class);
134136
}
135137

136-
public static void readNBytes(InputStream inputStream, int read, byte[] b, int off, int len) {
137-
SpanAndBuffer spanAndBuffer =
138-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
139-
if (spanAndBuffer == null) {
140-
return;
141-
}
138+
public static void readNBytes(
139+
InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b, int off, int len) {
142140
if (read == 0) {
143141
InputStreamUtils.addBody(
144142
spanAndBuffer.span,
@@ -149,5 +147,6 @@ public static void readNBytes(InputStream inputStream, int read, byte[] b, int o
149147
} else {
150148
spanAndBuffer.byteArrayBuffer.write(b, off, read);
151149
}
150+
CallDepthThreadLocalMap.reset(InputStream.class);
152151
}
153152
}

‎instrumentation/java-streams/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/java/outputstream/OutputStreamInstrumentationModule.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -63,7 +63,7 @@ static class OutputStreamInstrumentation implements TypeInstrumentation {
6363

6464
@Override
6565
public ElementMatcher<? super TypeDescription> typeMatcher() {
66-
return extendsClass(named("java.io.OutputStream"));
66+
return extendsClass(named(OutputStream.class.getName()));
6767
}
6868

6969
@Override

0 commit comments

Comments
 (0)