Add json serialization for DefaultTraceGroupFields (#5991)

* Add json serialization for DefaultTraceGroupFields

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

* Fixed code coverage by adding more tests

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

* Addressed review comments

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

---------

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>

Author: kkondaka

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/trace/DefaultTraceGroupFields.java

@@ -5,27 +5,65 @@
 
 package org.opensearch.dataprepper.model.trace;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+
+import java.io.ObjectOutputStream;
+import java.io.ObjectInputStream;
+
+import java.io.IOException;
+import java.io.Serializable;
+import java.util.Objects;
+
 /**
  * The default implementation of {@link TraceGroupFields}, the attributes associated with an entire trace.
  *
  * @since 1.2
  */
-public class DefaultTraceGroupFields implements TraceGroupFields {
+public class DefaultTraceGroupFields implements TraceGroupFields, Serializable {
 
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
     private String endTime;
     private Long durationInNanos;
     private Integer statusCode;
 
     // required for serialization
     DefaultTraceGroupFields() {}
 
+    DefaultTraceGroupFields(final String endTime, final Long durationInNanos, final Integer statusCode) {
+        this.endTime = endTime;
+        this.durationInNanos = durationInNanos;
+        this.statusCode = statusCode;
+    }
+
     private DefaultTraceGroupFields(final Builder builder) {
 
         this.endTime = builder.endTime;
         this.durationInNanos = builder.durationInNanos;
         this.statusCode = builder.statusCode;
     }
 
+    private void writeObject(ObjectOutputStream out) throws IOException {
+        String json = OBJECT_MAPPER.writeValueAsString(this);
+        out.writeUTF(json);
+    }
+
+    private void readObject(ObjectInputStream in) throws IOException, ClassNotFoundException {
+        String json = in.readUTF();
+        DefaultTraceGroupFields temp = OBJECT_MAPPER.readValue(json, DefaultTraceGroupFields.class);
+        this.endTime = temp.endTime;
+        this.durationInNanos = temp.durationInNanos;
+        this.statusCode = temp.statusCode;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        DefaultTraceGroupFields that = (DefaultTraceGroupFields) o;
+        return Objects.equals(endTime, that.endTime) && Objects.equals(durationInNanos, that.durationInNanos) && Objects.equals(statusCode, that.statusCode);
+    }
+
     @Override
     public String getEndTime() {
         return endTime;
@@ -41,15 +79,6 @@ public Integer getStatusCode() {
         return statusCode;
     }
 
-    @Override
-    public boolean equals(Object other) {
-        if (other instanceof DefaultTraceGroupFields) {
-            final DefaultTraceGroupFields o = (DefaultTraceGroupFields) other;
-            return (endTime.equals(o.endTime) && durationInNanos.equals(o.durationInNanos) && statusCode.equals(o.statusCode));
-        }
-        return false;
-    }
-
     public static Builder builder() {
         return new Builder();
     }

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/trace/JacksonSpan.java

@@ -232,7 +232,13 @@ public void setTraceGroup(final String traceGroup) {
 
     @Override
     public void setTraceGroupFields(final TraceGroupFields traceGroupFields) {
-        this.put(TRACE_GROUP_FIELDS_KEY, traceGroupFields);
+        EventMetadata metadata = getMetadata();
+        Object oldTraceGroupFields = metadata.getAttribute(TRACE_GROUP_FIELDS_KEY);
+        if (oldTraceGroupFields != null) {
+            metadata.setAttribute(TRACE_GROUP_FIELDS_KEY, traceGroupFields);
+        } else {
+            this.put(TRACE_GROUP_FIELDS_KEY, traceGroupFields);
+        }
     }
 
     public static Builder builder() {

data-prepper-api/src/test/java/org/opensearch/dataprepper/model/trace/DefaultTraceGroupFieldsTest.java

@@ -9,8 +9,14 @@
 import org.junit.jupiter.api.Test;
 import org.opensearch.dataprepper.model.event.TestObject;
 
+import java.io.ByteArrayOutputStream;
+import java.io.ByteArrayInputStream;
+import java.io.ObjectOutputStream;
+import java.io.ObjectInputStream;
 import java.util.UUID;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
+
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
@@ -19,6 +25,7 @@
 
 public class DefaultTraceGroupFieldsTest {
 
+    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
     private static final String TEST_END_TIME = UUID.randomUUID().toString();
     private static final Long TEST_DURATION = 123L;
     private static final Integer TEST_STATUS_CODE = 200;
@@ -70,6 +77,15 @@ public void testEquals_withDifferentDurationInNanos() {
         assertThat(defaultTraceGroupFields, is(not(equalTo(traceGroupFields))));
     }
 
+    @Test
+    public void testEquals_withDifferentValues() {
+        final DefaultTraceGroupFields traceGroupFields =
+            new DefaultTraceGroupFields(TEST_END_TIME, TEST_DURATION - 1, TEST_STATUS_CODE);
+        assertThat(traceGroupFields, is(not(equalTo(null))));
+        Integer tmp = 5;
+        assertThat(traceGroupFields, is(not(equalTo(tmp))));
+    }
+
     @Test
     public void testEquals_withDifferentStatusCode() {
         final DefaultTraceGroupFields traceGroupFields = DefaultTraceGroupFields.builder()
@@ -111,4 +127,25 @@ public void testBuilder_withAllParameters_createsTraceGroupFields() {
 
         assertThat(result, is(notNullValue()));
     }
+
+    @Test
+    public void testSerialization() throws Exception {
+        final DefaultTraceGroupFields result = DefaultTraceGroupFields.builder()
+                .withDurationInNanos(TEST_DURATION)
+                .withStatusCode(TEST_STATUS_CODE)
+                .withEndTime(TEST_END_TIME)
+                .build();
+        ByteArrayOutputStream baos = new ByteArrayOutputStream();
+        try (ObjectOutputStream oos = new ObjectOutputStream(baos)) {
+            oos.writeObject(result);
+        }
+
+        byte[] data = baos.toByteArray();
+    
+        try (ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(data))) {
+            DefaultTraceGroupFields result2 = (DefaultTraceGroupFields) ois.readObject();
+            assertThat(result, equalTo(result2));
+        }
+    }
+
 }

data-prepper-api/src/test/java/org/opensearch/dataprepper/model/trace/JacksonSpanTest.java

@@ -297,6 +297,20 @@ public void testSetAndGetTraceGroupFields() {
         assertThat(traceGroupFields, is(equalTo(testTraceGroupFields)));
     }
 
+    @Test
+    public void testSetAndGetTraceGroupFieldsInMetadata() {
+        final TraceGroupFields testTraceGroupFields = DefaultTraceGroupFields.builder()
+                .withDurationInNanos(200L)
+                .withStatusCode(404)
+                .withEndTime("Different end time")
+                .build();
+        jacksonSpan.getMetadata().setAttribute(JacksonSpan.TRACE_GROUP_FIELDS_KEY, new DefaultTraceGroupFields());
+        jacksonSpan.setTraceGroupFields(testTraceGroupFields);
+        final TraceGroupFields traceGroupFields = jacksonSpan.getTraceGroupFields();
+        assertThat(traceGroupFields, is(equalTo(traceGroupFields)));
+        assertThat(traceGroupFields, is(equalTo(testTraceGroupFields)));
+    }
+
     @Test
     public void testToJsonStringAllParameters() throws JsonProcessingException, JSONException {
         final String jsonResult = jacksonSpan.toJsonString();