feat(telemetry): cross-component async write tracing

build.gradle

@@ -38,7 +38,7 @@ buildscript {
   ext.springVersion = '6.1.14'
   ext.springBootVersion = '3.2.9'
   ext.springKafkaVersion = '3.1.6'
-  ext.openTelemetryVersion = '1.18.0'
+  ext.openTelemetryVersion = '1.45.0'
   ext.neo4jVersion = '5.14.0'
   ext.neo4jTestVersion = '5.14.0'
   ext.neo4jApocVersion = '5.14.0'
@@ -218,7 +218,10 @@ project.ext.externalDependency = [
     'neo4jApocCore': 'org.neo4j.procedure:apoc-core:' + neo4jApocVersion,
     'neo4jApocCommon': 'org.neo4j.procedure:apoc-common:' + neo4jApocVersion,
     'opentelemetryApi': 'io.opentelemetry:opentelemetry-api:' + openTelemetryVersion,
-    'opentelemetryAnnotations': 'io.opentelemetry:opentelemetry-extension-annotations:' + openTelemetryVersion,
+    'opentelemetrySdk': 'io.opentelemetry:opentelemetry-sdk:' + openTelemetryVersion,
+    'opentelemetrySdkTrace': 'io.opentelemetry:opentelemetry-sdk-trace:' + openTelemetryVersion,
+    'opentelemetryAutoConfig': 'io.opentelemetry:opentelemetry-sdk-extension-autoconfigure:' + openTelemetryVersion,
+    'opentelemetryAnnotations': 'io.opentelemetry.instrumentation:opentelemetry-instrumentation-annotations:2.11.0',
     'opentracingJdbc':'io.opentracing.contrib:opentracing-jdbc:0.2.15',
     'parquet': 'org.apache.parquet:parquet-avro:1.12.3',
     'parquetHadoop': 'org.apache.parquet:parquet-hadoop:1.13.1',

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/EntityCountsResolver.java

@@ -12,7 +12,7 @@
 import com.linkedin.metadata.service.ViewService;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
-import io.opentelemetry.extension.annotations.WithSpan;
+import io.opentelemetry.instrumentation.annotations.WithSpan;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.CompletableFuture;

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/mutate/MutableTypeBatchResolver.java

@@ -2,13 +2,14 @@
 
 import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
 
-import com.codahale.metrics.Timer;
+import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.concurrency.GraphQLConcurrencyUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.types.BatchMutableType;
 import com.linkedin.metadata.utils.metrics.MetricUtils;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
+import io.datahubproject.metadata.context.OperationContext;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import org.slf4j.Logger;
@@ -33,25 +34,29 @@ public MutableTypeBatchResolver(final BatchMutableType<I, B, T> batchMutableType
 
   @Override
   public CompletableFuture<List<T>> get(DataFetchingEnvironment environment) throws Exception {
+    final QueryContext context = environment.getContext();
+    final OperationContext opContext = context.getOperationContext();
+
     final B[] input =
         bindArgument(environment.getArgument("input"), _batchMutableType.batchInputClass());
 
-    return GraphQLConcurrencyUtils.supplyAsync(
-        () -> {
-          Timer.Context timer = MetricUtils.timer(this.getClass(), "batchMutate").time();
-
-          try {
-            return _batchMutableType.batchUpdate(input, environment.getContext());
-          } catch (AuthorizationException e) {
-            throw e;
-          } catch (Exception e) {
-            _logger.error("Failed to perform batchUpdate", e);
-            throw new IllegalArgumentException(e);
-          } finally {
-            timer.stop();
-          }
-        },
-        this.getClass().getSimpleName(),
-        "get");
+    return opContext.withSpan(
+        "batchMutate",
+        () ->
+            GraphQLConcurrencyUtils.supplyAsync(
+                () -> {
+                  try {
+                    return _batchMutableType.batchUpdate(input, environment.getContext());
+                  } catch (AuthorizationException e) {
+                    throw e;
+                  } catch (Exception e) {
+                    _logger.error("Failed to perform batchUpdate", e);
+                    throw new IllegalArgumentException(e);
+                  }
+                },
+                this.getClass().getSimpleName(),
+                "get"),
+        MetricUtils.DROPWIZARD_METRIC,
+        "true");
   }
 }

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/recommendation/ListRecommendationsResolver.java

@@ -27,7 +27,7 @@
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import io.datahubproject.metadata.context.OperationContext;
-import io.opentelemetry.extension.annotations.WithSpan;
+import io.opentelemetry.instrumentation.annotations.WithSpan;
 import java.net.URISyntaxException;
 import java.util.Collections;
 import java.util.List;

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/search/SearchResolver.java

@@ -19,7 +19,7 @@
 import com.linkedin.metadata.query.SearchFlags;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
-import io.opentelemetry.extension.annotations.WithSpan;
+import io.opentelemetry.instrumentation.annotations.WithSpan;
 import java.util.Collections;
 import java.util.concurrent.CompletableFuture;
 import lombok.RequiredArgsConstructor;

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/UpgradeCliApplication.java

@@ -5,7 +5,9 @@
 import com.linkedin.gms.factory.graphql.GraphQLEngineFactory;
 import com.linkedin.gms.factory.kafka.KafkaEventConsumerFactory;
 import com.linkedin.gms.factory.kafka.SimpleKafkaConsumerFactory;
+import com.linkedin.gms.factory.kafka.trace.KafkaTraceReaderFactory;
 import com.linkedin.gms.factory.telemetry.ScheduledAnalyticsFactory;
+import com.linkedin.gms.factory.trace.TraceServiceFactory;
 import org.springframework.boot.WebApplicationType;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
 import org.springframework.boot.autoconfigure.elasticsearch.ElasticsearchRestClientAutoConfiguration;
@@ -30,7 +32,9 @@
             DataHubAuthorizerFactory.class,
             SimpleKafkaConsumerFactory.class,
             KafkaEventConsumerFactory.class,
-            GraphQLEngineFactory.class
+            GraphQLEngineFactory.class,
+            KafkaTraceReaderFactory.class,
+            TraceServiceFactory.class
           })
     })
 public class UpgradeCliApplication {

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/config/SystemUpdateConfig.java

@@ -195,6 +195,7 @@ protected OperationContext javaSystemOperationContext(
                 .alternateValidation(
                     configurationProvider.getFeatureFlags().isAlternateMCPValidation())
                 .build(),
+            null,
             true);
 
     entityServiceAspectRetriever.setSystemOperationContext(systemOperationContext);

datahub-upgrade/src/main/java/com/linkedin/datahub/upgrade/impl/DefaultUpgradeManager.java

@@ -1,7 +1,6 @@
 package com.linkedin.datahub.upgrade.impl;
 
 import com.codahale.metrics.MetricRegistry;
-import com.codahale.metrics.Timer;
 import com.linkedin.datahub.upgrade.Upgrade;
 import com.linkedin.datahub.upgrade.UpgradeCleanupStep;
 import com.linkedin.datahub.upgrade.UpgradeContext;
@@ -119,44 +118,60 @@
   }
 
   private UpgradeStepResult executeStepInternal(UpgradeContext context, UpgradeStep step) {
-    int retryCount = step.retryCount();
-    UpgradeStepResult result = null;
-    int maxAttempts = retryCount + 1;
-    for (int i = 0; i < maxAttempts; i++) {
-      try (Timer.Context completionTimer =
-          MetricUtils.timer(MetricRegistry.name(step.id(), "completionTime")).time()) {
-        try (Timer.Context executionTimer =
-            MetricUtils.timer(MetricRegistry.name(step.id(), "executionTime")).time()) {
-          result = step.executable().apply(context);
-        }
-
-        if (result == null) {
-          // Failed to even retrieve a result. Create a default failure result.
-          result = new DefaultUpgradeStepResult(step.id(), DataHubUpgradeState.FAILED);
-          context
-              .report()
-              .addLine(String.format("Retrying %s more times...", maxAttempts - (i + 1)));
-          MetricUtils.counter(MetricRegistry.name(step.id(), "retry")).inc();
-        }
-
-        if (DataHubUpgradeState.SUCCEEDED.equals(result.result())) {
-          MetricUtils.counter(MetricRegistry.name(step.id(), "succeeded")).inc();
-          break;
-        }
-      } catch (Exception e) {
-        log.error("Caught exception during attempt {} of Step with id {}", i, step.id(), e);
-        context
-            .report()
-            .addLine(
-                String.format(
-                    "Caught exception during attempt %s of Step with id %s: %s", i, step.id(), e));
-        MetricUtils.counter(MetricRegistry.name(step.id(), "failed")).inc();
-        result = new DefaultUpgradeStepResult(step.id(), DataHubUpgradeState.FAILED);
-        context.report().addLine(String.format("Retrying %s more times...", maxAttempts - (i + 1)));
-      }
-    }
-
-    return result;
+    return context
+        .opContext()
+        .withSpan(
+            "completionTime",
+            () -> {
+              int retryCount = step.retryCount();
+              UpgradeStepResult result = null;
+              int maxAttempts = retryCount + 1;
+              for (int i = 0; i < maxAttempts; i++) {
+                try {
+                  result =
+                      context
+                          .opContext()
+                          .withSpan(
+                              "executionTime",
+                              () -> step.executable().apply(context),
+                              "step.id",
+                              step.id(),
+                              MetricUtils.DROPWIZARD_NAME,
+                              MetricUtils.name(step.id(), "executionTime"));
+
+                  if (result == null) {
+                    // Failed to even retrieve a result. Create a default failure result.
+                    result = new DefaultUpgradeStepResult(step.id(), DataHubUpgradeState.FAILED);
+                    context
+                        .report()
+                        .addLine(String.format("Retrying %s more times...", maxAttempts - (i + 1)));
+                    MetricUtils.counter(MetricRegistry.name(step.id(), "retry")).inc();
+                  }
+
+                  if (DataHubUpgradeState.SUCCEEDED.equals(result.result())) {
+                    MetricUtils.counter(MetricRegistry.name(step.id(), "succeeded")).inc();
+                    break;
+                  }
+                } catch (Exception e) {
+                  log.error(
+                      "Caught exception during attempt {} of Step with id {}", i, step.id(), e);
+                  context
+                      .report()
+                      .addLine(
+                          String.format(
+                              "Caught exception during attempt %s of Step with id %s: %s",
+                              i, step.id(), e));
+                  MetricUtils.counter(MetricRegistry.name(step.id(), "failed")).inc();
+                  result = new DefaultUpgradeStepResult(step.id(), DataHubUpgradeState.FAILED);
+                  context
+                      .report()
+                      .addLine(String.format("Retrying %s more times...", maxAttempts - (i + 1)));
+                }
+              }
+              return result;
+            },
+            MetricUtils.DROPWIZARD_METRIC,
+            "true");
   }
 
   private void executeCleanupInternal(UpgradeContext context, UpgradeResult result) {

entity-registry/src/main/java/com/linkedin/metadata/aspect/ReadItem.java

@@ -64,6 +64,13 @@ static <T> T getAspect(Class<T> clazz, @Nullable RecordTemplate recordTemplate)
   @Nullable
   SystemMetadata getSystemMetadata();
 
+  /**
+   * Set system metadata on the item
+   *
+   * @param systemMetadata
+   */
+  void setSystemMetadata(@Nonnull SystemMetadata systemMetadata);
+
   /**
    * The entity's schema
    *

entity-registry/src/main/java/com/linkedin/metadata/aspect/SystemAspect.java

@@ -6,6 +6,7 @@
 import java.sql.Timestamp;
 import java.util.Optional;
 import javax.annotation.Nonnull;
+import org.apache.commons.lang3.NotImplementedException;
 
 /**
  * An aspect along with system metadata and creation timestamp. Represents an aspect as stored in
@@ -36,4 +37,9 @@
         .map(SystemMetadata::getVersion)
         .map(Long::parseLong);
   }
+
+  @Override
+  default void setSystemMetadata(@Nonnull SystemMetadata systemMetadata) {
+    throw new NotImplementedException();
+  }
 }

entity-registry/src/main/java/com/linkedin/metadata/aspect/batch/MCLItem.java

@@ -37,6 +37,11 @@
     return getMetadataChangeLog().getSystemMetadata();
   }
 
+  @Override
+  default void setSystemMetadata(@Nonnull SystemMetadata systemMetadata) {
+    getMetadataChangeLog().setSystemMetadata(systemMetadata);
+  }
+
   default SystemMetadata getPreviousSystemMetadata() {
     return getMetadataChangeLog().getPreviousSystemMetadata();
   }

entity-registry/src/testFixtures/java/com/linkedin/test/metadata/aspect/batch/TestMCP.java

@@ -116,7 +116,7 @@ public static <T extends RecordTemplate> Set<ChangeMCP> ofOneMCP(
 
   private Urn urn;
   private RecordTemplate recordTemplate;
-  private SystemMetadata systemMetadata;
+  @Setter private SystemMetadata systemMetadata;
   private AuditStamp auditStamp;
   private ChangeType changeType;
   @Nonnull private final EntitySpec entitySpec;