Bounded Grpc payload sizes (#211)

pavolloffay · web-flow · commit 4989a264554f · 2021-01-08T15:22:41.000+01:00
* Buffer impl

Signed-off-by: Pavol Loffay &lt;p.loffay@gmail.com&gt;

* Add test

Signed-off-by: Pavol Loffay &lt;p.loffay@gmail.com&gt;
diff --git a/instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/ApacheHttpClientUtils.java b/instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/ApacheHttpClientUtils.java
@@ -32,8 +32,8 @@
 import org.hypertrace.agent.config.Config.AgentConfig;
 import org.hypertrace.agent.core.config.HypertraceConfig;
 import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeCharsetUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeUtils;
 import org.slf4j.Logger;
@@ -107,7 +107,7 @@ public static void traceEntity(
     if (entity.isRepeatable()) {
       try {
         BoundedByteArrayOutputStream byteArrayOutputStream =
-            BoundedByteArrayOutputStreamFactory.create(charset);
+            BoundedBuffersFactory.createStream(charset);
         entity.writeTo(byteArrayOutputStream);
 
         try {
diff --git a/instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/HttpEntityInstrumentation.java b/instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/HttpEntityInstrumentation.java
@@ -39,8 +39,8 @@
 import org.apache.http.HttpEntity;
 import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
 import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry.SpanAndBuffer;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.instrumentation.utils.ContentLengthUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeCharsetUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeUtils;
@@ -100,7 +100,7 @@ public static void exit(@Advice.This HttpEntity thizz, @Advice.Return InputStrea
       SpanAndBuffer spanAndBuffer =
           new SpanAndBuffer(
               clientSpan.span,
-              BoundedByteArrayOutputStreamFactory.create((int) contentSize, charset),
+              BoundedBuffersFactory.createStream((int) contentSize, charset),
               clientSpan.attributeKey,
               charset);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.put(inputStream, spanAndBuffer);
@@ -129,7 +129,7 @@ public static void enter(
       Charset charset = ContentTypeCharsetUtils.toCharset(charsetStr);
 
       BoundedByteArrayOutputStream byteArrayOutputStream =
-          BoundedByteArrayOutputStreamFactory.create((int) contentSize, charset);
+          BoundedBuffersFactory.createStream((int) contentSize, charset);
 
       GlobalObjectRegistry.outputStreamToBufferMap.put(outputStream, byteArrayOutputStream);
     }
diff --git a/instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/readall/ApacheClientReadAllInstrumentationModule.java b/instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/readall/ApacheClientReadAllInstrumentationModule.java
@@ -51,7 +51,7 @@
 import org.apache.http.HttpResponse;
 import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
 import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeUtils;
 
 @AutoService(InstrumentationModule.class)
@@ -198,8 +198,7 @@ public static void exit(@Return Object response) {
 
           BufferedInputStream bufferedInputStream = new BufferedInputStream(inputStream);
           ByteArrayOutputStream buffer =
-              BoundedByteArrayOutputStreamFactory.create(
-                  (int) contentSize, Charset.defaultCharset());
+              BoundedBuffersFactory.createStream((int) contentSize, Charset.defaultCharset());
           byte ch;
           while ((ch = (byte) bufferedInputStream.read()) != -1) {
             buffer.write(ch);
diff --git a/instrumentation/grpc-1.5/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/grpc/v1_5/GrpcSpanDecorator.java b/instrumentation/grpc-1.5/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/grpc/v1_5/GrpcSpanDecorator.java
@@ -16,17 +16,19 @@
 
 package io.opentelemetry.javaagent.instrumentation.hypertrace.grpc.v1_5;
 
-import com.google.protobuf.InvalidProtocolBufferException;
 import com.google.protobuf.Message;
 import com.google.protobuf.util.JsonFormat;
 import io.grpc.Metadata;
 import io.grpc.Metadata.Key;
 import io.opentelemetry.api.common.AttributeKey;
 import io.opentelemetry.api.trace.Span;
+import java.io.IOException;
 import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.function.Function;
 import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedCharArrayWriter;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -41,9 +43,10 @@ public static void addMessageAttribute(Object message, Span span, AttributeKey<S
     if (message instanceof Message) {
       Message mb = (Message) message;
       try {
-        String requestBody = PRINTER.print(mb);
-        span.setAttribute(key, requestBody);
-      } catch (InvalidProtocolBufferException e) {
+        BoundedCharArrayWriter writer = BoundedBuffersFactory.createWriter();
+        PRINTER.appendTo(mb, writer);
+        span.setAttribute(key, writer.toString());
+      } catch (IOException e) {
         log.error("Failed to decode message to JSON", e);
       }
     }
diff --git a/instrumentation/java-streams/src/test/java/io/opentelemetry/javaagent/instrumentation/hypertrace/java/outputstream/OutputStreamInstrumentationModuleTest.java b/instrumentation/java-streams/src/test/java/io/opentelemetry/javaagent/instrumentation/hypertrace/java/outputstream/OutputStreamInstrumentationModuleTest.java
@@ -22,8 +22,8 @@
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.testing.AbstractInstrumenterTest;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.Test;
@@ -87,7 +87,7 @@ public void writeBytesOffset() {
 
   private void write(OutputStream outputStream, Runnable read, String expected) {
     BoundedByteArrayOutputStream buffer =
-        BoundedByteArrayOutputStreamFactory.create(MAX_SIZE, DEFAULT_CHARSET);
+        BoundedBuffersFactory.createStream(MAX_SIZE, DEFAULT_CHARSET);
     GlobalObjectRegistry.outputStreamToBufferMap.put(outputStream, buffer);
     read.run();
     Assertions.assertEquals(expected, buffer.toString());
diff --git a/instrumentation/jaxrs-client-2.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/jaxrs/v2_0/JaxrsClientEntityInterceptor.java b/instrumentation/jaxrs-client-2.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/jaxrs/v2_0/JaxrsClientEntityInterceptor.java
@@ -35,8 +35,8 @@
 import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry;
 import org.hypertrace.agent.core.instrumentation.GlobalObjectRegistry.SpanAndBuffer;
 import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.instrumentation.utils.ContentLengthUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeCharsetUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeUtils;
@@ -86,7 +86,7 @@ public Object aroundReadFrom(ReaderInterceptorContext responseContext)
       Charset charset = ContentTypeCharsetUtils.toCharset(charsetStr);
 
       BoundedByteArrayOutputStream buffer =
-          BoundedByteArrayOutputStreamFactory.create(contentLength, charset);
+          BoundedBuffersFactory.createStream(contentLength, charset);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.put(
           entityStream,
           new SpanAndBuffer(
@@ -130,7 +130,7 @@ public void aroundWriteTo(WriterInterceptorContext requestContext)
     Charset charset = ContentTypeCharsetUtils.toCharset(charsetStr);
 
     // TODO length is not known
-    BoundedByteArrayOutputStream buffer = BoundedByteArrayOutputStreamFactory.create(charset);
+    BoundedByteArrayOutputStream buffer = BoundedBuffersFactory.createStream(charset);
     OutputStream entityStream = requestContext.getOutputStream();
     try {
       GlobalObjectRegistry.outputStreamToBufferMap.put(entityStream, buffer);
diff --git a/instrumentation/netty/netty-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_0/server/HttpServerRequestTracingHandler.java b/instrumentation/netty/netty-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_0/server/HttpServerRequestTracingHandler.java
@@ -34,8 +34,8 @@
 import org.hypertrace.agent.config.Config.AgentConfig;
 import org.hypertrace.agent.core.config.HypertraceConfig;
 import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.instrumentation.utils.ContentLengthUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeCharsetUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeUtils;
@@ -79,7 +79,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) {
         // the buffer is used byt captureBody method
         Attribute<BoundedByteArrayOutputStream> bufferAttr =
             ctx.channel().attr(AttributeKeys.REQUEST_BODY_BUFFER);
-        bufferAttr.set(BoundedByteArrayOutputStreamFactory.create(contentLength, charset));
+        bufferAttr.set(BoundedBuffersFactory.createStream(contentLength, charset));
       }
     }
 
diff --git a/instrumentation/netty/netty-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_0/server/HttpServerResponseTracingHandler.java b/instrumentation/netty/netty-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_0/server/HttpServerResponseTracingHandler.java
@@ -39,8 +39,8 @@
 import org.hypertrace.agent.config.Config.AgentConfig;
 import org.hypertrace.agent.core.config.HypertraceConfig;
 import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.instrumentation.utils.ContentLengthUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeCharsetUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeUtils;
@@ -79,7 +79,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) {
         // the buffer is used byt captureBody method
         Attribute<BoundedByteArrayOutputStream> bufferAttr =
             ctx.channel().attr(AttributeKeys.RESPONSE_BODY_BUFFER);
-        bufferAttr.set(BoundedByteArrayOutputStreamFactory.create(contentLength, charset));
+        bufferAttr.set(BoundedBuffersFactory.createStream(contentLength, charset));
       }
     }
 
diff --git a/instrumentation/netty/netty-4.1/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_1/server/HttpServerRequestTracingHandler.java b/instrumentation/netty/netty-4.1/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_1/server/HttpServerRequestTracingHandler.java
@@ -34,8 +34,8 @@
 import org.hypertrace.agent.config.Config.AgentConfig;
 import org.hypertrace.agent.core.config.HypertraceConfig;
 import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.instrumentation.utils.ContentLengthUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeCharsetUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeUtils;
@@ -79,7 +79,7 @@ public void channelRead(ChannelHandlerContext ctx, Object msg) {
         // the buffer is used byt captureBody method
         Attribute<BoundedByteArrayOutputStream> bufferAttr =
             ctx.channel().attr(AttributeKeys.REQUEST_BODY_BUFFER);
-        bufferAttr.set(BoundedByteArrayOutputStreamFactory.create(contentLength, charset));
+        bufferAttr.set(BoundedBuffersFactory.createStream(contentLength, charset));
       }
     }
 
diff --git a/instrumentation/netty/netty-4.1/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_1/server/HttpServerResponseTracingHandler.java b/instrumentation/netty/netty-4.1/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/netty/v4_1/server/HttpServerResponseTracingHandler.java
@@ -37,8 +37,8 @@
 import org.hypertrace.agent.config.Config.AgentConfig;
 import org.hypertrace.agent.core.config.HypertraceConfig;
 import org.hypertrace.agent.core.instrumentation.HypertraceSemanticAttributes;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStream;
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.instrumentation.utils.ContentLengthUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeCharsetUtils;
 import org.hypertrace.agent.core.instrumentation.utils.ContentTypeUtils;
@@ -77,7 +77,7 @@ public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise prm) {
         // the buffer is used byt captureBody method
         Attribute<BoundedByteArrayOutputStream> bufferAttr =
             ctx.channel().attr(AttributeKeys.RESPONSE_BODY_BUFFER);
-        bufferAttr.set(BoundedByteArrayOutputStreamFactory.create(contentLength, charset));
+        bufferAttr.set(BoundedBuffersFactory.createStream(contentLength, charset));
       }
     }
 
diff --git a/javaagent-core/src/main/java/org/hypertrace/agent/core/instrumentation/buffer/BoundedBuffersFactory.java b/javaagent-core/src/main/java/org/hypertrace/agent/core/instrumentation/buffer/BoundedBuffersFactory.java
@@ -19,19 +19,30 @@
 import java.nio.charset.Charset;
 import org.hypertrace.agent.core.config.HypertraceConfig;
 
-public class BoundedByteArrayOutputStreamFactory {
+public class BoundedBuffersFactory {
 
-  public static final int DEFAULT_SIZE =
+  public static final int MAX_SIZE =
       HypertraceConfig.get().getDataCapture().getBodyMaxSizeBytes().getValue();
 
-  public static BoundedByteArrayOutputStream create(Charset charset) {
-    return new BoundedByteArrayOutputStream(DEFAULT_SIZE, charset);
+  public static BoundedByteArrayOutputStream createStream(Charset charset) {
+    return new BoundedByteArrayOutputStream(MAX_SIZE, charset);
   }
 
-  public static BoundedByteArrayOutputStream create(int initialSize, Charset charset) {
-    if (initialSize > DEFAULT_SIZE) {
-      initialSize = DEFAULT_SIZE;
+  public static BoundedByteArrayOutputStream createStream(int initialSize, Charset charset) {
+    if (initialSize > MAX_SIZE) {
+      initialSize = MAX_SIZE;
     }
-    return new BoundedByteArrayOutputStream(DEFAULT_SIZE, initialSize, charset);
+    return new BoundedByteArrayOutputStream(MAX_SIZE, initialSize, charset);
+  }
+
+  public static BoundedCharArrayWriter createWriter() {
+    return new BoundedCharArrayWriter(MAX_SIZE);
+  }
+
+  public static BoundedCharArrayWriter createWriter(int initialSize) {
+    if (initialSize > MAX_SIZE) {
+      initialSize = MAX_SIZE;
+    }
+    return new BoundedCharArrayWriter(MAX_SIZE, initialSize);
   }
 }
diff --git a/javaagent-core/src/main/java/org/hypertrace/agent/core/instrumentation/buffer/BoundedCharArrayWriter.java b/javaagent-core/src/main/java/org/hypertrace/agent/core/instrumentation/buffer/BoundedCharArrayWriter.java
@@ -0,0 +1,100 @@
+/*
+ * Copyright The Hypertrace Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.hypertrace.agent.core.instrumentation.buffer;
+
+import java.io.CharArrayWriter;
+import java.io.IOException;
+import java.io.Writer;
+
+public class BoundedCharArrayWriter extends CharArrayWriter {
+
+  private final int maxCapacity;
+
+  BoundedCharArrayWriter(int maxCapacity) {
+    this.maxCapacity = maxCapacity;
+  }
+
+  BoundedCharArrayWriter(int maxCapacity, int initialSize) {
+    super(initialSize);
+    this.maxCapacity = maxCapacity;
+  }
+
+  @Override
+  public void write(int c) {
+    if (size() == maxCapacity) {
+      return;
+    }
+    super.write(c);
+  }
+
+  @Override
+  public void write(char[] c, int off, int len) {
+    int size = size();
+    if (size + len > maxCapacity) {
+      super.write(c, off, maxCapacity - size);
+      return;
+    }
+    super.write(c, off, len);
+  }
+
+  @Override
+  public void write(String str, int off, int len) {
+    int size = size();
+    if (size + len > maxCapacity) {
+      super.write(str, off, maxCapacity - size);
+      return;
+    }
+    super.write(str, off, len);
+  }
+
+  @Override
+  public void writeTo(Writer out) throws IOException {
+    // don't need to override, it will copy what is buffered
+    super.writeTo(out);
+  }
+
+  @Override
+  public CharArrayWriter append(CharSequence csq) {
+    // the same impl as the super class
+    String s = (csq == null ? "null" : csq.toString());
+    this.write(s, 0, s.length());
+    return this;
+  }
+
+  @Override
+  public CharArrayWriter append(CharSequence csq, int start, int end) {
+    String s = (csq == null ? "null" : csq).subSequence(start, end).toString();
+    write(s, 0, s.length());
+    return this;
+  }
+
+  @Override
+  public CharArrayWriter append(char c) {
+    write(c);
+    return this;
+  }
+
+  @Override
+  public void write(char[] cbuf) throws IOException {
+    this.write(cbuf, 0, cbuf.length);
+  }
+
+  @Override
+  public void write(String str) throws IOException {
+    this.write(str, 0, str.length());
+  }
+}
diff --git a/javaagent-core/src/main/java/org/hypertrace/agent/core/instrumentation/utils/ContentLengthUtils.java b/javaagent-core/src/main/java/org/hypertrace/agent/core/instrumentation/utils/ContentLengthUtils.java
@@ -16,13 +16,13 @@
 
 package org.hypertrace.agent.core.instrumentation.utils;
 
-import org.hypertrace.agent.core.instrumentation.buffer.BoundedByteArrayOutputStreamFactory;
+import org.hypertrace.agent.core.instrumentation.buffer.BoundedBuffersFactory;
 
 public class ContentLengthUtils {
   private ContentLengthUtils() {}
 
   // default content length
-  public static final int DEFAULT = BoundedByteArrayOutputStreamFactory.DEFAULT_SIZE;
+  public static final int DEFAULT = BoundedBuffersFactory.MAX_SIZE;
 
   public static int parseLength(CharSequence lengthStr) {
     if (lengthStr == null || lengthStr.length() == 0) {
diff --git a/javaagent-core/src/test/java/org/hypertrace/agent/core/instrumentation/buffer/BoundedCharArrayWriterTest.java b/javaagent-core/src/test/java/org/hypertrace/agent/core/instrumentation/buffer/BoundedCharArrayWriterTest.java
