METRICS-4792 : OpenTelemetry Traces Protobuf InputFormat

distribution/pom.xml

@@ -626,7 +626,7 @@
                                         <argument>-c</argument>
                                         <argument>io.confluent.druid.extensions:confluent-extensions</argument>
                                         <argument>-c</argument>
-                                        <argument>org.apache.druid.extensions.contrib:opentelemetry-emitter</argument>
+                                        <argument>org.apache.druid.extensions.contrib:druid-opentelemetry-extensions</argument>
                                     </arguments>
                                 </configuration>
                             </execution>

extensions-contrib/opencensus-extensions/src/main/java/org/apache/druid/data/input/opencensus/protobuf/HybridProtobufReader.java

@@ -26,7 +26,7 @@
 import org.apache.druid.data.input.MapBasedInputRow;
 import org.apache.druid.data.input.impl.ByteEntity;
 import org.apache.druid.data.input.impl.DimensionsSpec;
-import org.apache.druid.data.input.opentelemetry.protobuf.OpenTelemetryMetricsProtobufReader;
+import org.apache.druid.data.input.opentelemetry.protobuf.metrics.OpenTelemetryMetricsProtobufReader;
 import org.apache.druid.indexing.seekablestream.SettableByteEntity;
 import org.apache.druid.java.util.common.parsers.CloseableIterator;
 

extensions-contrib/opencensus-extensions/src/main/java/org/apache/druid/data/input/opencensus/protobuf/OpenCensusProtobufReader.java

@@ -19,8 +19,6 @@
 
 package org.apache.druid.data.input.opencensus.protobuf;
 
-import com.google.common.base.Supplier;
-import com.google.common.base.Suppliers;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
@@ -30,34 +28,27 @@
 import io.opencensus.proto.metrics.v1.Metric;
 import io.opencensus.proto.metrics.v1.Point;
 import io.opencensus.proto.metrics.v1.TimeSeries;
-import org.apache.druid.data.input.InputEntityReader;
 import org.apache.druid.data.input.InputRow;
-import org.apache.druid.data.input.InputRowListPlusRawValues;
 import org.apache.druid.data.input.MapBasedInputRow;
 import org.apache.druid.data.input.impl.ByteEntity;
 import org.apache.druid.data.input.impl.DimensionsSpec;
+import org.apache.druid.data.input.opentelemetry.protobuf.OpenXProtobufReader;
 import org.apache.druid.indexing.seekablestream.SettableByteEntity;
-import org.apache.druid.java.util.common.CloseableIterators;
-import org.apache.druid.java.util.common.parsers.CloseableIterator;
-import org.apache.druid.java.util.common.parsers.ParseException;
 import org.apache.druid.utils.CollectionUtils;
 
 import java.nio.ByteBuffer;
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.HashSet;
-import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
-public class OpenCensusProtobufReader implements InputEntityReader
+public class OpenCensusProtobufReader extends OpenXProtobufReader
 {
   private static final String SEPARATOR = "-";
   private static final String VALUE_COLUMN = "value";
 
-  private final DimensionsSpec dimensionsSpec;
-  private final SettableByteEntity<? extends ByteEntity> source;
   private final String metricDimension;
   private final String metricLabelPrefix;
   private final String resourceLabelPrefix;
@@ -70,8 +61,7 @@ public OpenCensusProtobufReader(
       String resourceLabelPrefix
   )
   {
-    this.dimensionsSpec = dimensionsSpec;
-    this.source = source;
+    super(dimensionsSpec, source);
     this.metricDimension = metricDimension;
     this.metricLabelPrefix = metricLabelPrefix;
     this.resourceLabelPrefix = resourceLabelPrefix;
@@ -83,36 +73,9 @@ private interface LabelContext
   }
 
   @Override
-  public CloseableIterator<InputRow> read()
+  public List<InputRow> parseData(ByteBuffer byteBuffer) throws InvalidProtocolBufferException
   {
-    Supplier<Iterator<InputRow>> supplier = Suppliers.memoize(() -> readAsList().iterator());
-    return CloseableIterators.withEmptyBaggage(new Iterator<InputRow>() {
-      @Override
-      public boolean hasNext()
-      {
-        return supplier.get().hasNext();
-      }
-      @Override
-      public InputRow next()
-      {
-        return supplier.get().next();
-      }
-    });
-  }
-
-  List<InputRow> readAsList()
-  {
-    try {
-      ByteBuffer buffer = source.getEntity().getBuffer();
-      List<InputRow> rows = parseMetric(Metric.parseFrom(buffer));
-      // Explicitly move the position assuming that all the remaining bytes have been consumed because the protobuf
-      // parser does not update the position itself
-      buffer.position(buffer.limit());
-      return rows;
-    }
-    catch (InvalidProtocolBufferException e) {
-      throw new ParseException(null, e, "Protobuf message could not be parsed");
-    }
+    return parseMetric(Metric.parseFrom(byteBuffer));
   }
 
   private List<InputRow> parseMetric(final Metric metric)
@@ -222,10 +185,4 @@ private void addPointRows(Point point, Metric metric, LabelContext labelContext)
       default:
     }
   }
-
-  @Override
-  public CloseableIterator<InputRowListPlusRawValues> sample()
-  {
-    return read().map(row -> InputRowListPlusRawValues.of(row, ((MapBasedInputRow) row).getEvent()));
-  }
 }

extensions-contrib/opentelemetry-extensions/pom.xml

@@ -67,6 +67,11 @@
             <artifactId>jackson-databind</artifactId>
             <scope>provided</scope>
         </dependency>
+        <dependency>
+            <groupId>commons-codec</groupId>
+            <artifactId>commons-codec</artifactId>
+            <scope>provided</scope>
+        </dependency>
         <dependency>
             <groupId>org.apache.druid</groupId>
             <artifactId>druid-core</artifactId>
@@ -85,6 +90,11 @@
             <artifactId>junit</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>nl.jqno.equalsverifier</groupId>
+            <artifactId>equalsverifier</artifactId>
+            <scope>test</scope>
+        </dependency>
         <!-- jmh -->
         <dependency>
             <groupId>org.openjdk.jmh</groupId>

extensions-contrib/opentelemetry-extensions/src/main/java/org/apache/druid/data/input/opentelemetry/protobuf/OpenTelemetryProtobufExtensionsModule.java

@@ -23,21 +23,23 @@
 import com.fasterxml.jackson.databind.jsontype.NamedType;
 import com.fasterxml.jackson.databind.module.SimpleModule;
 import com.google.inject.Binder;
+import org.apache.druid.data.input.opentelemetry.protobuf.metrics.OpenTelemetryMetricsProtobufInputFormat;
+import org.apache.druid.data.input.opentelemetry.protobuf.traces.OpenTelemetryTracesProtobufInputFormat;
 import org.apache.druid.initialization.DruidModule;
 
 import java.util.Collections;
 import java.util.List;
 
 public class OpenTelemetryProtobufExtensionsModule implements DruidModule
 {
-
   @Override
   public List<? extends Module> getJacksonModules()
   {
     return Collections.singletonList(
         new SimpleModule("OpenTelemetryProtobufInputFormat")
             .registerSubtypes(
-                new NamedType(OpenTelemetryMetricsProtobufInputFormat.class, "opentelemetry-metrics-protobuf")
+                new NamedType(OpenTelemetryMetricsProtobufInputFormat.class, "opentelemetry-metrics-protobuf"),
+                new NamedType(OpenTelemetryTracesProtobufInputFormat.class, "opentelemetry-traces-protobuf")
             )
     );
   }

extensions-contrib/opentelemetry-extensions/src/main/java/org/apache/druid/data/input/opentelemetry/protobuf/OpenXProtobufReader.java

@@ -0,0 +1,108 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.apache.druid.data.input.opentelemetry.protobuf;
+
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
+import com.google.common.collect.Sets;
+import com.google.protobuf.InvalidProtocolBufferException;
+import org.apache.druid.data.input.InputEntityReader;
+import org.apache.druid.data.input.InputRow;
+import org.apache.druid.data.input.InputRowListPlusRawValues;
+import org.apache.druid.data.input.MapBasedInputRow;
+import org.apache.druid.data.input.impl.ByteEntity;
+import org.apache.druid.data.input.impl.DimensionsSpec;
+import org.apache.druid.indexing.seekablestream.SettableByteEntity;
+import org.apache.druid.java.util.common.CloseableIterators;
+import org.apache.druid.java.util.common.parsers.CloseableIterator;
+import org.apache.druid.java.util.common.parsers.ParseException;
+
+import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+
+public abstract class OpenXProtobufReader implements InputEntityReader
+{
+
+  protected final SettableByteEntity<? extends ByteEntity> source;
+  protected final DimensionsSpec dimensionsSpec;
+
+  public OpenXProtobufReader(DimensionsSpec dimensionsSpec,
+                             SettableByteEntity<? extends ByteEntity> source)
+  {
+    this.dimensionsSpec = dimensionsSpec;
+    this.source = source;
+  }
+
+  public abstract List<InputRow> parseData(ByteBuffer byteBuffer) throws InvalidProtocolBufferException;
+
+  @Override
+  public CloseableIterator<InputRow> read()
+  {
+    Supplier<Iterator<InputRow>> supplier = Suppliers.memoize(() -> readAsList().iterator());
+    return CloseableIterators.withEmptyBaggage(new Iterator<InputRow>() {
+      @Override
+      public boolean hasNext()
+      {
+        return supplier.get().hasNext();
+      }
+      @Override
+      public InputRow next()
+      {
+        return supplier.get().next();
+      }
+    });
+  }
+
+  public List<InputRow> readAsList()
+  {
+    try {
+      ByteBuffer buffer = source.getEntity().getBuffer();
+      List<InputRow> rows = parseData(buffer);
+      // Explicitly move the position assuming that all the remaining bytes have been consumed because the protobuf
+      // parser does not update the position itself
+      buffer.position(buffer.limit());
+      return rows;
+    }
+    catch (InvalidProtocolBufferException e) {
+      throw new ParseException(null, e, "Protobuf message could not be parsed");
+    }
+  }
+
+  protected InputRow createRow(long timeUnixMilli, Map<String, Object> event)
+  {
+    final List<String> dimensions;
+    if (!dimensionsSpec.getDimensionNames().isEmpty()) {
+      dimensions = dimensionsSpec.getDimensionNames();
+    } else {
+      dimensions = new ArrayList<>(Sets.difference(event.keySet(), dimensionsSpec.getDimensionExclusions()));
+    }
+    return new MapBasedInputRow(timeUnixMilli, dimensions, event);
+  }
+
+  @Override
+  public CloseableIterator<InputRowListPlusRawValues> sample()
+  {
+    return read().map(row -> InputRowListPlusRawValues.of(row, ((MapBasedInputRow) row).getEvent()));
+  }
+}

extensions-contrib/opentelemetry-extensions/src/main/java/org/apache/druid/data/input/opentelemetry/protobuf/Utils.java

@@ -0,0 +1,84 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.apache.druid.data.input.opentelemetry.protobuf;
+
+import io.opentelemetry.proto.common.v1.AnyValue;
+import io.opentelemetry.proto.resource.v1.Resource;
+import org.apache.druid.data.input.InputEntity;
+import org.apache.druid.data.input.impl.ByteEntity;
+import org.apache.druid.indexing.seekablestream.SettableByteEntity;
+
+import javax.annotation.Nullable;
+import java.util.HashMap;
+import java.util.Map;
+
+public class Utils
+{
+  public static SettableByteEntity<? extends ByteEntity> getSettableEntity(InputEntity source)
+  {
+    // Sampler passes a KafkaRecordEntity directly, while the normal code path wraps the same entity in a
+    // SettableByteEntity
+    if (source instanceof SettableByteEntity) {
+      return (SettableByteEntity<? extends ByteEntity>) source;
+    } else {
+      SettableByteEntity<ByteEntity> wrapper = new SettableByteEntity<>();
+      wrapper.setEntity((ByteEntity) source);
+      return wrapper;
+    }
+  }
+
+  @Nullable
+  public static Object parseAnyValue(AnyValue value)
+  {
+    switch (value.getValueCase()) {
+      case INT_VALUE:
+        return value.getIntValue();
+      case BOOL_VALUE:
+        return value.getBoolValue();
+      case DOUBLE_VALUE:
+        return value.getDoubleValue();
+      case STRING_VALUE:
+        return value.getStringValue();
+
+      // TODO: Support KVLIST_VALUE, ARRAY_VALUE and BYTES_VALUE
+
+      default:
+        // VALUE_NOT_SET
+        return null;
+    }
+  }
+
+  public static Map<String, Object> getResourceAttributes(Resource resource, String resourceAttributePrefix)
+  {
+    return resource
+        .getAttributesList()
+        .stream()
+        .collect(
+            HashMap::new,
+            (m, kv) -> {
+              Object value = Utils.parseAnyValue(kv.getValue());
+              if (value != null) {
+                m.put(resourceAttributePrefix + kv.getKey(), value);
+              }
+            },
+            HashMap::putAll
+        );
+  }
+}