diff --git a/hypertrace-trace-enricher/hypertrace-trace-enricher-api/src/main/java/org/hypertrace/traceenricher/trace/util/ApiTraceGraph.java b/hypertrace-trace-enricher/hypertrace-trace-enricher-api/src/main/java/org/hypertrace/traceenricher/trace/util/ApiTraceGraph.java index 69027113c..4e3f27273 100644 --- a/hypertrace-trace-enricher/hypertrace-trace-enricher-api/src/main/java/org/hypertrace/traceenricher/trace/util/ApiTraceGraph.java +++ b/hypertrace-trace-enricher/hypertrace-trace-enricher-api/src/main/java/org/hypertrace/traceenricher/trace/util/ApiTraceGraph.java @@ -33,9 +33,9 @@ public class ApiTraceGraph { EnrichedSpanConstants.getValue(AttributeValue.ATTRIBUTE_VALUE_UNKNOWN); private List> nodeList; - private List apiNodeEventEdgeList; - private StructuredTrace trace; - private Map eventIdToIndexInTrace; + private final List apiNodeEventEdgeList; + private final StructuredTrace trace; + private final Map eventIdToIndexInTrace; public ApiTraceGraph(StructuredTrace trace) { this.trace = trace; @@ -44,6 +44,10 @@ public ApiTraceGraph(StructuredTrace trace) { eventIdToIndexInTrace = buildEventIdToIndexInTrace(trace); } + public StructuredTrace getTrace() { + return trace; + } + public List> getNodeList() { return nodeList; } diff --git a/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/BaseViewGenerator.java b/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/BaseViewGenerator.java index 1816e8051..3fce703ba 100644 --- a/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/BaseViewGenerator.java +++ b/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/BaseViewGenerator.java @@ -23,6 +23,7 @@ import org.hypertrace.core.viewgenerator.JavaCodeBasedViewGenerator; import org.hypertrace.traceenricher.enrichedspan.constants.EnrichedSpanConstants; import org.hypertrace.traceenricher.enrichedspan.constants.v1.CommonAttribute; +import org.hypertrace.viewgenerator.generators.ViewGeneratorState.TraceState; /** * Pre-processing for View Generators, for data that are mostly needed by the the view generators @@ -66,43 +67,17 @@ static String getTransactionName(StructuredTrace trace) { public List process(StructuredTrace trace) { DataflowMetricUtils.reportArrivalLagAndInsertTimestamp(trace, viewGeneratorArrivalTimer, VIEW_GENERATION_ARRIVAL_TIME); - Map entityMap = new HashMap<>(); - Map eventMap = new HashMap<>(); - Map> parentToChildrenEventIds = new HashMap<>(); - Map childToParentEventIds = new HashMap<>(); - - for (Entity entity : trace.getEntityList()) { - entityMap.put(entity.getEntityId(), entity); - } - for (Event event : trace.getEventList()) { - eventMap.put(event.getEventId(), event); - } - - for (Event event : trace.getEventList()) { - ByteBuffer childEventId = event.getEventId(); - List eventRefs = eventMap.get(childEventId).getEventRefList(); - if (eventRefs != null) { - eventRefs.stream() - .filter(eventRef -> EventRefType.CHILD_OF == eventRef.getRefType()) - .forEach( - eventRef -> { - ByteBuffer parentEventId = eventRef.getEventId(); - if (!parentToChildrenEventIds.containsKey(parentEventId)) { - parentToChildrenEventIds.put(parentEventId, new ArrayList<>()); - } - parentToChildrenEventIds.get(parentEventId).add(childEventId); - childToParentEventIds.put(childEventId, parentEventId); - }); - } - // expected only 1 childOf relationship - } + TraceState traceState = ViewGeneratorState.getTraceState(trace); + Map entityMap = Collections.unmodifiableMap(traceState.getEntityMap()); + Map eventMap = Collections.unmodifiableMap(traceState.getEventMap()); + Map> parentToChildrenEventIds = Collections.unmodifiableMap(traceState.getParentToChildrenEventIds()); + Map childToParentEventIds = Collections.unmodifiableMap(traceState.getChildToParentEventIds()); return generateView( trace, - Collections.unmodifiableMap(entityMap), - Collections.unmodifiableMap(eventMap), - Collections.unmodifiableMap(parentToChildrenEventIds), - Collections.unmodifiableMap(childToParentEventIds)); + entityMap, eventMap, + parentToChildrenEventIds, + childToParentEventIds); } abstract List generateView( diff --git a/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/RawServiceViewGenerator.java b/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/RawServiceViewGenerator.java index 9e6e97aca..9ceb7b411 100644 --- a/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/RawServiceViewGenerator.java +++ b/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/RawServiceViewGenerator.java @@ -29,8 +29,8 @@ List generateView( List list = new ArrayList<>(); // Construct ApiTraceGraph and look at all the head spans within each ApiNode - ApiTraceGraph apiTraceGraph = new ApiTraceGraph(structuredTrace); - apiTraceGraph = apiTraceGraph.build(); + ApiTraceGraph apiTraceGraph = ViewGeneratorState.getApiTraceGraph(structuredTrace); + List> apiNodes = apiTraceGraph.getNodeList(); for (ApiNode apiNode : apiNodes) { Event event = apiNode.getHeadEvent(); diff --git a/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/ServiceCallViewGenerator.java b/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/ServiceCallViewGenerator.java index ee609d755..48d77ba99 100644 --- a/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/ServiceCallViewGenerator.java +++ b/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/ServiceCallViewGenerator.java @@ -63,7 +63,7 @@ List generateView( Map eventMap, Map> parentToChildrenEventIds, Map childToParentEventIds) { - ApiTraceGraph apiTraceGraph = new ApiTraceGraph(structuredTrace).build(); + ApiTraceGraph apiTraceGraph = ViewGeneratorState.getApiTraceGraph(structuredTrace); // Scenario #1: Go through the apiNode edges and create a record for each edge. Should be easy // to get to the events diff --git a/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/ViewGeneratorState.java b/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/ViewGeneratorState.java new file mode 100644 index 000000000..6872953ac --- /dev/null +++ b/hypertrace-view-generator/hypertrace-view-generator/src/main/java/org/hypertrace/viewgenerator/generators/ViewGeneratorState.java @@ -0,0 +1,94 @@ +package org.hypertrace.viewgenerator.generators; + +import java.nio.ByteBuffer; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.hypertrace.core.datamodel.Entity; +import org.hypertrace.core.datamodel.Event; +import org.hypertrace.core.datamodel.EventRef; +import org.hypertrace.core.datamodel.EventRefType; +import org.hypertrace.core.datamodel.StructuredTrace; +import org.hypertrace.traceenricher.trace.util.ApiTraceGraph; + +public class ViewGeneratorState { + + private static final ThreadLocal traceStateThreadLocal = new ThreadLocal<>(); + private static final ThreadLocal apiTraceGraphThreadLocal = new ThreadLocal<>(); + + public static ApiTraceGraph getApiTraceGraph(StructuredTrace trace) { + if (apiTraceGraphThreadLocal.get() == null + || isDifferentTrace(apiTraceGraphThreadLocal.get().getTrace(), trace)) { + apiTraceGraphThreadLocal.set(new ApiTraceGraph(trace).build()); + } + return apiTraceGraphThreadLocal.get(); + } + + public static TraceState getTraceState(StructuredTrace trace) { + if (traceStateThreadLocal.get() == null + || isDifferentTrace(traceStateThreadLocal.get().getTrace(), trace)) { + traceStateThreadLocal.set(new TraceState(trace)); + } + return traceStateThreadLocal.get(); + } + + private static boolean isDifferentTrace(StructuredTrace cached, StructuredTrace trace) { + return cached != trace; + } + + public static class TraceState { + private final StructuredTrace trace; + private final Map entityMap = new HashMap<>(); + private final Map eventMap = new HashMap<>(); + private final Map> parentToChildrenEventIds = new HashMap<>(); + private final Map childToParentEventIds = new HashMap<>(); + + public TraceState(StructuredTrace trace) { + this.trace = trace; + for (Entity entity : trace.getEntityList()) { + entityMap.put(entity.getEntityId(), entity); + } + for (Event event : trace.getEventList()) { + eventMap.put(event.getEventId(), event); + } + + for (Event event : trace.getEventList()) { + ByteBuffer childEventId = event.getEventId(); + List eventRefs = eventMap.get(childEventId).getEventRefList(); + if (eventRefs != null) { + eventRefs.stream() + .filter(eventRef -> EventRefType.CHILD_OF == eventRef.getRefType()) + .forEach( + eventRef -> { + ByteBuffer parentEventId = eventRef.getEventId(); + parentToChildrenEventIds.computeIfAbsent( + parentEventId, v -> new ArrayList<>()).add(childEventId); + childToParentEventIds.put(childEventId, parentEventId); + }); + } + // expected only 1 childOf relationship + } + } + + public StructuredTrace getTrace() { + return trace; + } + + public Map getEntityMap() { + return entityMap; + } + + public Map getEventMap() { + return eventMap; + } + + public Map> getParentToChildrenEventIds() { + return parentToChildrenEventIds; + } + + public Map getChildToParentEventIds() { + return childToParentEventIds; + } + } +} diff --git a/hypertrace-view-generator/hypertrace-view-generator/src/test/java/org/hypertrace/viewgenerator/generators/ViewGeneratorStateTest.java b/hypertrace-view-generator/hypertrace-view-generator/src/test/java/org/hypertrace/viewgenerator/generators/ViewGeneratorStateTest.java new file mode 100644 index 000000000..da5fc5806 --- /dev/null +++ b/hypertrace-view-generator/hypertrace-view-generator/src/test/java/org/hypertrace/viewgenerator/generators/ViewGeneratorStateTest.java @@ -0,0 +1,137 @@ +package org.hypertrace.viewgenerator.generators; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.nio.ByteBuffer; +import java.util.Arrays; +import java.util.Collections; +import org.hypertrace.core.datamodel.Entity; +import org.hypertrace.core.datamodel.Event; +import org.hypertrace.core.datamodel.EventRef; +import org.hypertrace.core.datamodel.EventRefType; +import org.hypertrace.core.datamodel.StructuredTrace; +import org.hypertrace.traceenricher.trace.util.ApiTraceGraph; +import org.hypertrace.viewgenerator.generators.ViewGeneratorState.TraceState; +import org.junit.jupiter.api.Test; + +public class ViewGeneratorStateTest { + + ByteBuffer span1 = ByteBuffer.wrap(("span-1".getBytes())), span2 = ByteBuffer.wrap(("span-2".getBytes())); + String customerId = "customer-1"; + ByteBuffer traceId1 = ByteBuffer.wrap(("trace-1".getBytes())), traceId2 = ByteBuffer.wrap(("trace-2".getBytes())); + + @Test + public void testTraceState() { + TraceState traceState = new TraceState(getTestTrace(customerId, traceId1)); + assertEquals(1, traceState.getEntityMap().size()); + assertEquals(2, traceState.getEventMap().size()); + assertEquals(1, traceState.getChildToParentEventIds().size()); + assertEquals(1, traceState.getParentToChildrenEventIds().size()); + + assertTrue(traceState.getParentToChildrenEventIds().containsKey(span1)); + assertEquals(1, traceState.getParentToChildrenEventIds().get(span1).size()); + assertTrue(traceState.getChildToParentEventIds().containsKey(span2)); + assertEquals(span1, traceState.getChildToParentEventIds().get(span2)); + } + + @Test + public void testGetTraceState() { + StructuredTrace trace = getTestTrace(customerId, traceId1); + TraceState traceState = ViewGeneratorState.getTraceState(trace); + assertNotNull(traceState); + assertEquals(trace, traceState.getTrace()); + + TraceState sameTraceState = ViewGeneratorState.getTraceState(trace); + assertEquals(sameTraceState, traceState); + + StructuredTrace modifiedTrace = getTestTrace(customerId, traceId1); + modifiedTrace.setEntityList(Arrays.asList( + Entity.newBuilder() + .setCustomerId(customerId) + .setEntityId("entity-2") + .setEntityName("entity-2") + .setEntityType("service") + .build())); + + // same trace id but different object should result in rebuilding of trace state + TraceState differentTraceState1 = ViewGeneratorState.getTraceState(modifiedTrace); + assertNotEquals(traceState, differentTraceState1); + + StructuredTrace differentTrace = getTestTrace(customerId, traceId2); + TraceState differentTraceState2 = ViewGeneratorState.getTraceState(differentTrace); + assertEquals(differentTrace, differentTraceState2.getTrace()); + } + + @Test + public void testGetApiTraceGraph() { + StructuredTrace trace = getTestTrace(customerId, traceId1); + ApiTraceGraph apiTraceGraph = ViewGeneratorState.getApiTraceGraph(trace); + assertNotNull(apiTraceGraph); + + ApiTraceGraph sameApiTraceGraph = ViewGeneratorState.getApiTraceGraph(trace); + assertEquals(sameApiTraceGraph, apiTraceGraph); + + StructuredTrace modifiedTrace = getTestTrace(customerId, traceId1); + modifiedTrace.setEntityList(Arrays.asList( + Entity.newBuilder() + .setCustomerId(customerId) + .setEntityId("entity-2") + .setEntityName("entity-2") + .setEntityType("service") + .build())); + + + // same trace id but different object should result in rebuilding of trace state + ApiTraceGraph differentApiTraceGraph1 = ViewGeneratorState.getApiTraceGraph(modifiedTrace); + assertNotEquals(apiTraceGraph, differentApiTraceGraph1); + + StructuredTrace differentTrace = getTestTrace(customerId, traceId2); + ApiTraceGraph differentApiTraceGraph2 = ViewGeneratorState.getApiTraceGraph(differentTrace); + assertNotEquals(apiTraceGraph, differentApiTraceGraph2); + } + + private StructuredTrace getTestTrace(String customerId, ByteBuffer traceId) { + return StructuredTrace.newBuilder() + .setCustomerId(customerId) + .setTraceId(traceId) + .setStartTimeMillis(20) + .setEndTimeMillis(30) + .setEntityList(Arrays.asList( + Entity.newBuilder() + .setCustomerId(customerId) + .setEntityId("entity-1") + .setEntityName("entity-1") + .setEntityType("service") + .build())) + .setEventList(Arrays.asList( + Event.newBuilder() + .setCustomerId(customerId) + .setEventId(ByteBuffer.wrap(("span-1".getBytes()))) + .setEventName("span-1") + .build(), + Event.newBuilder() + .setCustomerId(customerId) + .setEventId(ByteBuffer.wrap(("span-2".getBytes()))) + .setEventName("span-2") + .setEventRefList(Arrays.asList( + EventRef.newBuilder() + .setTraceId(traceId) + .setRefType(EventRefType.CHILD_OF) + .setEventId(ByteBuffer.wrap(("span-1".getBytes()))) + .build() + )) + .build() + )) + .setEntityEdgeList(Collections.EMPTY_LIST) + .setEntityEventEdgeList(Collections.EMPTY_LIST) + .setEventEdgeList(Collections.EMPTY_LIST) + .setEntityEntityGraph(null) + .setEventEventGraph(null) + .setEntityEventGraph(null) + .setMetrics(null) + .build(); + } +}