Use bounded size buffers (#181)

pavolloffay · web-flow · commit 18ed5481c3c5 · 2020-12-18T14:52:27.000+01:00
* Use bounded size buffers

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

* Use config

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/ApacheClientInstrumentationModule.java b/instrumentation/apache-httpclient-4.0/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/apachehttpclient/v4_0/ApacheClientInstrumentationModule.java
@@ -51,6 +51,7 @@
 import org.apache.http.HttpMessage;
 import org.apache.http.HttpResponse;
 import org.hypertrace.agent.config.Config.AgentConfig;
+import org.hypertrace.agent.core.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.ContentEncodingUtils;
 import org.hypertrace.agent.core.ContentLengthUtils;
 import org.hypertrace.agent.core.ContentTypeUtils;
@@ -238,7 +239,7 @@ public static void exit(@Advice.This HttpEntity thizz, @Advice.Return InputStrea
       SpanAndBuffer spanAndBuffer =
           new SpanAndBuffer(
               clientSpan,
-              new ByteArrayOutputStream((int) contentSize),
+              BoundedByteArrayOutputStreamFactory.create((int) contentSize),
               HypertraceSemanticAttributes.HTTP_RESPONSE_BODY,
               charset);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.put(inputStream, spanAndBuffer);
@@ -257,7 +258,8 @@ public static void enter(
       if (contentSize <= 0 || contentSize == Long.MAX_VALUE) {
         contentSize = ContentLengthUtils.DEFAULT;
       }
-      ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream((int) contentSize);
+      ByteArrayOutputStream byteArrayOutputStream =
+          BoundedByteArrayOutputStreamFactory.create((int) contentSize);
 
       GlobalObjectRegistry.outputStreamToBufferMap.put(outputStream, byteArrayOutputStream);
     }
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
@@ -29,6 +29,7 @@
 import org.apache.http.HttpEntityEnclosingRequest;
 import org.apache.http.HttpMessage;
 import org.hypertrace.agent.config.Config.AgentConfig;
+import org.hypertrace.agent.core.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.ContentEncodingUtils;
 import org.hypertrace.agent.core.HypertraceConfig;
 import org.hypertrace.agent.core.HypertraceSemanticAttributes;
@@ -81,7 +82,7 @@ public static void traceEntity(Span span, String bodyAttributeKey, HttpEntity en
 
     if (entity.isRepeatable()) {
       try {
-        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+        ByteArrayOutputStream byteArrayOutputStream = BoundedByteArrayOutputStreamFactory.create();
         entity.writeTo(byteArrayOutputStream);
         String encoding =
             entity.getContentEncoding() != null ? entity.getContentEncoding().getValue() : "";
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
@@ -48,6 +48,7 @@
 import org.apache.http.Header;
 import org.apache.http.HttpEntity;
 import org.apache.http.HttpResponse;
+import org.hypertrace.agent.core.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.ContentTypeUtils;
 import org.hypertrace.agent.core.GlobalObjectRegistry;
 import org.hypertrace.agent.core.HypertraceSemanticAttributes;
@@ -195,7 +196,8 @@ public static void exit(@Return Object response) {
           }
 
           BufferedInputStream bufferedInputStream = new BufferedInputStream(inputStream);
-          ByteArrayOutputStream buffer = new ByteArrayOutputStream((int) contentSize);
+          ByteArrayOutputStream buffer =
+              BoundedByteArrayOutputStreamFactory.create((int) contentSize);
           byte ch;
           while ((ch = (byte) bufferedInputStream.read()) != -1) {
             buffer.write(ch);
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
@@ -31,6 +31,7 @@
 import javax.ws.rs.ext.WriterInterceptor;
 import javax.ws.rs.ext.WriterInterceptorContext;
 import org.hypertrace.agent.config.Config.AgentConfig;
+import org.hypertrace.agent.core.BoundedByteArrayOutputStreamFactory;
 import org.hypertrace.agent.core.ContentEncodingUtils;
 import org.hypertrace.agent.core.ContentLengthUtils;
 import org.hypertrace.agent.core.ContentTypeUtils;
@@ -78,7 +79,7 @@ public Object aroundReadFrom(ReaderInterceptorContext responseContext)
       String contentLengthStr = responseContext.getHeaders().getFirst(HttpHeaders.CONTENT_LENGTH);
       int contentLength = ContentLengthUtils.parseLength(contentLengthStr);
 
-      ByteArrayOutputStream buffer = new ByteArrayOutputStream(contentLength);
+      ByteArrayOutputStream buffer = BoundedByteArrayOutputStreamFactory.create(contentLength);
       GlobalObjectRegistry.inputStreamToSpanAndBufferMap.put(
           entityStream,
           new SpanAndBuffer(
@@ -118,7 +119,7 @@ public void aroundWriteTo(WriterInterceptorContext requestContext)
     }
 
     // TODO length is not known
-    ByteArrayOutputStream buffer = new ByteArrayOutputStream();
+    ByteArrayOutputStream buffer = BoundedByteArrayOutputStreamFactory.create();
     OutputStream entityStream = requestContext.getOutputStream();
     try {
       GlobalObjectRegistry.outputStreamToBufferMap.put(entityStream, buffer);
diff --git a/instrumentation/servlet/servlet-common/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/common/ByteBufferData.java b/instrumentation/servlet/servlet-common/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/common/ByteBufferData.java
@@ -19,11 +19,13 @@
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
+import org.hypertrace.agent.core.HypertraceConfig;
 
 public class ByteBufferData {
 
   private static final int MIN_BUFFER_SIZE = 128;
-  private static final int MAX_BUFFER_SIZE = 1048576; // 1MB
+  private static final int MAX_BUFFER_SIZE =
+      HypertraceConfig.get().getDataCapture().getBodyMaxSizeBytes().getValue();
   private static final int GROW_FACTOR = 4;
   private static final Charset ISO_8859_1 = StandardCharsets.ISO_8859_1;
   private Charset charset;
diff --git a/instrumentation/servlet/servlet-common/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/common/CharBufferData.java b/instrumentation/servlet/servlet-common/src/main/java/io/opentelemetry/javaagent/instrumentation/hypertrace/servlet/common/CharBufferData.java
@@ -17,11 +17,13 @@
 package io.opentelemetry.javaagent.instrumentation.hypertrace.servlet.common;
 
 import java.util.Arrays;
+import org.hypertrace.agent.core.HypertraceConfig;
 
 public class CharBufferData {
 
   private static final int MIN_BUFFER_SIZE = 128;
-  private static final int MAX_BUFFER_SIZE = 1048576; // 1MB
+  private static final int MAX_BUFFER_SIZE =
+      HypertraceConfig.get().getDataCapture().getBodyMaxSizeBytes().getValue();
   private static final int GROW_FACTOR = 4;
   private char[] buffer;
   private int bufferLen;
diff --git a/javaagent-core/src/main/java/org/hypertrace/agent/core/BoundedByteArrayOutputStream.java b/javaagent-core/src/main/java/org/hypertrace/agent/core/BoundedByteArrayOutputStream.java
@@ -0,0 +1,62 @@
+/*
+ * 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;
+
+import java.io.ByteArrayOutputStream;
+import java.io.IOException;
+
+/**
+ * {@link ByteArrayOutputStream} with a bounded capacity. Write methods are no-op if the size
+ * reaches the maximum capacity.
+ */
+public class BoundedByteArrayOutputStream extends ByteArrayOutputStream {
+
+  private final int maxCapacity;
+
+  BoundedByteArrayOutputStream(int maxCapacity) {
+    this.maxCapacity = maxCapacity;
+  }
+
+  BoundedByteArrayOutputStream(int maxCapacity, int size) {
+    super(size);
+    this.maxCapacity = maxCapacity;
+  }
+
+  @Override
+  public synchronized void write(int b) {
+    if (size() == maxCapacity) {
+      return;
+    }
+    super.write(b);
+  }
+
+  @Override
+  public void write(byte[] b) throws IOException {
+    // ByteArrayOutputStream calls write(b, off, len)
+    super.write(b);
+  }
+
+  @Override
+  public synchronized void write(byte[] b, int off, int len) {
+    int size = size();
+    if (size + len > maxCapacity) {
+      super.write(b, off, maxCapacity - size);
+      return;
+    }
+    super.write(b, off, len);
+  }
+}
diff --git a/javaagent-core/src/main/java/org/hypertrace/agent/core/BoundedByteArrayOutputStreamFactory.java b/javaagent-core/src/main/java/org/hypertrace/agent/core/BoundedByteArrayOutputStreamFactory.java
@@ -0,0 +1,33 @@
+/*
+ * 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;
+
+public class BoundedByteArrayOutputStreamFactory {
+
+  private static final int DEFAULT_SIZE = 128;
+
+  public static BoundedByteArrayOutputStream create() {
+    return new BoundedByteArrayOutputStream(DEFAULT_SIZE);
+  }
+
+  public static BoundedByteArrayOutputStream create(int initialSize) {
+    if (initialSize > DEFAULT_SIZE) {
+      initialSize = DEFAULT_SIZE;
+    }
+    return new BoundedByteArrayOutputStream(DEFAULT_SIZE, initialSize);
+  }
+}
diff --git a/javaagent-core/src/main/java/org/hypertrace/agent/core/EnvironmentConfig.java b/javaagent-core/src/main/java/org/hypertrace/agent/core/EnvironmentConfig.java
@@ -47,6 +47,7 @@ private EnvironmentConfig() {}
   static final String OPA_POLL_PERIOD = OPA_PREFIX + "poll.period.seconds";
 
   private static final String CAPTURE_PREFIX = HT_PREFIX + "data.capture.";
+  public static final String CAPTURE_BODY_MAX_SIZE_BYTES = CAPTURE_PREFIX + "body.max.size.bytes";
   public static final String CAPTURE_HTTP_HEADERS_PREFIX = CAPTURE_PREFIX + "http.headers.";
   public static final String CAPTURE_HTTP_BODY_PREFIX = CAPTURE_PREFIX + "http.body.";
   public static final String CAPTURE_RPC_METADATA_PREFIX = CAPTURE_PREFIX + "rpc.metadata.";
@@ -106,6 +107,11 @@ private static Opa.Builder applyOpa(Opa.Builder builder) {
   }
 
   private static DataCapture.Builder setDefaultsToDataCapture(DataCapture.Builder builder) {
+    String bodyMaxSizeBytes = getProperty(CAPTURE_BODY_MAX_SIZE_BYTES);
+    if (bodyMaxSizeBytes != null) {
+      builder.setBodyMaxSizeBytes(
+          Int32Value.newBuilder().setValue(Integer.valueOf(bodyMaxSizeBytes)).build());
+    }
     builder.setHttpHeaders(
         applyMessageDefaults(builder.getHttpHeaders().toBuilder(), CAPTURE_HTTP_HEADERS_PREFIX));
     builder.setHttpBody(
diff --git a/javaagent-core/src/main/java/org/hypertrace/agent/core/HypertraceConfig.java b/javaagent-core/src/main/java/org/hypertrace/agent/core/HypertraceConfig.java
@@ -57,6 +57,8 @@ private HypertraceConfig() {}
   static final String DEFAULT_REPORTING_ENDPOINT = "http://localhost:9411/api/v2/spans";
   static final String DEFAULT_OPA_ENDPOINT = "http://opa.traceableai:8181/";
   static final int DEFAULT_OPA_POLL_PERIOD_SECONDS = 30;
+  // 128 KiB
+  static final int DEFAULT_BODY_MAX_SIZE_BYTES = 128 * 1024;
 
   public static AgentConfig get() {
     if (agentConfig == null) {
@@ -199,6 +201,10 @@ private static DataCapture.Builder setDefaultsToDataCapture(DataCapture.Builder
     builder.setHttpBody(applyMessageDefaults(builder.getHttpBody().toBuilder()));
     builder.setRpcMetadata(applyMessageDefaults(builder.getRpcMetadata().toBuilder()));
     builder.setRpcBody(applyMessageDefaults(builder.getRpcBody().toBuilder()));
+    if (!builder.hasBodyMaxSizeBytes()) {
+      builder.setBodyMaxSizeBytes(
+          Int32Value.newBuilder().setValue(DEFAULT_BODY_MAX_SIZE_BYTES).build());
+    }
     return builder;
   }
 
diff --git a/javaagent-core/src/main/proto b/javaagent-core/src/main/proto
@@ -1 +1 @@
-Subproject commit 82c6d2c9e3c23bed5295a0d3bdf2812b23259ab6
+Subproject commit 56f67a79df6c0533de73967fda50da9b8bb5a29e
diff --git a/javaagent-core/src/test/java/org/hypertrace/agent/core/BoundedByteArrayOutputStreamTest.java b/javaagent-core/src/test/java/org/hypertrace/agent/core/BoundedByteArrayOutputStreamTest.java
@@ -0,0 +1,65 @@
+/*
+ * 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;
+
+import java.io.IOException;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class BoundedByteArrayOutputStreamTest {
+
+  private static final String ONE_TO_TEN = "0123456789";
+
+  @Test
+  public void writeArrTheSameSizeAsBuffer() throws IOException {
+    BoundedByteArrayOutputStream boundedBuffer = new BoundedByteArrayOutputStream(10);
+
+    boundedBuffer.write(ONE_TO_TEN.getBytes());
+    Assertions.assertEquals(10, boundedBuffer.size());
+    boundedBuffer.write("01234".getBytes());
+    Assertions.assertEquals(10, boundedBuffer.size());
+    Assertions.assertEquals(ONE_TO_TEN, boundedBuffer.toString());
+  }
+
+  @Test
+  public void writeArrSmallerSizeAsBuffer() throws IOException {
+    BoundedByteArrayOutputStream boundedBuffer = new BoundedByteArrayOutputStream(15);
+
+    boundedBuffer.write(ONE_TO_TEN.getBytes());
+    Assertions.assertEquals(10, boundedBuffer.size());
+    boundedBuffer.write("0123456".getBytes());
+    Assertions.assertEquals(15, boundedBuffer.size());
+    Assertions.assertEquals(ONE_TO_TEN + "01234", boundedBuffer.toString());
+  }
+
+  @Test
+  public void writeByteSmallerSizeAsBuffer() throws IOException {
+    BoundedByteArrayOutputStream boundedBuffer = new BoundedByteArrayOutputStream(5);
+
+    boundedBuffer.write('0');
+    boundedBuffer.write('1');
+    boundedBuffer.write('2');
+    boundedBuffer.write('3');
+    boundedBuffer.write('4');
+    Assertions.assertEquals(5, boundedBuffer.size());
+    boundedBuffer.write('5');
+    Assertions.assertEquals(5, boundedBuffer.size());
+    boundedBuffer.write("01234".getBytes());
+    Assertions.assertEquals(5, boundedBuffer.size());
+    Assertions.assertEquals("01234", boundedBuffer.toString());
+  }
+}
diff --git a/javaagent-core/src/test/java/org/hypertrace/agent/core/EnvironmentConfigTest.java b/javaagent-core/src/test/java/org/hypertrace/agent/core/EnvironmentConfigTest.java
@@ -33,6 +33,7 @@ class EnvironmentConfigTest {
   @ClearSystemProperty(key = EnvironmentConfig.OPA_POLL_PERIOD)
   @ClearSystemProperty(key = EnvironmentConfig.PROPAGATION_FORMATS)
   @ClearSystemProperty(key = EnvironmentConfig.CAPTURE_HTTP_BODY_PREFIX + "request")
+  @ClearSystemProperty(key = EnvironmentConfig.CAPTURE_BODY_MAX_SIZE_BYTES)
   public void systemProperties() {
     // when tests are run in parallel the env vars/sys props set it junit-pioneer are visible to
     // parallel tests
@@ -42,6 +43,7 @@ public void systemProperties() {
     System.setProperty(EnvironmentConfig.OPA_ENDPOINT, "http://azkaban:9090");
     System.setProperty(EnvironmentConfig.OPA_POLL_PERIOD, "10");
     System.setProperty(EnvironmentConfig.PROPAGATION_FORMATS, "B3,TRACECONTEXT");
+    System.setProperty(EnvironmentConfig.CAPTURE_BODY_MAX_SIZE_BYTES, "512");
 
     AgentConfig.Builder configBuilder = AgentConfig.newBuilder();
     configBuilder.setServiceName(StringValue.newBuilder().setValue("foo"));
@@ -56,6 +58,7 @@ public void systemProperties() {
         "http://azkaban:9090", agentConfig.getReporting().getOpa().getEndpoint().getValue());
     Assertions.assertEquals(
         10, agentConfig.getReporting().getOpa().getPollPeriodSeconds().getValue());
+    Assertions.assertEquals(512, agentConfig.getDataCapture().getBodyMaxSizeBytes().getValue());
     Assertions.assertEquals(true, agentConfig.getReporting().getSecure().getValue());
     Assertions.assertEquals(
         true, agentConfig.getDataCapture().getHttpBody().getRequest().getValue());
diff --git a/javaagent-core/src/test/java/org/hypertrace/agent/core/HypertraceConfigTest.java b/javaagent-core/src/test/java/org/hypertrace/agent/core/HypertraceConfigTest.java
@@ -48,6 +48,9 @@ public void defaultValues() throws IOException {
     Assertions.assertEquals(
         HypertraceConfig.DEFAULT_OPA_POLL_PERIOD_SECONDS,
         agentConfig.getReporting().getOpa().getPollPeriodSeconds().getValue());
+    Assertions.assertEquals(
+        HypertraceConfig.DEFAULT_BODY_MAX_SIZE_BYTES,
+        agentConfig.getDataCapture().getBodyMaxSizeBytes().getValue());
     Assertions.assertEquals(
         true, agentConfig.getDataCapture().getHttpHeaders().getRequest().getValue());
     Assertions.assertEquals(
@@ -99,6 +102,7 @@ private void assertConfig(AgentConfig agentConfig) {
         "http://opa.localhost:8181/", agentConfig.getReporting().getOpa().getEndpoint().getValue());
     Assertions.assertEquals(
         12, agentConfig.getReporting().getOpa().getPollPeriodSeconds().getValue());
+    Assertions.assertEquals(16, agentConfig.getDataCapture().getBodyMaxSizeBytes().getValue());
     Assertions.assertEquals(
         true, agentConfig.getDataCapture().getHttpHeaders().getRequest().getValue());
     Assertions.assertEquals(
diff --git a/javaagent-core/src/test/resources/config.yaml b/javaagent-core/src/test/resources/config.yaml
@@ -8,6 +8,7 @@ reporting:
     endpoint: http://opa.localhost:8181/
     pollPeriodSeconds: 12
 dataCapture:
+  bodyMaxSizeBytes: 16
   httpHeaders:
     request: true
     response: false
