diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/enrichers/ResourceAttributeEnricher.java b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/enrichers/ResourceAttributeEnricher.java new file mode 100644 index 000000000..c2757275f --- /dev/null +++ b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/enrichers/ResourceAttributeEnricher.java @@ -0,0 +1,80 @@ +package org.hypertrace.traceenricher.enrichment.enrichers; + +import static org.hypertrace.traceenricher.util.EnricherUtil.getResourceAttribute; + +import com.typesafe.config.Config; +import java.util.*; +import java.util.Map.Entry; +import java.util.stream.Collectors; +import org.hypertrace.core.datamodel.AttributeValue; +import org.hypertrace.core.datamodel.Event; +import org.hypertrace.core.datamodel.StructuredTrace; +import org.hypertrace.traceenricher.enrichment.AbstractTraceEnricher; +import org.hypertrace.traceenricher.enrichment.clients.ClientRegistry; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Enricher to add resource attributes to the spans. As of now resource attributes are attached from + * process tags. + */ +public class ResourceAttributeEnricher extends AbstractTraceEnricher { + + private static final Logger LOGGER = LoggerFactory.getLogger(ResourceAttributeEnricher.class); + private static final String RESOURCE_ATTRIBUTES_CONFIG_KEY = "attributes"; + private static final String NODE_SELECTOR_KEY = "node.selector"; + private static final String ATTRIBUTES_TO_MATCH_CONFIG_KEY = "attributesToMatch"; + private List resourceAttributesToAdd = new ArrayList<>(); + private Map resourceAttributeKeysToMatch = new HashMap<>(); + + @Override + public void init(Config enricherConfig, ClientRegistry clientRegistry) { + resourceAttributesToAdd = enricherConfig.getStringList(RESOURCE_ATTRIBUTES_CONFIG_KEY); + resourceAttributeKeysToMatch = + enricherConfig.getConfig(ATTRIBUTES_TO_MATCH_CONFIG_KEY).entrySet().stream() + .collect(Collectors.toMap(Entry::getKey, e -> e.getValue().unwrapped().toString())); + } + + @Override + public void enrichEvent(StructuredTrace trace, Event event) { + try { + if (!isValidEvent(event)) { + return; + } + Map attributeMap = event.getAttributes().getAttributeMap(); + for (String resourceAttributeKey : resourceAttributesToAdd) { + String resourceAttributeKeyToMatch = resourceAttributeKey; + if (resourceAttributeKeysToMatch.containsKey(resourceAttributeKey)) { + resourceAttributeKeyToMatch = resourceAttributeKeysToMatch.get(resourceAttributeKey); + } + Optional resourceAttributeMaybe = + getResourceAttribute(trace, event, resourceAttributeKeyToMatch); + + resourceAttributeMaybe.ifPresent( + attributeValue -> + attributeMap.computeIfAbsent( + resourceAttributeKey, + key -> { + if (resourceAttributeKey.equals(NODE_SELECTOR_KEY)) { + attributeValue.setValue( + attributeValue + .getValue() + .substring(attributeValue.getValue().lastIndexOf('/') + 1)); + } + return attributeValue; + })); + } + } catch (Exception e) { + LOGGER.error( + "Exception while enriching event with resource attributes having event id: {}", + event.getEventId(), + e); + } + } + + private boolean isValidEvent(Event event) { + return (event.getResourceIndex() >= 0) + && (event.getAttributes() != null) + && (event.getAttributes().getAttributeMap() != null); + } +} diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/util/EnricherUtil.java b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/util/EnricherUtil.java index 2a4b32be6..e8568aee4 100644 --- a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/util/EnricherUtil.java +++ b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/util/EnricherUtil.java @@ -1,11 +1,11 @@ package org.hypertrace.traceenricher.util; -import java.util.Collections; -import java.util.HashMap; -import java.util.List; -import java.util.Map; +import java.util.*; import org.apache.commons.lang3.StringUtils; +import org.hypertrace.core.datamodel.Attributes; import org.hypertrace.core.datamodel.Event; +import org.hypertrace.core.datamodel.Resource; +import org.hypertrace.core.datamodel.StructuredTrace; import org.hypertrace.core.datamodel.shared.trace.AttributeValueCreator; import org.hypertrace.core.span.constants.RawSpanConstants; import org.hypertrace.core.span.constants.v1.SpanNamePrefix; @@ -39,6 +39,24 @@ public static Map getAttributesForFirstExistingKey( return Collections.unmodifiableMap(attributes); } + public static Optional getAttribute( + Attributes attributes, String key) { + return Optional.ofNullable(attributes) + .map(Attributes::getAttributeMap) + .map(attributeMap -> attributeMap.get(key)); + } + + public static Optional getResourceAttribute( + StructuredTrace trace, Event span, String key) { + if (span.getResourceIndex() < 0 || span.getResourceIndex() >= trace.getResourceList().size()) { + return Optional.empty(); + } + + return Optional.of(trace.getResourceList().get(span.getResourceIndex())) + .map(Resource::getAttributes) + .flatMap(attributes -> getAttribute(attributes, key)); + } + public static void setAttributeForFirstExistingKey( Event event, Builder entityBuilder, List attributeKeys) { for (String attributeKey : attributeKeys) { diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/enrichment/enrichers/ResourceAttributeEnricherTest.java b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/enrichment/enrichers/ResourceAttributeEnricherTest.java new file mode 100644 index 000000000..899c7e5b1 --- /dev/null +++ b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/enrichment/enrichers/ResourceAttributeEnricherTest.java @@ -0,0 +1,234 @@ +package org.hypertrace.traceenricher.enrichment.enrichers; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.typesafe.config.Config; +import com.typesafe.config.ConfigFactory; +import java.io.File; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.hypertrace.core.datamodel.*; +import org.hypertrace.traceenricher.enrichment.clients.ClientRegistry; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.TestInstance; + +@TestInstance(TestInstance.Lifecycle.PER_CLASS) +public class ResourceAttributeEnricherTest extends AbstractAttributeEnricherTest { + + private final ResourceAttributeEnricher resourceAttributeEnricher = + new ResourceAttributeEnricher(); + + private List resourceAttributesToAddList; + + @BeforeAll + public void setup() { + String configFilePath = + Thread.currentThread().getContextClassLoader().getResource("enricher.conf").getPath(); + if (configFilePath == null) { + throw new RuntimeException("Cannot find enricher config file enricher.conf in the classpath"); + } + + Config fileConfig = ConfigFactory.parseFile(new File(configFilePath)); + Config configs = ConfigFactory.load(fileConfig); + if (!configs.hasPath("enricher.ResourceAttributeEnricher")) { + throw new RuntimeException( + "Cannot find enricher config for ResourceAttributeEnricher in " + configs); + } + Config enricherConfig = configs.getConfig("enricher.ResourceAttributeEnricher"); + resourceAttributesToAddList = enricherConfig.getStringList("attributes"); + resourceAttributeEnricher.init(enricherConfig, mock(ClientRegistry.class)); + } + + @Test + public void noResourceInTrace() { + // This trace has no resource attributes. + StructuredTrace trace = getBigTrace(); + for (Event event : trace.getEventList()) { + int attributeMapSize = 0; + if (event.getAttributes() != null && event.getAttributes().getAttributeMap() != null) { + attributeMapSize = event.getAttributes().getAttributeMap().size(); + } + resourceAttributeEnricher.enrichEvent(trace, event); + if (event.getAttributes() != null && event.getAttributes().getAttributeMap() != null) { + assertEquals(attributeMapSize, event.getAttributes().getAttributeMap().size()); + } + } + } + + @Test + public void traceWithResource() { + StructuredTrace structuredTrace = mock(StructuredTrace.class); + List resourceList = new ArrayList<>(); + + resourceList.add(getResource1()); + resourceList.add(getResource2()); + resourceList.add(getResource3()); + resourceList.add(getResource4()); + when(structuredTrace.getResourceList()).thenReturn(resourceList); + + Attributes attributes = Attributes.newBuilder().setAttributeMap(new HashMap<>()).build(); + Event event = + Event.newBuilder() + .setAttributes(attributes) + .setEventId(createByteBuffer("event1")) + .setCustomerId(TENANT_ID) + .build(); + event.setResourceIndex(0); + resourceAttributeEnricher.enrichEvent(structuredTrace, event); + assertEquals( + resourceAttributesToAddList.size() - 2, event.getAttributes().getAttributeMap().size()); + assertEquals( + "test-56f5d554c-5swkj", event.getAttributes().getAttributeMap().get("pod.name").getValue()); + assertEquals( + "01188498a468b5fef1eb4accd63533297c195a73", + event.getAttributes().getAttributeMap().get("service.version").getValue()); + assertEquals("10.21.18.171", event.getAttributes().getAttributeMap().get("ip").getValue()); + assertEquals( + "worker-hypertrace", + event.getAttributes().getAttributeMap().get("node.selector").getValue()); + + Event event2 = + Event.newBuilder() + .setAttributes(Attributes.newBuilder().setAttributeMap(new HashMap<>()).build()) + .setEventId(createByteBuffer("event2")) + .setCustomerId(TENANT_ID) + .build(); + event2.setResourceIndex(1); + addAttribute(event2, "service.version", "123"); + addAttribute(event2, "cluster.name", "default"); + resourceAttributeEnricher.enrichEvent(structuredTrace, event2); + assertEquals( + resourceAttributesToAddList.size(), event2.getAttributes().getAttributeMap().size()); + assertEquals("123", event2.getAttributes().getAttributeMap().get("service.version").getValue()); + assertEquals( + "default", event2.getAttributes().getAttributeMap().get("cluster.name").getValue()); + assertEquals( + "worker-generic", event2.getAttributes().getAttributeMap().get("node.name").getValue()); + assertEquals( + "worker-generic", event2.getAttributes().getAttributeMap().get("node.selector").getValue()); + + Event event3 = + Event.newBuilder() + .setAttributes(Attributes.newBuilder().setAttributeMap(new HashMap<>()).build()) + .setEventId(createByteBuffer("event3")) + .setCustomerId(TENANT_ID) + .build(); + event3.setResourceIndex(2); + resourceAttributeEnricher.enrichEvent(structuredTrace, event3); + assertEquals("", event3.getAttributes().getAttributeMap().get("node.selector").getValue()); + + Event event4 = + Event.newBuilder() + .setAttributes(Attributes.newBuilder().setAttributeMap(new HashMap<>()).build()) + .setEventId(createByteBuffer("event4")) + .setCustomerId(TENANT_ID) + .build(); + event4.setResourceIndex(3); + resourceAttributeEnricher.enrichEvent(structuredTrace, event4); + assertEquals( + "worker-generic", event4.getAttributes().getAttributeMap().get("node.selector").getValue()); + assertEquals("pod1", event4.getAttributes().getAttributeMap().get("pod.name").getValue()); + } + + private Resource getResource4() { + Map resourceAttributeMap = + new HashMap<>() { + { + put("node.selector", AttributeValue.newBuilder().setValue("worker-generic").build()); + put("host.name", AttributeValue.newBuilder().setValue("pod1").build()); + } + }; + return Resource.newBuilder() + .setAttributes(Attributes.newBuilder().setAttributeMap(resourceAttributeMap).build()) + .build(); + } + + private Resource getResource3() { + Map resourceAttributeMap = + new HashMap<>() { + { + put( + "node.selector", + AttributeValue.newBuilder() + .setValue("node-role.kubernetes.io/worker-generic/") + .build()); + } + }; + return Resource.newBuilder() + .setAttributes(Attributes.newBuilder().setAttributeMap(resourceAttributeMap).build()) + .build(); + } + + private Resource getResource2() { + Map resourceAttributeMap = + new HashMap<>() { + { + put( + "service.version", + AttributeValue.newBuilder() + .setValue("018a468b5fef1eb4accd63533297c195a73") + .build()); + put("environment", AttributeValue.newBuilder().setValue("stage").build()); + put( + "opencensus.exporterversion", + AttributeValue.newBuilder().setValue("Jaeger-Go-2.23.1").build()); + put("host.name", AttributeValue.newBuilder().setValue("test1-56f5d554c-5swkj").build()); + put("ip", AttributeValue.newBuilder().setValue("10.21.18.1712").build()); + put("client-uuid", AttributeValue.newBuilder().setValue("53a112a715bdf86").build()); + put("node.name", AttributeValue.newBuilder().setValue("worker-generic").build()); + put( + "cluster.name", + AttributeValue.newBuilder().setValue("worker-generic-cluster").build()); + put( + "node.selector", + AttributeValue.newBuilder() + .setValue("node-role.kubernetes.io/worker-generic") + .build()); + } + }; + return Resource.newBuilder() + .setAttributes(Attributes.newBuilder().setAttributeMap(resourceAttributeMap).build()) + .build(); + } + + private Resource getResource1() { + // In ideal scenarios below resource tags are present in spans. + Map resourceAttributeMap = + new HashMap<>() { + { + put( + "service.version", + AttributeValue.newBuilder() + .setValue("01188498a468b5fef1eb4accd63533297c195a73") + .build()); + put("environment", AttributeValue.newBuilder().setValue("stage").build()); + put( + "opencensus.exporterversion", + AttributeValue.newBuilder().setValue("Jaeger-Go-2.23.1").build()); + put("host.name", AttributeValue.newBuilder().setValue("test-56f5d554c-5swkj").build()); + put("ip", AttributeValue.newBuilder().setValue("10.21.18.171").build()); + put("client-uuid", AttributeValue.newBuilder().setValue("53a112a715bda986").build()); + put( + "node.selector", + AttributeValue.newBuilder() + .setValue("node-role.kubernetes.io/worker-hypertrace") + .build()); + } + }; + return Resource.newBuilder() + .setAttributes(Attributes.newBuilder().setAttributeMap(resourceAttributeMap).build()) + .build(); + } + + private void addAttribute(Event event, String key, String val) { + event + .getAttributes() + .getAttributeMap() + .put(key, AttributeValue.newBuilder().setValue(val).build()); + } +} diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/util/EnricherUtilTest.java b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/util/EnricherUtilTest.java index 701982a3e..d4f105926 100644 --- a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/util/EnricherUtilTest.java +++ b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/util/EnricherUtilTest.java @@ -7,6 +7,8 @@ import java.util.Arrays; import java.util.Map; +import java.util.Optional; +import org.hypertrace.core.datamodel.AttributeValue; import org.hypertrace.core.datamodel.Attributes; import org.hypertrace.core.datamodel.Event; import org.hypertrace.entity.data.service.v1.Entity; @@ -51,4 +53,25 @@ public void testSetAttributeForFirstExistingKey() { EnricherUtil.setAttributeForFirstExistingKey(e, entityBuilder, Arrays.asList("a", "b", "c")); Assertions.assertTrue(entityBuilder.getAttributesMap().containsKey("a")); } + + @Test + public void testGetAttribute() { + Attributes attributes = + Attributes.newBuilder() + .setAttributeMap( + Map.of( + "a", + TestUtil.buildAttributeValue("a-value"), + "b", + TestUtil.buildAttributeValue("b-value"))) + .build(); + Optional val = EnricherUtil.getAttribute(attributes, "a"); + Assertions.assertEquals("a-value", val.get().getValue()); + val = EnricherUtil.getAttribute(attributes, "c"); + Assertions.assertTrue(val.isEmpty()); + + attributes = null; + val = EnricherUtil.getAttribute(attributes, "a"); + Assertions.assertTrue(val.isEmpty()); + } } diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/resources/enricher.conf b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/resources/enricher.conf index ea6b6bd98..e8552f70b 100644 --- a/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/resources/enricher.conf +++ b/hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/test/resources/enricher.conf @@ -61,4 +61,12 @@ enricher { class = "org.hypertrace.traceenricher.enrichment.enrichers.GrpcAttributeEnricher" dependencies = ["SpanTypeAttributeEnricher", "ApiBoundaryTypeAttributeEnricher"] } + + ResourceAttributeEnricher { + class = "org.hypertrace.traceenricher.enrichment.enrichers.ResourceAttributeEnricher" + attributes = ["pod.name","node.name","cluster.name","ip","service.version","node.selector"] + attributesToMatch { + pod.name = "host.name" + } + } } \ No newline at end of file