Skip to content

Commit 0648fc6

Browse files
GurtejSohiGurtejSohi
authored
chore: read tenant id from process tags (#255)
* chore: read tenant id from process tags * chore: address review comments
1 parent 6593b7e commit 0648fc6

File tree

9 files changed

+91
-29
lines changed

9 files changed

+91
-29
lines changed

‎span-normalizer/span-normalizer/src/main/java/org/hypertrace/core/spannormalizer/jaeger/JaegerResourceNormalizer.java‎

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
package org.hypertrace.core.spannormalizer.jaeger;
22

3-
import static java.util.function.Predicate.not;
43
import static org.hypertrace.core.spannormalizer.util.JaegerHTTagsConverter.createFromJaegerKeyValue;
54

65
import io.jaegertracing.api_v2.JaegerSpanInternalModel.KeyValue;
@@ -22,23 +21,28 @@ class JaegerResourceNormalizer {
2221
MAP_COLLECTOR =
2322
Collectors.toUnmodifiableMap(Entry::getKey, Entry::getValue, (first, second) -> first);
2423

25-
Optional<Resource> normalize(Span span) {
24+
Optional<Resource> normalize(Span span, Optional<String> tenantIdKey) {
2625
return Optional.of(span.getProcess())
2726
.map(Process::getTagsList)
28-
.filter(not(List::isEmpty))
29-
.map(this::buildResource);
27+
.flatMap(keyValueList -> buildResource(keyValueList, tenantIdKey));
3028
}
3129

32-
private Resource buildResource(List<KeyValue> keyValueList) {
30+
private Optional<Resource> buildResource(
31+
List<KeyValue> keyValueList, Optional<String> tenantIdKey) {
3332
return keyValueList.stream()
33+
.filter(kv -> tenantIdKey.isEmpty() || !tenantIdKey.get().equalsIgnoreCase(kv.getKey()))
3434
.map(this::buildResourceValue)
3535
.collect(Collectors.collectingAndThen(MAP_COLLECTOR, this::buildResource));
3636
}
3737

38-
private Resource buildResource(Map<String, AttributeValue> resourceValueMap) {
39-
return Resource.newBuilder()
40-
.setAttributesBuilder(Attributes.newBuilder().setAttributeMap(resourceValueMap))
41-
.build();
38+
private Optional<Resource> buildResource(Map<String, AttributeValue> resourceValueMap) {
39+
if (resourceValueMap.isEmpty()) {
40+
return Optional.empty();
41+
}
42+
return Optional.of(
43+
Resource.newBuilder()
44+
.setAttributesBuilder(Attributes.newBuilder().setAttributeMap(resourceValueMap))
45+
.build());
4246
}
4347

4448
private Entry<String, AttributeValue> buildResourceValue(KeyValue keyValue) {

‎span-normalizer/span-normalizer/src/main/java/org/hypertrace/core/spannormalizer/jaeger/JaegerSpanNormalizer.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -112,7 +112,9 @@ private Callable<RawSpan> getRawSpanNormalizerCallable(
112112
tenantIdHandler.getTenantIdProvider().getTenantIdTagKey());
113113
rawSpanBuilder.setEvent(event);
114114
rawSpanBuilder.setReceivedTimeMillis(System.currentTimeMillis());
115-
resourceNormalizer.normalize(jaegerSpan).ifPresent(rawSpanBuilder::setResource);
115+
resourceNormalizer
116+
.normalize(jaegerSpan, tenantIdHandler.getTenantIdProvider().getTenantIdTagKey())
117+
.ifPresent(rawSpanBuilder::setResource);
116118

117119
// build raw span
118120
RawSpan rawSpan = rawSpanBuilder.build();

‎span-normalizer/span-normalizer/src/main/java/org/hypertrace/core/spannormalizer/jaeger/JaegerSpanPreProcessor.java‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,18 +85,22 @@ public KeyValue<byte[], PreProcessedSpan> transform(byte[] key, Span value) {
8585

8686
@VisibleForTesting
8787
PreProcessedSpan preProcessSpan(Span span) {
88-
Map<String, JaegerSpanInternalModel.KeyValue> tags =
88+
Map<String, JaegerSpanInternalModel.KeyValue> spanTags =
8989
span.getTagsList().stream()
9090
.collect(Collectors.toMap(t -> t.getKey().toLowerCase(), t -> t, (v1, v2) -> v2));
91+
Map<String, JaegerSpanInternalModel.KeyValue> processTags =
92+
span.getProcess().getTagsList().stream()
93+
.collect(Collectors.toMap(t -> t.getKey().toLowerCase(), t -> t, (v1, v2) -> v2));
9194

92-
Optional<String> maybeTenantId = tenantIdHandler.getAllowedTenantId(span, tags);
95+
Optional<String> maybeTenantId =
96+
tenantIdHandler.getAllowedTenantId(span, spanTags, processTags);
9397
if (maybeTenantId.isEmpty()) {
9498
return null;
9599
}
96100

97101
String tenantId = maybeTenantId.get();
98102

99-
if (spanFilter.shouldDropSpan(span, tags)) {
103+
if (spanFilter.shouldDropSpan(span, spanTags)) {
100104
// increment dropped counter at tenant level
101105
tenantToSpansDroppedCount
102106
.computeIfAbsent(

‎span-normalizer/span-normalizer/src/main/java/org/hypertrace/core/spannormalizer/jaeger/TenantIdHandler.java‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -84,8 +84,9 @@ private TenantIdProvider getTenantIdProvider(Config config) {
8484
}
8585
}
8686

87-
Optional<String> getAllowedTenantId(Span jaegerSpan, Map<String, KeyValue> tags) {
88-
Optional<String> maybeTenantId = this.tenantIdProvider.getTenantId(tags);
87+
Optional<String> getAllowedTenantId(
88+
Span jaegerSpan, Map<String, KeyValue> spanTags, Map<String, KeyValue> processTags) {
89+
Optional<String> maybeTenantId = this.tenantIdProvider.getTenantId(spanTags, processTags);
8990

9091
if (maybeTenantId.isEmpty()) {
9192
tenantIdProvider.logWarning(LOG, jaegerSpan);

‎span-normalizer/span-normalizer/src/main/java/org/hypertrace/core/spannormalizer/jaeger/tenant/DefaultTenantIdProvider.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ public Optional<String> getTenantIdTagKey() {
2323
}
2424

2525
@Override
26-
public Optional<String> getTenantId(Map<String, KeyValue> tags) {
26+
public Optional<String> getTenantId(
27+
Map<String, KeyValue> spanTags, Map<String, KeyValue> processTags) {
2728
return defaultTenantId;
2829
}
2930

‎span-normalizer/span-normalizer/src/main/java/org/hypertrace/core/spannormalizer/jaeger/tenant/JaegerKeyBasedTenantIdProvider.java‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,8 @@
11
package org.hypertrace.core.spannormalizer.jaeger.tenant;
22

33
import com.google.common.util.concurrent.RateLimiter;
4-
import io.jaegertracing.api_v2.JaegerSpanInternalModel;
4+
import io.jaegertracing.api_v2.JaegerSpanInternalModel.KeyValue;
5+
import io.jaegertracing.api_v2.JaegerSpanInternalModel.Span;
56
import java.util.Map;
67
import java.util.Optional;
78
import org.slf4j.Logger;
@@ -11,6 +12,7 @@
1112
* given tenant id tag key.
1213
*/
1314
public class JaegerKeyBasedTenantIdProvider implements TenantIdProvider {
15+
1416
private static final RateLimiter LOG_LIMITER = RateLimiter.create(0.1);
1517

1618
private final String tenantIdKey;
@@ -24,14 +26,24 @@ public Optional<String> getTenantIdTagKey() {
2426
return Optional.of(tenantIdKey);
2527
}
2628

29+
/**
30+
* gets the tenant id from process tags, with a fallback to getting it from the span tags.
31+
*
32+
* @param spanTags
33+
* @param processTags
34+
* @return the tenant id if it is present in the process tags or span tags; otherwise, returns an
35+
* empty optional.
36+
*/
2737
@Override
28-
public Optional<String> getTenantId(Map<String, JaegerSpanInternalModel.KeyValue> tags) {
29-
JaegerSpanInternalModel.KeyValue value = tags.get(tenantIdKey);
30-
return Optional.ofNullable(value != null ? value.getVStr() : null);
38+
public Optional<String> getTenantId(
39+
Map<String, KeyValue> spanTags, Map<String, KeyValue> processTags) {
40+
return Optional.ofNullable(processTags.get(tenantIdKey))
41+
.or(() -> Optional.ofNullable(spanTags.get(tenantIdKey)))
42+
.map(KeyValue::getVStr);
3143
}
3244

3345
@Override
34-
public void logWarning(Logger logger, JaegerSpanInternalModel.Span span) {
46+
public void logWarning(Logger logger, Span span) {
3547
if (LOG_LIMITER.tryAcquire()) {
3648
logger.warn(
3749
"Dropping span without tenant id. tenantIdTagKey: {}, span: {}", tenantIdKey, span);

‎span-normalizer/span-normalizer/src/main/java/org/hypertrace/core/spannormalizer/jaeger/tenant/TenantIdProvider.java‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,16 @@
11
package org.hypertrace.core.spannormalizer.jaeger.tenant;
22

3-
import io.jaegertracing.api_v2.JaegerSpanInternalModel;
3+
import io.jaegertracing.api_v2.JaegerSpanInternalModel.KeyValue;
4+
import io.jaegertracing.api_v2.JaegerSpanInternalModel.Span;
45
import java.util.Map;
56
import java.util.Optional;
67
import org.slf4j.Logger;
78

89
public interface TenantIdProvider {
10+
911
Optional<String> getTenantIdTagKey();
1012

11-
Optional<String> getTenantId(Map<String, JaegerSpanInternalModel.KeyValue> tags);
13+
Optional<String> getTenantId(Map<String, KeyValue> spanTags, Map<String, KeyValue> processTags);
1214

13-
void logWarning(Logger logger, JaegerSpanInternalModel.Span span);
15+
void logWarning(Logger logger, Span span);
1416
}

‎span-normalizer/span-normalizer/src/test/java/org/hypertrace/core/spannormalizer/jaeger/JaegerResourceNormalizerTest.java‎

Lines changed: 29 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@ void producesResourceWithAllKeys() {
2626
List.of(
2727
Map.entry("first", "first-value"),
2828
Map.entry("second", "second-value"),
29-
Map.entry("third", "third-value"))))
29+
Map.entry("third", "third-value"))),
30+
Optional.empty())
3031
.orElseThrow();
3132

3233
assertEquals(3, createdResource.getAttributes().getAttributeMap().size());
@@ -44,7 +45,8 @@ void ignoresDuplicateKeys() {
4445
normalizer
4546
.normalize(
4647
buildInputSpanWithResourceAttributes(
47-
List.of(Map.entry("foo", "bar"), Map.entry("foo", "baz"))))
48+
List.of(Map.entry("foo", "bar"), Map.entry("foo", "baz"))),
49+
Optional.empty())
4850
.orElseThrow();
4951

5052
assertEquals(1, createdResource.getAttributes().getAttributeMap().size());
@@ -55,7 +57,31 @@ void ignoresDuplicateKeys() {
5557
void returnsEmptyOptionalIfNoResourceAttributes() {
5658
assertEquals(
5759
Optional.empty(),
58-
normalizer.normalize(buildInputSpanWithResourceAttributes(Collections.emptyList())));
60+
normalizer.normalize(
61+
buildInputSpanWithResourceAttributes(Collections.emptyList()), Optional.empty()));
62+
}
63+
64+
@Test
65+
void returnsEmptyOptionalWithOnlyTenantIdKey() {
66+
assertEquals(
67+
Optional.empty(),
68+
normalizer.normalize(
69+
buildInputSpanWithResourceAttributes(List.of(Map.entry("tenant-key", "tenant-id"))),
70+
Optional.of("tenant-key")));
71+
}
72+
73+
@Test
74+
void ignoresTenantIdKey() {
75+
Resource createdResource =
76+
normalizer
77+
.normalize(
78+
buildInputSpanWithResourceAttributes(
79+
List.of(Map.entry("foo", "bar"), Map.entry("tenant-key", "tenant-id"))),
80+
Optional.of("tenant-key"))
81+
.orElseThrow();
82+
83+
assertEquals(1, createdResource.getAttributes().getAttributeMap().size());
84+
assertEquals("bar", createdResource.getAttributes().getAttributeMap().get("foo").getValue());
5985
}
6086

6187
// Take a list of pairs instead of a map so we can test dupe behavior

‎span-normalizer/span-normalizer/src/test/java/org/hypertrace/core/spannormalizer/jaeger/JaegerSpanPreProcessorTest.java‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@
1313
import org.junit.jupiter.api.Assertions;
1414
import org.junit.jupiter.api.Test;
1515

16-
public class JaegerSpanPreProcessorTest {
16+
class JaegerSpanPreProcessorTest {
1717

1818
private final Random random = new Random();
1919

@@ -57,7 +57,7 @@ void testPreProcessSpan_validTenantId() {
5757
PreProcessedSpan preProcessedSpan = jaegerSpanPreProcessor.preProcessSpan(span1);
5858
Assertions.assertEquals("default-tenant", preProcessedSpan.getTenantId());
5959

60-
// provided tenant id
60+
// provided tenant id in span tags
6161
configs = new HashMap<>(getCommonConfig());
6262
configs.putAll(Map.of("processor", Map.of("tenantIdTagKey", "tenant-key")));
6363
jaegerSpanPreProcessor = new JaegerSpanPreProcessor(ConfigFactory.parseMap(configs));
@@ -69,6 +69,16 @@ void testPreProcessSpan_validTenantId() {
6969
.build();
7070
preProcessedSpan = jaegerSpanPreProcessor.preProcessSpan(span2);
7171
Assertions.assertEquals(tenantId, preProcessedSpan.getTenantId());
72+
73+
// provided tenant id in process tags
74+
process =
75+
Process.newBuilder()
76+
.addTags(KeyValue.newBuilder().setKey("tenant-key").setVStr(tenantId).build())
77+
.build();
78+
79+
Span span3 = Span.newBuilder().setProcess(process).build();
80+
preProcessedSpan = jaegerSpanPreProcessor.preProcessSpan(span3);
81+
Assertions.assertEquals(tenantId, preProcessedSpan.getTenantId());
7282
}
7383

7484
@Test

0 commit comments

Comments
 (0)