Skip to content

Commit 0dc7b00

Browse files
pavolloffaypavolloffay
authored
Remove global registry and use instrumentation context (#318)
Signed-off-by: Pavol Loffay <[email protected]>
1 parent 9c996de commit 0dc7b00

File tree

13 files changed

+324
-134
lines changed

13 files changed

+324
-134
lines changed

‎instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/HttpEntityInstrumentation.java‎

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
2424
import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
2525

26+
import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
27+
import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
2628
import io.opentelemetry.javaagent.instrumentation.hypertrace.apachehttpclient.v4_0.ApacheHttpClientObjectRegistry.SpanAndAttributeKey;
2729
import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
2830
import java.io.InputStream;
@@ -37,8 +39,7 @@
3739
import net.bytebuddy.matcher.ElementMatcher;
3840
import org.apache.http.Header;
3941
import org.apache.http.HttpEntity;
40-
import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
41-
import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry.SpanAndBuffer;
42+
import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
4243
import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
4344
import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
4445
import org.hypertrace.agent.core.instrumentation.utils.ContentLengthUtils;
@@ -103,7 +104,8 @@ public static void exit(@Advice.This HttpEntity thizz, @Advice.Return InputStrea
103104
BoundedBuffersFactory.createStream((int) contentSize, charset),
104105
clientSpan.attributeKey,
105106
charset);
106-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.put(inputStream, spanAndBuffer);
107+
InstrumentationContext.get(InputStream.class, SpanAndBuffer.class)
108+
.put(inputStream, spanAndBuffer);
107109
}
108110
}
109111

@@ -131,7 +133,8 @@ public static void enter(
131133
BoundedByteArrayOutputStream byteArrayOutputStream =
132134
BoundedBuffersFactory.createStream((int) contentSize, charset);
133135

134-
GlobalObjectRegistry.outputStreamToBufferMap.put(outputStream, byteArrayOutputStream);
136+
InstrumentationContext.get(OutputStream.class, BoundedByteArrayOutputStream.class)
137+
.put(outputStream, byteArrayOutputStream);
135138
}
136139

137140
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -144,9 +147,10 @@ public static void exit(
144147
return;
145148
}
146149

147-
BoundedByteArrayOutputStream bufferedOutStream =
148-
GlobalObjectRegistry.outputStreamToBufferMap.get(outputStream);
149-
GlobalObjectRegistry.outputStreamToBufferMap.remove(outputStream);
150+
ContextStore<OutputStream, BoundedByteArrayOutputStream> contextStore =
151+
InstrumentationContext.get(OutputStream.class, BoundedByteArrayOutputStream.class);
152+
BoundedByteArrayOutputStream bufferedOutStream = contextStore.get(outputStream);
153+
contextStore.put(outputStream, null);
150154
try {
151155
String requestBody = bufferedOutStream.toStringWithSuppliedCharset();
152156
spanAndAttributeKey.span.setAttribute(spanAndAttributeKey.attributeKey, requestBody);

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

Lines changed: 60 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,8 @@
2727

2828
import com.google.auto.service.AutoService;
2929
import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
30+
import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
31+
import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
3032
import io.opentelemetry.javaagent.tooling.InstrumentationModule;
3133
import io.opentelemetry.javaagent.tooling.TypeInstrumentation;
3234
import java.io.IOException;
@@ -39,18 +41,17 @@
3941
import net.bytebuddy.description.method.MethodDescription;
4042
import net.bytebuddy.description.type.TypeDescription;
4143
import net.bytebuddy.matcher.ElementMatcher;
42-
import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
43-
import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry.SpanAndBuffer;
44+
import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
4445

4546
/**
4647
* {@link InputStream} instrumentation. The type matcher applies to all implementations. However
47-
* only streams that are in the {@link GlobalObjectRegistry#inputStreamToSpanAndBufferMap} are
48-
* instrumented, otherwise the instrumentation is noop.
48+
* only streams that are in the {@link ContextStore} are instrumented, otherwise the instrumentation
49+
* is noop.
4950
*
50-
* <p>If the stream is in the {@link GlobalObjectRegistry#inputStreamToSpanAndBufferMap} then result
51-
* of read methods is also passed to the buffered stream (value) from the map. The instrumentation
52-
* adds buffer to span from the map when read is finished (return -1), creates new span with buffer
53-
* when the original span is not recording.
51+
* <p>If the stream is in the {@link ContextStore} then result of read methods is also passed to the
52+
* buffered stream (value) from the map. The instrumentation adds buffer to span from the map when
53+
* read is finished (return -1), creates new span with buffer when the original span is not
54+
* recording.
5455
*
5556
* <p>Maybe we could add optimization to instrument the input streams only when certain classes are
5657
* present in classloader e.g. classes from frameworks that we instrument.
@@ -117,7 +118,8 @@ public Map<? extends ElementMatcher<? super MethodDescription>, String> transfor
117118
public static class InputStream_ReadNoArgsAdvice {
118119
@Advice.OnMethodEnter(suppress = Throwable.class)
119120
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
120-
return InputStreamUtils.check(thizz);
121+
return InputStreamUtils.check(
122+
thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
121123
}
122124

123125
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -133,14 +135,19 @@ public static void exit(
133135
return;
134136
}
135137

136-
InputStreamUtils.read(thizz, spanAndBuffer, read);
138+
InputStreamUtils.read(
139+
thizz,
140+
spanAndBuffer,
141+
InstrumentationContext.get(InputStream.class, SpanAndBuffer.class),
142+
read);
137143
}
138144
}
139145

140146
public static class InputStream_ReadByteArrayAdvice {
141147
@Advice.OnMethodEnter(suppress = Throwable.class)
142148
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
143-
return InputStreamUtils.check(thizz);
149+
return InputStreamUtils.check(
150+
thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
144151
}
145152

146153
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -164,7 +171,8 @@ public static void exit(
164171
public static class InputStream_ReadByteArrayOffsetAdvice {
165172
@Advice.OnMethodEnter(suppress = Throwable.class)
166173
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
167-
return InputStreamUtils.check(thizz);
174+
return InputStreamUtils.check(
175+
thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
168176
}
169177

170178
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -183,14 +191,22 @@ public static void exit(
183191
return;
184192
}
185193

186-
InputStreamUtils.read(thizz, spanAndBuffer, read, b, off, len);
194+
InputStreamUtils.read(
195+
thizz,
196+
spanAndBuffer,
197+
InstrumentationContext.get(InputStream.class, SpanAndBuffer.class),
198+
read,
199+
b,
200+
off,
201+
len);
187202
}
188203
}
189204

190205
public static class InputStream_ReadAllBytes {
191206
@Advice.OnMethodEnter(suppress = Throwable.class)
192207
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
193-
return InputStreamUtils.check(thizz);
208+
return InputStreamUtils.check(
209+
thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
194210
}
195211

196212
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -207,14 +223,19 @@ public static void exit(
207223
return;
208224
}
209225

210-
InputStreamUtils.readAll(thizz, spanAndBuffer, b);
226+
InputStreamUtils.readAll(
227+
thizz,
228+
spanAndBuffer,
229+
InstrumentationContext.get(InputStream.class, SpanAndBuffer.class),
230+
b);
211231
}
212232
}
213233

214234
public static class InputStream_ReadNBytes {
215235
@Advice.OnMethodEnter(suppress = Throwable.class)
216236
public static SpanAndBuffer enter(@Advice.This InputStream thizz) {
217-
return InputStreamUtils.check(thizz);
237+
return InputStreamUtils.check(
238+
thizz, InstrumentationContext.get(InputStream.class, SpanAndBuffer.class));
218239
}
219240

220241
@Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
@@ -232,14 +253,35 @@ public static void exit(
232253
if (callDepth > 0) {
233254
return;
234255
}
235-
InputStreamUtils.readNBytes(thizz, spanAndBuffer, read, b, off, len);
256+
InputStreamUtils.readNBytes(
257+
thizz,
258+
spanAndBuffer,
259+
InstrumentationContext.get(InputStream.class, SpanAndBuffer.class),
260+
read,
261+
b,
262+
off,
263+
len);
236264
}
237265
}
238266

239267
public static class InputStream_Available {
240268
@Advice.OnMethodExit(suppress = Throwable.class)
241269
public static void exit(@Advice.This InputStream thizz, @Advice.Return int available) {
242-
InputStreamUtils.available(thizz, available);
270+
if (available != 0) {
271+
return;
272+
}
273+
ContextStore<InputStream, SpanAndBuffer> contextStore =
274+
InstrumentationContext.get(InputStream.class, SpanAndBuffer.class);
275+
276+
SpanAndBuffer spanAndBuffer = contextStore.get(thizz);
277+
if (spanAndBuffer != null) {
278+
InputStreamUtils.addBody(
279+
spanAndBuffer.span,
280+
spanAndBuffer.attributeKey,
281+
spanAndBuffer.byteArrayBuffer,
282+
spanAndBuffer.charset);
283+
contextStore.put(thizz, null);
284+
}
243285
}
244286
}
245287
}

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

Lines changed: 35 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -22,14 +22,15 @@
2222
import io.opentelemetry.api.trace.Tracer;
2323
import io.opentelemetry.context.Context;
2424
import io.opentelemetry.javaagent.instrumentation.api.CallDepthThreadLocalMap;
25+
import io.opentelemetry.javaagent.instrumentation.api.ContextStore;
26+
import io.opentelemetry.javaagent.instrumentation.api.InstrumentationContext;
2527
import java.io.ByteArrayOutputStream;
2628
import java.io.IOException;
2729
import java.io.InputStream;
2830
import java.io.UnsupportedEncodingException;
2931
import java.nio.charset.Charset;
30-
import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
31-
import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry.SpanAndBuffer;
3232
import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
33+
import org.hypertrace.agent.core.instrumentation.SpanAndBuffer;
3334
import org.slf4j.Logger;
3435
import org.slf4j.LoggerFactory;
3536

@@ -69,9 +70,9 @@ public static void addBody(
6970
}
7071
}
7172

72-
public static SpanAndBuffer check(InputStream inputStream) {
73-
SpanAndBuffer spanAndBuffer =
74-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
73+
public static SpanAndBuffer check(
74+
InputStream inputStream, ContextStore<InputStream, SpanAndBuffer> contextStore) {
75+
SpanAndBuffer spanAndBuffer = contextStore.get(inputStream);
7576
if (spanAndBuffer == null) {
7677
return null;
7778
}
@@ -80,7 +81,11 @@ public static SpanAndBuffer check(InputStream inputStream) {
8081
return spanAndBuffer;
8182
}
8283

83-
public static void read(InputStream inputStream, SpanAndBuffer spanAndBuffer, int read) {
84+
public static void read(
85+
InputStream inputStream,
86+
SpanAndBuffer spanAndBuffer,
87+
ContextStore<InputStream, SpanAndBuffer> contextStore,
88+
int read) {
8489
if (read != -1) {
8590
spanAndBuffer.byteArrayBuffer.write((byte) read);
8691
} else if (read == -1) {
@@ -89,7 +94,7 @@ public static void read(InputStream inputStream, SpanAndBuffer spanAndBuffer, in
8994
spanAndBuffer.attributeKey,
9095
spanAndBuffer.byteArrayBuffer,
9196
spanAndBuffer.charset);
92-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
97+
contextStore.put(inputStream, null);
9398
}
9499
}
95100

@@ -103,12 +108,18 @@ public static void read(
103108
spanAndBuffer.attributeKey,
104109
spanAndBuffer.byteArrayBuffer,
105110
spanAndBuffer.charset);
106-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
111+
InstrumentationContext.get(InputStream.class, SpanAndBuffer.class).put(inputStream, null);
107112
}
108113
}
109114

110115
public static void read(
111-
InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b, int off, int len) {
116+
InputStream inputStream,
117+
SpanAndBuffer spanAndBuffer,
118+
ContextStore<InputStream, SpanAndBuffer> contextStore,
119+
int read,
120+
byte[] b,
121+
int off,
122+
int len) {
112123
if (read > 0) {
113124
spanAndBuffer.byteArrayBuffer.write(b, off, read);
114125
} else if (read == -1) {
@@ -117,43 +128,37 @@ public static void read(
117128
spanAndBuffer.attributeKey,
118129
spanAndBuffer.byteArrayBuffer,
119130
spanAndBuffer.charset);
120-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
131+
contextStore.put(inputStream, null);
121132
}
122133
}
123134

124-
public static void readAll(InputStream inputStream, SpanAndBuffer spanAndBuffer, byte[] b)
135+
public static void readAll(
136+
InputStream inputStream,
137+
SpanAndBuffer spanAndBuffer,
138+
ContextStore<InputStream, SpanAndBuffer> contextStore,
139+
byte[] b)
125140
throws IOException {
126141
spanAndBuffer.byteArrayBuffer.write(b);
127-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
142+
contextStore.put(inputStream, null);
128143
}
129144

130145
public static void readNBytes(
131-
InputStream inputStream, SpanAndBuffer spanAndBuffer, int read, byte[] b, int off, int len) {
146+
InputStream inputStream,
147+
SpanAndBuffer spanAndBuffer,
148+
ContextStore<InputStream, SpanAndBuffer> contextStore,
149+
int read,
150+
byte[] b,
151+
int off,
152+
int len) {
132153
if (read == 0) {
133154
InputStreamUtils.addBody(
134155
spanAndBuffer.span,
135156
spanAndBuffer.attributeKey,
136157
spanAndBuffer.byteArrayBuffer,
137158
spanAndBuffer.charset);
138-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
159+
contextStore.put(inputStream, null);
139160
} else {
140161
spanAndBuffer.byteArrayBuffer.write(b, off, read);
141162
}
142163
}
143-
144-
public static void available(InputStream inputStream, int available) {
145-
if (available != 0) {
146-
return;
147-
}
148-
SpanAndBuffer spanAndBuffer =
149-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.get(inputStream);
150-
if (spanAndBuffer != null) {
151-
InputStreamUtils.addBody(
152-
spanAndBuffer.span,
153-
spanAndBuffer.attributeKey,
154-
spanAndBuffer.byteArrayBuffer,
155-
spanAndBuffer.charset);
156-
GlobalObjectRegistry.inputStreamToSpanAndBufferMap.remove(inputStream);
157-
}
158-
}
159164
}

0 commit comments

Comments
 (0)