fix for performance observed during enrichment process (#100)

* fix for performance

* functional validation

* changes logs

* cleaning PR by removing measuing logging codes

* Removes un-used logger

* Adds comments for redability

Author: kotharironak

hypertrace-trace-enricher/hypertrace-trace-enricher-api/src/main/java/org/hypertrace/traceenricher/trace/util/ApiTraceGraph.java

@@ -178,7 +178,7 @@ private Pair<List<Event>, List<Event>> getEventsInApiBoundary(
   }
 
   public ApiTraceGraph build() {
-    StructuredTraceGraph graph = StructuredTraceGraph.createGraph(trace);
+    StructuredTraceGraph graph = StructuredTraceGraphBuilder.buildGraph(trace);
 
     List<ApiNode<Event>> apiNodes = buildApiNodes(graph);
     this.setNodeList(apiNodes);

hypertrace-trace-enricher/hypertrace-trace-enricher-api/src/main/java/org/hypertrace/traceenricher/trace/util/StructuredTraceGraphBuilder.java

@@ -0,0 +1,49 @@
+package org.hypertrace.traceenricher.trace.util;
+
+import org.hypertrace.core.datamodel.StructuredTrace;
+import org.hypertrace.core.datamodel.shared.StructuredTraceGraph;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class StructuredTraceGraphBuilder {
+  private static final Logger LOG = LoggerFactory.getLogger(StructuredTraceGraphBuilder.class);
+
+  private static ThreadLocal<StructuredTraceGraph> cachedGraph = new ThreadLocal<>();
+  private static ThreadLocal<StructuredTrace> cachedTrace = new ThreadLocal<>();
+
+  public static StructuredTraceGraph buildGraph(StructuredTrace trace) {
+    // trace doesn't exist
+    if (cachedTrace.get() == null) {
+      LOG.info("Building structured trace graph. Reason: no cached trace");
+      StructuredTraceGraph graph = StructuredTraceGraph.createGraph(trace);
+      cachedTrace.set(StructuredTrace.newBuilder(trace).build());
+      cachedGraph.set(graph);
+      return graph;
+    }
+
+    // is processed and cached are same trace?
+    if (!cachedTrace.get().getCustomerId().equals(trace.getCustomerId()) ||
+        !cachedTrace.get().getTraceId().equals(trace.getTraceId())) {
+      LOG.info("Building structured trace graph. Reason: cached trace and current trace doesn't not match");
+      StructuredTraceGraph graph = StructuredTraceGraph.createGraph(trace);
+      cachedTrace.set(StructuredTrace.newBuilder(trace).build());
+      cachedGraph.set(graph);
+      return graph;
+    }
+
+    // trace internally changed
+    if(cachedTrace.get().getEntityList().size()  != trace.getEntityList().size() ||
+        cachedTrace.get().getEventList().size() != trace.getEventList().size() ||
+        cachedTrace.get().getEntityEdgeList().size() != trace.getEntityEdgeList().size() ||
+        cachedTrace.get().getEntityEventEdgeList().size() != trace.getEntityEventEdgeList().size() ||
+        cachedTrace.get().getEventEdgeList().size() != trace.getEventEdgeList().size()) {
+      LOG.info("Building structured trace graph. Reason: cached trace and current trace have different size");
+      StructuredTraceGraph graph = StructuredTraceGraph.createGraph(trace);
+      cachedTrace.set(StructuredTrace.newBuilder(trace).build());
+      cachedGraph.set(graph);
+      return graph;
+    }
+
+    return cachedGraph.get();
+  }
+}

hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/AbstractTraceEnricher.java

@@ -16,6 +16,7 @@
 import org.hypertrace.core.datamodel.StructuredTrace;
 import org.hypertrace.core.datamodel.shared.StructuredTraceGraph;
 import org.hypertrace.entity.data.service.client.EntityDataServiceClientProvider;
+import org.hypertrace.traceenricher.trace.util.StructuredTraceGraphBuilder;
 
 public abstract class AbstractTraceEnricher implements Enricher {
 
@@ -52,7 +53,7 @@ public void enrichTrace(StructuredTrace trace) {
    * Wrapper to the structure graph factory for testing
    */
   public StructuredTraceGraph buildGraph(StructuredTrace trace) {
-    return StructuredTraceGraph.createGraph(trace);
+    return StructuredTraceGraphBuilder.buildGraph(trace);
   }
 
   @Nullable

hypertrace-trace-enricher/hypertrace-trace-enricher-impl/src/main/java/org/hypertrace/traceenricher/enrichment/enrichers/DefaultServiceEntityEnricher.java

@@ -10,8 +10,6 @@
 import org.hypertrace.core.datamodel.shared.SpanAttributeUtils;
 import org.hypertrace.core.datamodel.shared.StructuredTraceGraph;
 import org.hypertrace.core.datamodel.shared.trace.AttributeValueCreator;
-import org.hypertrace.core.span.constants.RawSpanConstants;
-import org.hypertrace.core.span.constants.v1.JaegerAttribute;
 import org.hypertrace.entity.constants.v1.ServiceAttribute;
 import org.hypertrace.entity.data.service.client.EdsClient;
 import org.hypertrace.entity.data.service.client.EntityDataServiceClientProvider;
@@ -69,7 +67,7 @@ public void enrichEvent(StructuredTrace trace, Event event) {
       // 2. Enrich the exit span with the parent span's service entity.
       //    This will enable creating an edge between the exit span and the backend
 
-      StructuredTraceGraph graph = StructuredTraceGraph.createGraph(trace);
+      StructuredTraceGraph graph = buildGraph(trace);
       if (EnrichedSpanUtils.isExitSpan(event) &&
           SpanAttributeUtils.isLeafSpan(graph, event)) {
         String parentSvcName = findServiceNameOfFirstAncestorThatIsNotAnExitSpanAndBelongsToADifferentService(event,

hypertrace-trace-enricher/hypertrace-trace-enricher/src/main/java/org/hypertrace/traceenricher/trace/enricher/StructuredTraceEnrichProcessor.java

@@ -5,6 +5,8 @@
 import static org.hypertrace.traceenricher.trace.enricher.StructuredTraceEnricherConstants.STRUCTURED_TRACES_ENRICHMENT_JOB_CONFIG_KEY;
 
 import com.typesafe.config.Config;
+import java.time.Duration;
+import java.time.Instant;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -15,10 +17,13 @@
 import org.hypertrace.entity.data.service.client.DefaultEdsClientProvider;
 import org.hypertrace.traceenricher.enrichment.EnrichmentProcessor;
 import org.hypertrace.traceenricher.enrichment.EnrichmentRegistry;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 public class StructuredTraceEnrichProcessor implements
     Transformer<String, StructuredTrace, KeyValue<String, StructuredTrace>> {
 
+  private static final Logger logger = LoggerFactory.getLogger(StructuredTraceEnrichProcessor.class);
   private static EnrichmentProcessor processor = null;
 
   @Override
@@ -38,7 +43,11 @@ public void init(ProcessorContext context) {
 
   @Override
   public KeyValue<String, StructuredTrace> transform(String key, StructuredTrace value) {
+    Instant start = Instant.now();
     processor.process(value);
+    Instant finish = Instant.now();
+    long timeElapsed = Duration.between(start, finish).toMillis();
+    logger.debug("Time taken by enrichment process for key:{} is :{}", key, timeElapsed);
     return new KeyValue<>(null, value);
   }