Issue 1958 - Instrument for AWS X-Ray with OpenTelemetry

example/full/instance.properties

@@ -129,6 +129,9 @@ sleeper.table.batching.lambdas.memory=1024
 # create compaction jobs, run garbage collection, perform partition splitting.
 sleeper.table.batching.lambdas.timeout.seconds=60
 
+# This specifies whether OpenTelemetry tracing is enabled.
+sleeper.opentelemetry.tracing.enabled=true
+
 
 ## The following properties relate to standard ingest.
 

java/athena/pom.xml

@@ -27,6 +27,27 @@
     <artifactId>athena</artifactId>
 
     <dependencies>
+        <dependency>
+            <groupId>com.fasterxml.jackson.core</groupId>
+            <artifactId>jackson-annotations</artifactId>
+            <version>${jackson.version}</version>
+        </dependency>
+        <!-- AWS Dependencies -->
+        <dependency>
+            <groupId>com.amazonaws</groupId>
+            <artifactId>aws-athena-federation-sdk</artifactId>
+            <version>${athena.version}</version>
+        </dependency>
+        <!-- Arrow dependencies -->
+        <dependency>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>arrow-vector</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.arrow</groupId>
+            <artifactId>arrow-memory-unsafe</artifactId>
+        </dependency>
+        <!-- Sleeper Dependencies -->
         <dependency>
             <groupId>sleeper</groupId>
             <artifactId>configuration</artifactId>
@@ -47,25 +68,6 @@
             <artifactId>statestore</artifactId>
             <version>${project.parent.version}</version>
         </dependency>
-        <dependency>
-            <groupId>com.fasterxml.jackson.core</groupId>
-            <artifactId>jackson-annotations</artifactId>
-            <version>${jackson.version}</version>
-        </dependency>
-        <dependency>
-            <groupId>com.amazonaws</groupId>
-            <artifactId>aws-athena-federation-sdk</artifactId>
-            <version>${athena.version}</version>
-        </dependency>
-        <!-- Arrow dependencies -->
-        <dependency>
-            <groupId>org.apache.arrow</groupId>
-            <artifactId>arrow-vector</artifactId>
-        </dependency>
-        <dependency>
-            <groupId>org.apache.arrow</groupId>
-            <artifactId>arrow-memory-unsafe</artifactId>
-        </dependency>
         <!-- Test dependencies -->
         <dependency>
             <groupId>org.testcontainers</groupId>

java/bulk-import/bulk-import-runner/docker/eks-native/Dockerfile

@@ -17,6 +17,8 @@ ARG BUILDER_IMAGE_TAG=3.8-openjdk-8-slim
 ARG BASE_IMAGE_NAME=amazoncorretto
 ARG BASE_IMAGE_TAG=11
 
+ARG MODE=no_tracing
+
 ARG SPARK_VERSION=3.4.1
 ARG HADOOP_VERSION=3.3.3
 ARG SPARK_DOWNLOAD_FILENAME=spark-${SPARK_VERSION}-bin-hadoop3
@@ -57,7 +59,14 @@ RUN echo "Before slimming: $(du -sh /opt/${SPARK_DIRNAME})" && \
 # Add workdir
 RUN mkdir /opt/${SPARK_DIRNAME}/workdir
 
-FROM ${BASE_IMAGE_NAME}:${BASE_IMAGE_TAG}
+FROM ${BASE_IMAGE_NAME}:${BASE_IMAGE_TAG} as base_no_tracing
+
+FROM base_no_tracing as base_tracing
+
+ONBUILD ADD https://github.com/aws-observability/aws-otel-java-instrumentation/releases/latest/download/aws-opentelemetry-agent.jar /opt/aws-opentelemetry-agent.jar
+ONBUILD ENV JAVA_TOOL_OPTIONS=-javaagent:/opt/aws-opentelemetry-agent.jar
+
+FROM base_${MODE}
 
 ARG SPARK_VERSION
 ARG HADOOP_VERSION

java/bulk-import/bulk-import-runner/docker/eks/Dockerfile

@@ -11,8 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+ARG MODE=no_tracing
 
-FROM apache/spark:3.4.1-scala2.12-java11-ubuntu
+FROM apache/spark:3.4.1-scala2.12-java11-ubuntu as build_no_tracing
+
+FROM build_no_tracing as build_tracing
+
+ONBUILD ADD https://github.com/aws-observability/aws-otel-java-instrumentation/releases/latest/download/aws-opentelemetry-agent.jar /opt/aws-opentelemetry-agent.jar
+ONBUILD ENV JAVA_TOOL_OPTIONS=-javaagent:/opt/aws-opentelemetry-agent.jar
+
+FROM build_${MODE}
 
 ENV PATH="$PATH:/opt/spark/bin"
 USER root

java/bulk-import/bulk-import-runner/docker/emr-serverless/Dockerfile

@@ -11,8 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+ARG MODE=no_tracing
 
-FROM public.ecr.aws/emr-serverless/spark/emr-6.13.0:latest
+FROM public.ecr.aws/emr-serverless/spark/emr-6.13.0:latest as build_no_tracing
+
+FROM build_no_tracing as build_tracing
+
+ONBUILD ADD https://github.com/aws-observability/aws-otel-java-instrumentation/releases/latest/download/aws-opentelemetry-agent.jar /opt/aws-opentelemetry-agent.jar
+ONBUILD ENV JAVA_TOOL_OPTIONS=-javaagent:/opt/aws-opentelemetry-agent.jar
+
+FROM build_${MODE}
 
 USER root
 
@@ -24,4 +32,4 @@ RUN yum update -y && \
 COPY ./bulk-import-runner.jar /workdir
 
 # EMR Severless will run the image as hadoop
-USER hadoop:hadoop
+USER hadoop:hadoop

java/cdk/src/main/java/sleeper/cdk/SleeperCdkApp.java

@@ -64,7 +64,6 @@
 
 import static sleeper.configuration.properties.instance.CommonProperty.ACCOUNT;
 import static sleeper.configuration.properties.instance.CommonProperty.ID;
-import static sleeper.configuration.properties.instance.CommonProperty.JARS_BUCKET;
 import static sleeper.configuration.properties.instance.CommonProperty.OPTIONAL_STACKS;
 import static sleeper.configuration.properties.instance.CommonProperty.REGION;
 
@@ -322,7 +321,7 @@ public static void main(String[] args) {
                 .account(instanceProperties.get(ACCOUNT))
                 .region(instanceProperties.get(REGION))
                 .build();
-        BuiltJars jars = new BuiltJars(AmazonS3ClientBuilder.defaultClient(), instanceProperties.get(JARS_BUCKET));
+        BuiltJars jars = new BuiltJars(AmazonS3ClientBuilder.defaultClient(), instanceProperties);
 
         new SleeperCdkApp(app, id, StackProps.builder()
                 .stackName(id)

java/cdk/src/main/java/sleeper/cdk/TracingUtils.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2022-2024 Crown Copyright
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package sleeper.cdk;
+
+import software.amazon.awscdk.services.lambda.Tracing;
+
+import sleeper.configuration.properties.instance.InstanceProperties;
+
+import static sleeper.configuration.properties.instance.CommonProperty.TRACING_ENABLED;
+
+public class TracingUtils {
+
+    private TracingUtils() {
+    }
+
+    public static Tracing active(InstanceProperties properties) {
+        if (properties.getBoolean(TRACING_ENABLED)) {
+            return Tracing.ACTIVE;
+        } else {
+            return Tracing.DISABLED;
+        }
+    }
+
+    public static Tracing passThrough(InstanceProperties properties) {
+        if (properties.getBoolean(TRACING_ENABLED)) {
+            return Tracing.PASS_THROUGH;
+        } else {
+            return Tracing.DISABLED;
+        }
+    }
+}

java/cdk/src/main/java/sleeper/cdk/Utils.java

@@ -76,8 +76,10 @@ private Utils() {
 
     public static Map<String, String> createDefaultEnvironment(InstanceProperties instanceProperties) {
         Map<String, String> environmentVariables = new HashMap<>();
-        environmentVariables.put(CONFIG_BUCKET.toEnvironmentVariable(),
-                instanceProperties.get(CONFIG_BUCKET));
+        if (instanceProperties.isSet(CONFIG_BUCKET)) {
+            environmentVariables.put(CONFIG_BUCKET.toEnvironmentVariable(),
+                    instanceProperties.get(CONFIG_BUCKET));
+        }
 
         environmentVariables.put("JAVA_TOOL_OPTIONS", createToolOptions(instanceProperties,
                 LOGGING_LEVEL,

java/cdk/src/main/java/sleeper/cdk/jars/BuiltJars.java

@@ -18,25 +18,44 @@
 import com.amazonaws.services.s3.AmazonS3;
 import software.amazon.awscdk.services.s3.IBucket;
 
+import sleeper.configuration.properties.instance.InstanceProperties;
+
+import java.util.HashMap;
+import java.util.Map;
+
+import static sleeper.configuration.properties.instance.CommonProperty.JARS_BUCKET;
+
 public class BuiltJars {
 
     private final AmazonS3 s3;
     private final String bucketName;
+    private final LambdaBuilder.Configuration globalConfig;
+    private final Map<String, String> jarFilenameToVersionId = new HashMap<>();
+
+    public BuiltJars(AmazonS3 s3, InstanceProperties instanceProperties) {
+        this(s3, instanceProperties.get(JARS_BUCKET), new GlobalLambdaConfiguration(instanceProperties));
+    }
 
     public BuiltJars(AmazonS3 s3, String bucketName) {
+        this(s3, bucketName, LambdaBuilder.Configuration.none());
+    }
+
+    private BuiltJars(AmazonS3 s3, String bucketName, LambdaBuilder.Configuration globalConfig) {
         this.s3 = s3;
         this.bucketName = bucketName;
+        this.globalConfig = globalConfig;
     }
 
     public String bucketName() {
         return bucketName;
     }
 
     public LambdaCode lambdaCode(BuiltJar jar, IBucket bucketConstruct) {
-        return new LambdaCode(bucketConstruct, jar.getFileName(), getLatestVersionId(jar));
+        return new LambdaCode(bucketConstruct, jar.getFileName(), getLatestVersionId(jar), globalConfig);
     }
 
     public String getLatestVersionId(BuiltJar jar) {
-        return s3.getObjectMetadata(bucketName, jar.getFileName()).getVersionId();
+        return jarFilenameToVersionId.computeIfAbsent(jar.getFileName(),
+                filename -> s3.getObjectMetadata(bucketName, filename).getVersionId());
     }
 }

java/cdk/src/main/java/sleeper/cdk/jars/GlobalLambdaConfiguration.java

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2022-2024 Crown Copyright
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package sleeper.cdk.jars;
+
+import software.amazon.awscdk.services.lambda.LayerVersion;
+import software.constructs.Construct;
+
+import sleeper.cdk.Utils;
+import sleeper.configuration.properties.instance.InstanceProperties;
+
+import static sleeper.configuration.properties.instance.CommonProperty.REGION;
+import static sleeper.configuration.properties.instance.CommonProperty.TRACING_ENABLED;
+
+public class GlobalLambdaConfiguration implements LambdaBuilder.Configuration {
+
+    private final InstanceProperties instanceProperties;
+
+    public GlobalLambdaConfiguration(InstanceProperties instanceProperties) {
+        this.instanceProperties = instanceProperties;
+    }
+
+    @Override
+    public void apply(Construct scope, String functionId, LambdaBuilder builder) {
+        builder.environmentVariables(Utils.createDefaultEnvironment(instanceProperties));
+        if (instanceProperties.getBoolean(TRACING_ENABLED)) {
+            String region = instanceProperties.get(REGION);
+            String arn = "arn:aws:lambda:" + region + ":901920570463:layer:aws-otel-java-agent-amd64-ver-1-32-0:1";
+            builder.layer(LayerVersion.fromLayerVersionArn(scope, functionId + "Tracing", arn));
+            builder.environmentVariable("AWS_LAMBDA_EXEC_WRAPPER", "/opt/otel-handler");
+        }
+    }
+
+}

java/cdk/src/main/java/sleeper/cdk/jars/LambdaBuilder.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2022-2024 Crown Copyright
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package sleeper.cdk.jars;
+
+import software.amazon.awscdk.services.lambda.Function;
+import software.amazon.awscdk.services.lambda.ILayerVersion;
+import software.amazon.awscdk.services.lambda.IVersion;
+import software.constructs.Construct;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Consumer;
+
+public class LambdaBuilder {
+    private final Function.Builder builder;
+    private final Map<String, String> environment = new HashMap<>();
+    private final List<ILayerVersion> layers = new ArrayList<>();
+
+    public LambdaBuilder(Function.Builder builder) {
+        this.builder = builder;
+    }
+
+    public LambdaBuilder config(Consumer<Function.Builder> config) {
+        config.accept(builder);
+        return this;
+    }
+
+    public LambdaBuilder environmentVariable(String key, String value) {
+        this.environment.put(key, value);
+        return this;
+    }
+
+    public LambdaBuilder environmentVariables(Map<String, String> environment) {
+        this.environment.putAll(environment);
+        return this;
+    }
+
+    public LambdaBuilder layer(ILayerVersion layer) {
+        layers.add(layer);
+        return this;
+    }
+
+    public IVersion build() {
+        Function function = builder.environment(environment).layers(layers).build();
+
+        // This is needed to tell the CDK to update the functions with new code when it changes in the jars bucket.
+        // See the following:
+        // https://www.define.run/posts/cdk-not-updating-lambda/
+        // https://awsteele.com/blog/2020/12/24/aws-lambda-latest-is-dangerous.html
+        // https://docs.aws.amazon.com/cdk/api/v1/java/software/amazon/awscdk/services/lambda/Version.html
+        return function.getCurrentVersion();
+    }
+
+    public interface Configuration {
+        void apply(Construct scope, String functionId, LambdaBuilder builder);
+
+        static Configuration none() {
+            return (scope, functionId, builder) -> {
+            };
+        }
+    }
+}