Skip to content

Adds resource attribute enricher #376

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Adds resource attribute enricher #376

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 80 additions & 0 deletions ...ain/java/org/hypertrace/traceenricher/enrichment/enrichers/ResourceAttributeEnricher.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<String> resourceAttributesToAdd = new ArrayList<>();
private Map<String, String> 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<String, AttributeValue> attributeMap = event.getAttributes().getAttributeMap();
for (String resourceAttributeKey : resourceAttributesToAdd) {
String resourceAttributeKeyToMatch = resourceAttributeKey;
if (resourceAttributeKeysToMatch.containsKey(resourceAttributeKey)) {
resourceAttributeKeyToMatch = resourceAttributeKeysToMatch.get(resourceAttributeKey);
}
Optional<AttributeValue> 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);
}
}
26 changes: 22 additions & 4 deletions ...ace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/util/EnricherUtil.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -39,6 +39,24 @@ public static Map<String, AttributeValue> getAttributesForFirstExistingKey(
return Collections.unmodifiableMap(attributes);
}

public static Optional<org.hypertrace.core.datamodel.AttributeValue> getAttribute(
Attributes attributes, String key) {
return Optional.ofNullable(attributes)
.map(Attributes::getAttributeMap)
.map(attributeMap -> attributeMap.get(key));
}

public static Optional<org.hypertrace.core.datamodel.AttributeValue> 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<String> attributeKeys) {
for (String attributeKey : attributeKeys) {
Expand Down
234 changes: 234 additions & 0 deletions ...java/org/hypertrace/traceenricher/enrichment/enrichers/ResourceAttributeEnricherTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<String> 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<Resource> 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<String, AttributeValue> 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<String, AttributeValue> 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<String, AttributeValue> 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<String, AttributeValue> 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());
}
}
23 changes: 23 additions & 0 deletions ...trace-enricher-impl/src/test/java/org/hypertrace/traceenricher/util/EnricherUtilTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<AttributeValue> 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());
}
}
Loading