Works

nadav3396 · nadav3396 · commit 4b1538dd4897 · 2024-05-20T13:54:34.000+03:00
diff --git a/pom.xml b/pom.xml
@@ -80,7 +80,7 @@
         <dependency>
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-lambda-java-events</artifactId>
-            <version>2.2.6</version>
+            <version>3.11.5</version>
         </dependency>
         <dependency>
             <groupId>com.amazonaws</groupId>
diff --git a/src/main/java/io/lumigo/core/SpansContainer.java b/src/main/java/io/lumigo/core/SpansContainer.java
@@ -35,6 +35,7 @@
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.clients.producer.internals.ProducerMetadata;
 import org.apache.kafka.clients.producer.internals.Sender;
+import org.apache.kafka.common.serialization.Serializer;
 import org.pmw.tinylog.Logger;
 import software.amazon.awssdk.awscore.AwsResponse;
 import software.amazon.awssdk.core.SdkResponse;
@@ -452,8 +453,8 @@ public void addHttpSpan(
         httpSpans.add(httpSpan);
     }
 
-    public void addKafkaProduceSpan(Long startTime, ProducerMetadata producerMetadata, ProducerRecord<?, ?> record, RecordMetadata recordMetadata, Exception exception) {
-        this.kafkaSpans.add(KafkaSpan.createProduce(this.baseSpan, startTime, producerMetadata, record, recordMetadata, exception));
+    public <K, V>  void addKafkaProduceSpan(Long startTime, Serializer<K>  keySerializer, Serializer<V> valueSerializer, ProducerMetadata producerMetadata, ProducerRecord<K, V> record, RecordMetadata recordMetadata, Exception exception) {
+        this.kafkaSpans.add(KafkaSpan.createProduce(this.baseSpan, startTime, keySerializer, valueSerializer, producerMetadata, record, recordMetadata, exception));
     }
 
     public void addKafkaConsumeSpan(Long startTime, KafkaConsumer<?, ?> consumer, ConsumerMetadata consumerMetadata, ConsumerRecords<?, ?> consumerRecords) {
diff --git a/src/main/java/io/lumigo/core/instrumentation/impl/ApacheKafkaProducerInstrumentation.java b/src/main/java/io/lumigo/core/instrumentation/impl/ApacheKafkaProducerInstrumentation.java
@@ -9,11 +9,11 @@
 import net.bytebuddy.asm.Advice;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.matcher.ElementMatcher;
-import org.apache.kafka.clients.consumer.internals.ConsumerMetadata;
 import org.apache.kafka.clients.producer.Callback;
 import org.apache.kafka.clients.producer.ProducerRecord;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.clients.producer.internals.ProducerMetadata;
+import org.apache.kafka.common.serialization.Serializer;
 import org.pmw.tinylog.Logger;
 
 import java.util.UUID;
@@ -49,13 +49,15 @@ public static class ApacheKafkaProducerAdvice {
         public static final LRUCache<Integer, Boolean> handled = new LRUCache<>(1000);
 
         @Advice.OnMethodEnter
-        public static void methodEnter(
+        public static <K, V> void methodEnter(
                 @Advice.FieldValue("metadata") ProducerMetadata metadata,
-                @Advice.Argument(value = 0, readOnly = false) ProducerRecord<?, ?> record,
+                @Advice.FieldValue("keySerializer") Serializer<K> keySerializer,
+                @Advice.FieldValue("valueSerializer") Serializer<V> valueSerializer,
+                @Advice.Argument(value = 0, readOnly = false) ProducerRecord<K, V> record,
                 @Advice.Argument(value = 1, readOnly = false) Callback callback) {
             try {
                 System.out.println("Inside ApacheKafkaProducerAdvice.methodEnter()");
-                callback = new KafkaProducerCallback(callback, metadata, record, System.currentTimeMillis());
+                callback = new KafkaProducerCallback<>(callback, keySerializer, valueSerializer, metadata, record, System.currentTimeMillis());
 
                 // Try to inject correlation id to the kafka record headers
                 record.headers().add("lumigoMessageId", UUID.randomUUID().toString().substring(0, 10).getBytes());
@@ -65,10 +67,12 @@ public static void methodEnter(
         }
 
         @AllArgsConstructor
-        public static class KafkaProducerCallback implements Callback {
+        public static class KafkaProducerCallback<K, V> implements Callback {
             private final Callback callback;
+            private final Serializer<K> keySerializer;
+            private final Serializer<V> valueSerializer;
             private final ProducerMetadata producerMetadata;
-            private final ProducerRecord<?, ?> record;
+            private final ProducerRecord<K, V> record;
             private final long startTime;
 
             @Override
@@ -80,7 +84,7 @@ public void onCompletion(RecordMetadata recordMetadata, Exception exception) {
                     System.out.println("Inside KafkaProducerCallback.onCompletion()");
 
                     Logger.info("Handling kafka request {} from host {}", record.hashCode());
-                    spansContainer.addKafkaProduceSpan(startTime, producerMetadata, record, recordMetadata, exception);
+                    spansContainer.addKafkaProduceSpan(startTime, keySerializer, valueSerializer, producerMetadata, record, recordMetadata, exception);
                     handled.put(record.hashCode(), true);
                 } catch (Throwable error) {
                     Logger.error(error, "Failed to add kafka span");
diff --git a/src/main/java/io/lumigo/core/parsers/event/EventParserFactory.java b/src/main/java/io/lumigo/core/parsers/event/EventParserFactory.java
@@ -3,6 +3,7 @@
 import com.amazonaws.services.lambda.runtime.events.APIGatewayProxyRequestEvent;
 import com.amazonaws.services.lambda.runtime.events.SNSEvent;
 import com.amazonaws.services.lambda.runtime.events.SQSEvent;
+import com.amazonaws.services.lambda.runtime.events.KafkaEvent;
 import org.pmw.tinylog.Logger;
 
 public interface EventParserFactory {
@@ -14,6 +15,8 @@ static Object parseEvent(Object event) {
                 return new SnsEventParser().parse((SNSEvent) event);
             } else if (event instanceof SQSEvent) {
                 return new SqsEventParser().parse((SQSEvent) event);
+            } else if (event instanceof KafkaEvent) {
+                return new KafkaEventParser().parse((KafkaEvent) event);
             } else {
                 return event;
             }
diff --git a/src/main/java/io/lumigo/core/parsers/event/KafkaEventParser.java b/src/main/java/io/lumigo/core/parsers/event/KafkaEventParser.java
@@ -0,0 +1,41 @@
+package io.lumigo.core.parsers.event;
+
+import com.amazonaws.services.lambda.runtime.events.KafkaEvent;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class KafkaEventParser implements IEventParser<KafkaEvent> {
+    @Override
+    public Object parse(KafkaEvent event) {
+        List<ParsedKafkaEvent.Record> records = new ArrayList<>();
+        for (Map.Entry<String, List<KafkaEvent.KafkaEventRecord>> entry : event.getRecords().entrySet()) {
+            String messageId = null;
+            for (KafkaEvent.KafkaEventRecord record : entry.getValue()) {
+                for (Map<String, byte[]> headers : record.getHeaders()) {
+                    if (headers.containsKey("lumigoMessageId")) {
+                        messageId = new String(headers.get("lumigoMessageId"));
+                        break;
+                    }
+                }
+
+                records.add(
+                        ParsedKafkaEvent.Record.builder()
+                                .topic(record.getTopic())
+                                .partition(record.getPartition())
+                                .offset(record.getOffset())
+                                .key(record.getKey())
+                                .value(record.getValue())
+                                .headers(record.getHeaders())
+                                .messageId(messageId)
+                                .build());
+            }
+        }
+
+        return ParsedKafkaEvent.builder()
+                .bootstrapServers(event.getBootstrapServers())
+                .records(records)
+                .build();
+    }
+}
diff --git a/src/main/java/io/lumigo/core/parsers/event/ParsedKafkaEvent.java b/src/main/java/io/lumigo/core/parsers/event/ParsedKafkaEvent.java
@@ -0,0 +1,29 @@
+package io.lumigo.core.parsers.event;
+
+import lombok.AllArgsConstructor;
+import lombok.Builder;
+import lombok.Data;
+
+import java.util.List;
+import java.util.Map;
+
+@AllArgsConstructor
+@Builder(toBuilder = true)
+@Data(staticConstructor = "of")
+public class ParsedKafkaEvent {
+    private String bootstrapServers;
+    private List<ParsedKafkaEvent.Record> records;
+
+    @AllArgsConstructor
+    @Builder(toBuilder = true)
+    @Data(staticConstructor = "of")
+    public static class Record {
+        private String topic;
+        private Integer partition;
+        private Long offset;
+        private String key;
+        private String value;
+        private List<Map<String, byte[]>> headers;
+        private String messageId;
+    }
+}
diff --git a/src/main/java/io/lumigo/core/utils/AwsSdkV2Utils.java b/src/main/java/io/lumigo/core/utils/AwsSdkV2Utils.java
@@ -10,8 +10,7 @@
 @UtilityClass
 public class AwsSdkV2Utils {
 
-    public String calculateItemHash(
-            Map<String, software.amazon.awssdk.services.dynamodb.model.AttributeValue> item) {
+    public String calculateItemHash(Map<String, AttributeValue> item) {
         Map<String, Object> simpleMap = AwsSdkV2Utils.convertAttributeMapToSimpleMap(item);
         return StringUtils.buildMd5Hash(JsonUtils.getObjectAsJsonString(simpleMap));
     }
diff --git a/src/main/java/io/lumigo/core/utils/AwsUtils.java b/src/main/java/io/lumigo/core/utils/AwsUtils.java
@@ -1,20 +1,22 @@
 package io.lumigo.core.utils;
 
-import static io.lumigo.core.utils.StringUtils.dynamodbItemToHash;
-
 import com.amazonaws.services.lambda.runtime.events.*;
+import com.amazonaws.services.lambda.runtime.events.models.dynamodb.AttributeValue;
 import com.fasterxml.jackson.annotation.JsonInclude;
+import io.lumigo.core.parsers.event.KafkaEventParser;
+import io.lumigo.core.parsers.event.ParsedKafkaEvent;
 import io.lumigo.models.Span;
-import java.util.Collections;
-import java.util.List;
-import java.util.Objects;
+
+import java.util.*;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import lombok.Data;
 import lombok.NoArgsConstructor;
 import org.pmw.tinylog.Logger;
 
+import static io.lumigo.core.utils.StringUtils.buildMd5Hash;
+
 public class AwsUtils {
 
     public static final String COLD_START_KEY = "LUMIGO_COLD_START_KEY";
@@ -155,6 +157,10 @@ public static TriggeredBy extractTriggeredByFromEvent(Object event) {
                 triggeredBy.setTriggeredBy("lex");
             } else if (event instanceof CognitoEvent) {
                 triggeredBy.setTriggeredBy("cognito");
+            } else if (event instanceof KafkaEvent) {
+                triggeredBy.setTriggeredBy("kafka");
+                Object parsed = new KafkaEventParser().parse((KafkaEvent) event);
+                triggeredBy.setMessageIds(((ParsedKafkaEvent) parsed).getRecords().stream().map(ParsedKafkaEvent.Record::getMessageId).collect(Collectors.toList()));
             } else {
                 Logger.info(
                         "Failed to found relevant triggered by found for event {} ",
@@ -282,10 +288,44 @@ private static String extractMessageIdFromDynamodbRecord(
             DynamodbEvent.DynamodbStreamRecord record) {
         if (record.getEventName() == null) return null;
         if (record.getEventName().equals("INSERT")) {
-            return dynamodbItemToHash(record.getDynamodb().getNewImage());
+            return calculateItemHash(record.getDynamodb().getNewImage());
         } else if (record.getEventName().equals("MODIFY")
                 || record.getEventName().equals("REMOVE")) {
-            return dynamodbItemToHash(record.getDynamodb().getKeys());
+            return calculateItemHash(record.getDynamodb().getKeys());
+        }
+        return null;
+    }
+
+    private static String calculateItemHash(Map<String, AttributeValue> item) {
+        Map<String, Object> simpleMap = convertAttributeMapToSimpleMap(item);
+        return StringUtils.buildMd5Hash(JsonUtils.getObjectAsJsonString(simpleMap));
+    }
+
+    private static Map<String, Object> convertAttributeMapToSimpleMap(
+            Map<String, AttributeValue> attributeValueMap) {
+        Map<String, Object> simpleMap = new HashMap<>();
+        attributeValueMap.forEach(
+                (key, value) -> simpleMap.put(key, attributeValueToObject(value)));
+        return simpleMap;
+    }
+
+    private static Object attributeValueToObject(AttributeValue value) {
+        if (value == null) {
+            return null;
+        } else if (value.getS() != null) {
+            return value.getS();
+        } else if (value.getN() != null) {
+            return value.getN();
+        } else if (value.getBOOL() != null) {
+            return value.getBOOL();
+        } else if (value.getL() != null && !value.getL().isEmpty()) {
+            List<Object> list = new ArrayList<>();
+            for (AttributeValue v : value.getL()) {
+                list.add(attributeValueToObject(v));
+            }
+            return list;
+        } else if (value.getM() != null && !value.getM().isEmpty()) {
+            return convertAttributeMapToSimpleMap(value.getM());
         }
         return null;
     }
diff --git a/src/main/java/io/lumigo/models/KafkaSpan.java b/src/main/java/io/lumigo/models/KafkaSpan.java
@@ -12,6 +12,7 @@
 import org.apache.kafka.clients.producer.internals.ProducerMetadata;
 import org.apache.kafka.common.header.Header;
 import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.serialization.Serializer;
 
 import java.util.*;
 import java.util.stream.Collectors;
@@ -81,8 +82,8 @@ public static class KafkaProducerInfo implements KafkaInfo {
     @Builder(toBuilder = true)
     @Data(staticConstructor = "of")
     public static class KafkaProducerRecord {
-        private String key;
-        private String value;
+        private byte[] key;
+        private byte[] value;
         private Map<String, byte[]> headers;
     }
 
@@ -127,12 +128,12 @@ public static class KafkaConsumerRecord {
         private Map<String, byte[]> headers;
     }
 
-    public static KafkaSpan createProduce(Span baseSpan, Long startTime, ProducerMetadata producerMetadata, ProducerRecord<?, ?> record, RecordMetadata recordMetadata, Exception exception) {
+    public static <K, V> KafkaSpan createProduce(Span baseSpan, Long startTime, Serializer<K> keySerializer, Serializer<V> valueSerializer, ProducerMetadata producerMetadata, ProducerRecord<K, V> record, RecordMetadata recordMetadata, Exception exception) {
         List<String> bootstrapServers = producerMetadata.fetch().nodes().stream().map(node -> node.host() + ":" + node.port()).collect(Collectors.toList());
         String topic = record.topic();
         KafkaProducerRecord producerRecord = KafkaProducerRecord.builder()
-                .key(record.key().toString()) // TODO USE serializer
-                .value(record.value().toString()) // TODO USE serializer
+                .key(keySerializer.serialize(record.topic(), record.headers(), record.key()))
+                .value(valueSerializer.serialize(record.topic(), record.headers(), record.value()))
                 .headers(extractHeaders(record.headers()))
                 .build();
 
diff --git a/src/test/java/infa/AwsLambdaEventGenerator.java b/src/test/java/infa/AwsLambdaEventGenerator.java
@@ -25,7 +25,7 @@ public S3Event s3Event() {
         when(s3EventNotificationRecord.getS3()).thenReturn(s3Entity);
         when(s3EventNotificationRecord.getS3().getBucket()).thenReturn(bucketEntity);
         when(s3EventNotificationRecord.getS3().getBucket().getArn()).thenReturn("s3-arn");
-        when(s3Event.getRecords()).thenReturn(Collections.singletonList(s3EventNotificationRecord));
+//        when(s3Event.getRecords()).thenReturn(Collections.singletonList(s3EventNotificationRecord));
         return s3Event;
     }
 
@@ -55,12 +55,12 @@ public DynamodbEvent dynamodbEvent() {
         when(dynamodbEvent.getRecords()).thenReturn(List.of(record, record2));
         when(record.getEventSourceARN()).thenReturn("dynamodb-arn");
         when(record.getEventName()).thenReturn("INSERT");
-        when(record.getDynamodb()).thenReturn(streamRecord);
+//        when(record.getDynamodb()).thenReturn(streamRecord);
         when(streamRecord.getApproximateCreationDateTime()).thenReturn(new Date(769554000));
         when(streamRecord.getNewImage())
                 .thenReturn(Collections.singletonMap("k", new AttributeValue("v")));
         when(record2.getEventName()).thenReturn("MODIFY");
-        when(record2.getDynamodb()).thenReturn(streamRecord2);
+//        when(record2.getDynamodb()).thenReturn(streamRecord2);
         when(streamRecord2.getKeys())
                 .thenReturn(Collections.singletonMap("k2", new AttributeValue("v2")));
         return dynamodbEvent;

Original file line number	Diff line number	Diff line change
`@@ -10,8 +10,7 @@`
`10`	`10`	`@UtilityClass`
`11`	`11`	`public class AwsSdkV2Utils {`
`12`	`12`
`13`		`- public String calculateItemHash(`
`14`		`- Map<String, software.amazon.awssdk.services.dynamodb.model.AttributeValue> item) {`
	`13`	`+ public String calculateItemHash(Map<String, AttributeValue> item) {`
`15`	`14`	`Map<String, Object> simpleMap = AwsSdkV2Utils.convertAttributeMapToSimpleMap(item);`
`16`	`15`	`return StringUtils.buildMd5Hash(JsonUtils.getObjectAsJsonString(simpleMap));`
`17`	`16`	`}`