Add ByteBuffer field type marshaling support to exporter. (#6686)

Author: jhalliday

exporters/common/src/main/java/io/opentelemetry/exporter/internal/marshal/CodedOutputStream.java

@@ -44,6 +44,7 @@
 import io.opentelemetry.api.internal.ConfigUtil;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 
 /**
  * Protobuf wire encoder.
@@ -56,7 +57,7 @@
 //
 // Differences
 // - No support for Message/Lite
-// - No support for ByteString or ByteBuffer
+// - No support for ByteString
 // - No support for message set extensions
 // - No support for Unsafe
 // - No support for Java String, only UTF-8 bytes
@@ -329,6 +330,11 @@ public static int computeByteArraySizeNoTag(final byte[] value) {
     return computeLengthDelimitedFieldSize(value.length);
   }
 
+  /** Compute the number of bytes that would be needed to encode a {@code bytes} field. */
+  public static int computeByteBufferSizeNoTag(final ByteBuffer value) {
+    return computeLengthDelimitedFieldSize(value.capacity());
+  }
+
   static int computeLengthDelimitedFieldSize(int fieldLength) {
     return computeUInt32SizeNoTag(fieldLength) + fieldLength;
   }
@@ -375,6 +381,8 @@ static long encodeZigZag64(final long n) {
   abstract void writeByteArrayNoTag(final byte[] value, final int offset, final int length)
       throws IOException;
 
+  abstract void writeByteBufferNoTag(final ByteBuffer value) throws IOException;
+
   // =================================================================
 
   /** Abstract base class for buffered encoders. */
@@ -487,6 +495,49 @@ void writeByteArrayNoTag(final byte[] value, int offset, int length) throws IOEx
       write(value, offset, length);
     }
 
+    @Override
+    void writeByteBufferNoTag(final ByteBuffer value) throws IOException {
+      writeUInt32NoTag(value.capacity());
+      if (value.hasArray()) {
+        write(value.array(), value.arrayOffset(), value.capacity());
+      } else {
+        write((ByteBuffer) value.duplicate().clear());
+      }
+    }
+
+    void write(ByteBuffer value) throws IOException {
+      int length = value.remaining();
+      if (limit - position >= length) {
+        // We have room in the current buffer.
+        value.get(buffer, position, length);
+        position += length;
+        totalBytesWritten += length;
+      } else {
+        // Write extends past current buffer.  Fill the rest of this buffer and
+        // flush.
+        final int bytesWritten = limit - position;
+        value.get(buffer, position, bytesWritten);
+        length -= bytesWritten;
+        position = limit;
+        totalBytesWritten += bytesWritten;
+        doFlush();
+
+        // Now deal with the rest.
+        // Since we have an output stream, this is our buffer
+        // and buffer offset == 0
+        while (length > limit) {
+          // Copy data into the buffer before writing it to OutputStream.
+          value.get(buffer, 0, limit);
+          out.write(buffer, 0, limit);
+          length -= limit;
+          totalBytesWritten += limit;
+        }
+        value.get(buffer, 0, length);
+        position = length;
+        totalBytesWritten += length;
+      }
+    }
+
     @Override
     void write(byte value) throws IOException {
       if (position == limit) {

exporters/common/src/main/java/io/opentelemetry/exporter/internal/marshal/JsonSerializer.java

@@ -9,6 +9,7 @@
 import com.fasterxml.jackson.core.JsonGenerator;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 
@@ -126,6 +127,13 @@ public void writeBytes(ProtoFieldInfo field, byte[] value) throws IOException {
     generator.writeBinaryField(field.getJsonName(), value);
   }
 
+  @Override
+  public void writeByteBuffer(ProtoFieldInfo field, ByteBuffer value) throws IOException {
+    byte[] data = new byte[value.capacity()];
+    ((ByteBuffer) value.duplicate().clear()).get(data);
+    generator.writeBinaryField(field.getJsonName(), data);
+  }
+
   @Override
   protected void writeStartMessage(ProtoFieldInfo field, int protoMessageSize) throws IOException {
     generator.writeObjectFieldStart(field.getJsonName());

exporters/common/src/main/java/io/opentelemetry/exporter/internal/marshal/MarshalerUtil.java

@@ -13,6 +13,7 @@
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.Collection;
@@ -346,6 +347,14 @@ public static int sizeBytes(ProtoFieldInfo field, byte[] message) {
     return field.getTagSize() + CodedOutputStream.computeByteArraySizeNoTag(message);
   }
 
+  /** Returns the size of a bytes field based on the buffer's capacity. */
+  public static int sizeByteBuffer(ProtoFieldInfo field, ByteBuffer message) {
+    if (message.capacity() == 0) {
+      return 0;
+    }
+    return field.getTagSize() + CodedOutputStream.computeByteBufferSizeNoTag(message);
+  }
+
   /** Returns the size of a enum field. */
   // Assumes OTLP always defines the first item in an enum with number 0, which it does and will.
   public static int sizeEnum(ProtoFieldInfo field, ProtoEnumInfo enumValue) {

exporters/common/src/main/java/io/opentelemetry/exporter/internal/marshal/ProtoSerializer.java

@@ -10,6 +10,7 @@
 import io.opentelemetry.api.trace.TraceId;
 import java.io.IOException;
 import java.io.OutputStream;
+import java.nio.ByteBuffer;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -168,6 +169,12 @@ public void writeBytes(ProtoFieldInfo field, byte[] value) throws IOException {
     output.writeByteArrayNoTag(value);
   }
 
+  @Override
+  public void writeByteBuffer(ProtoFieldInfo field, ByteBuffer value) throws IOException {
+    output.writeUInt32NoTag(field.getTag());
+    output.writeByteBufferNoTag(value);
+  }
+
   @Override
   protected void writeStartMessage(ProtoFieldInfo field, int protoMessageSize) throws IOException {
     output.writeUInt32NoTag(field.getTag());

exporters/common/src/main/java/io/opentelemetry/exporter/internal/marshal/Serializer.java

@@ -9,6 +9,7 @@
 import io.opentelemetry.api.common.Attributes;
 import io.opentelemetry.sdk.internal.DynamicPrimitiveLongList;
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
@@ -253,8 +254,22 @@ public void serializeBytes(ProtoFieldInfo field, byte[] value) throws IOExceptio
     writeBytes(field, value);
   }
 
+  /**
+   * Serializes a protobuf {@code bytes} field. Writes all content of the ByteBuffer regardless of
+   * the current position and limit. Does not alter the position or limit of the provided
+   * ByteBuffer.
+   */
+  public void serializeByteBuffer(ProtoFieldInfo field, ByteBuffer value) throws IOException {
+    if (value.capacity() == 0) {
+      return;
+    }
+    writeByteBuffer(field, value);
+  }
+
   public abstract void writeBytes(ProtoFieldInfo field, byte[] value) throws IOException;
 
+  public abstract void writeByteBuffer(ProtoFieldInfo field, ByteBuffer value) throws IOException;
+
   protected abstract void writeStartMessage(ProtoFieldInfo field, int protoMessageSize)
       throws IOException;
 

exporters/otlp/profiles/src/main/java/io/opentelemetry/exporter/otlp/profiles/ProfileContainerMarshaler.java

@@ -11,6 +11,7 @@
 import io.opentelemetry.exporter.internal.otlp.KeyValueMarshaler;
 import io.opentelemetry.proto.profiles.v1experimental.internal.ProfileContainer;
 import java.io.IOException;
+import java.nio.ByteBuffer;
 
 final class ProfileContainerMarshaler extends MarshalerWithSize {
 
@@ -20,17 +21,19 @@ final class ProfileContainerMarshaler extends MarshalerWithSize {
   private final KeyValueMarshaler[] attributeMarshalers;
   private final int droppedAttributesCount;
   private final byte[] originalPayloadFormatUtf8;
-  private final byte[] originalPayload;
+  private final ByteBuffer originalPayload;
   private final ProfileMarshaler profileMarshaler;
 
   static ProfileContainerMarshaler create(ProfileContainerData profileContainerData) {
     int droppedAttributesCount =
         profileContainerData.getTotalAttributeCount() - profileContainerData.getAttributes().size();
 
-    // Not ideal, but this will do for now. ByteBuffer support in
-    // Serialzer/CodedOutputStream/MarshalerUtilwill follow in a separate step.
-    byte[] originalPayload = new byte[profileContainerData.getOriginalPayload().remaining()];
-    profileContainerData.getOriginalPayload().get(originalPayload);
+    ByteBuffer originalPayload = profileContainerData.getOriginalPayload();
+    if (originalPayload == null) {
+      originalPayload = ByteBuffer.allocate(0);
+    } else {
+      originalPayload = originalPayload.duplicate().asReadOnlyBuffer();
+    }
 
     return new ProfileContainerMarshaler(
         profileContainerData.getProfileIdBytes(),
@@ -50,7 +53,7 @@ private ProfileContainerMarshaler(
       KeyValueMarshaler[] attributeMarshalers,
       int droppedAttributesCount,
       byte[] originalPayloadFormat,
-      byte[] originalPayload,
+      ByteBuffer originalPayload,
       ProfileMarshaler profileMarshaler) {
     super(
         calculateSize(
@@ -80,7 +83,7 @@ protected void writeTo(Serializer output) throws IOException {
     output.serializeRepeatedMessage(ProfileContainer.ATTRIBUTES, attributeMarshalers);
     output.serializeUInt32(ProfileContainer.DROPPED_ATTRIBUTES_COUNT, droppedAttributesCount);
     output.serializeString(ProfileContainer.ORIGINAL_PAYLOAD_FORMAT, originalPayloadFormatUtf8);
-    output.serializeBytes(ProfileContainer.ORIGINAL_PAYLOAD, originalPayload);
+    output.serializeByteBuffer(ProfileContainer.ORIGINAL_PAYLOAD, originalPayload);
     output.serializeMessage(ProfileContainer.PROFILE, profileMarshaler);
   }
 
@@ -91,7 +94,7 @@ private static int calculateSize(
       KeyValueMarshaler[] attributeMarshalers,
       int droppedAttributesCount,
       byte[] originalPayloadFormat,
-      byte[] originalPayload,
+      ByteBuffer originalPayload,
       ProfileMarshaler profileMarshaler) {
     int size;
     size = 0;
@@ -103,7 +106,7 @@ private static int calculateSize(
         MarshalerUtil.sizeUInt32(ProfileContainer.DROPPED_ATTRIBUTES_COUNT, droppedAttributesCount);
     size +=
         MarshalerUtil.sizeBytes(ProfileContainer.ORIGINAL_PAYLOAD_FORMAT, originalPayloadFormat);
-    size += MarshalerUtil.sizeBytes(ProfileContainer.ORIGINAL_PAYLOAD, originalPayload);
+    size += MarshalerUtil.sizeByteBuffer(ProfileContainer.ORIGINAL_PAYLOAD, originalPayload);
     size += MarshalerUtil.sizeMessage(ProfileContainer.PROFILE, profileMarshaler);
     return size;
   }