feat: support codecbuffer in SnapshotDiffJobCodec

Author: YutaLin

hadoop-ozone/common/src/main/java/org/apache/hadoop/ozone/om/helpers/SnapshotDiffJob.java

@@ -21,10 +21,14 @@
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.google.protobuf.InvalidProtocolBufferException;
+import jakarta.annotation.Nonnull;
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.Objects;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.hdds.utils.db.Codec;
+import org.apache.hadoop.hdds.utils.db.CodecBuffer;
+import org.apache.hadoop.hdds.utils.db.CodecException;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.SnapshotDiffJobProto;
 import org.apache.hadoop.ozone.snapshot.SnapshotDiffResponse.JobStatus;
 import org.apache.hadoop.ozone.snapshot.SnapshotDiffResponse.SubStatus;
@@ -310,6 +314,51 @@ private static final class SnapshotDiffJobCodec
         .setSerializationInclusion(JsonInclude.Include.NON_NULL)
         .configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
 
+    @Override
+    public boolean supportCodecBuffer() {
+      return true;
+    }
+
+    @Override
+    public CodecBuffer toCodecBuffer(@Nonnull SnapshotDiffJob object,
+        CodecBuffer.Allocator allocator) throws CodecException {
+
+      SnapshotDiffJobProto proto = object.toProtoBuf();
+      final int size = proto.getSerializedSize();
+      final CodecBuffer buffer = allocator.apply(size);
+
+      buffer.put(out -> {
+        try {
+          proto.writeTo(out);
+          return size;
+        } catch (IOException e) {
+          throw new IllegalStateException("Failed to write protobuf to buffer", e);
+        }
+      });
+      return buffer;
+    }
+
+    @Override
+    public SnapshotDiffJob fromCodecBuffer(@Nonnull CodecBuffer buffer)
+        throws CodecException {
+      // Direct protobuf parsing from InputStream to avoid byte array copy
+      try (java.io.InputStream in = buffer.getInputStream()) {
+        SnapshotDiffJobProto proto = SnapshotDiffJobProto.parseFrom(in);
+        return SnapshotDiffJob.getFromProtoBuf(proto);
+      } catch (InvalidProtocolBufferException e) {
+        ByteBuffer bb = buffer.asReadOnlyByteBuffer();
+        byte[] data = new byte[bb.remaining()];
+        bb.get(data);
+        try {
+          return MAPPER.readValue(data, SnapshotDiffJob.class);
+        } catch (IOException ex) {
+          throw new CodecException("Failed to deserialize SnapshotDiffJob from JSON", ex);
+        }
+      } catch (IOException e) {
+        throw new CodecException("Failed to read from CodecBuffer", e);
+      }
+    }
+
     @Override
     public Class<SnapshotDiffJob> getTypeClass() {
       return SnapshotDiffJob.class;

hadoop-ozone/common/src/test/java/org/apache/hadoop/ozone/om/helpers/TestOmSnapshotDiffJobCodec.java

@@ -18,8 +18,11 @@
 package org.apache.hadoop.ozone.om.helpers;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.nio.ByteBuffer;
 import org.apache.hadoop.hdds.utils.db.Codec;
+import org.apache.hadoop.hdds.utils.db.CodecBuffer;
 import org.apache.hadoop.ozone.snapshot.SnapshotDiffResponse.JobStatus;
 import org.apache.hadoop.ozone.snapshot.SnapshotDiffResponse.SubStatus;
 import org.junit.jupiter.api.Test;
@@ -70,4 +73,88 @@ public void testOldJsonSerializedDataCanBeReadByNewCodec() throws Exception {
 
     assertEquals(0.0, parsed.getKeysProcessedPct());
   }
+
+  @Test
+  public void testCodecBufferSupport() throws Exception {
+    assertTrue(newCodec.supportCodecBuffer());
+
+    SnapshotDiffJob original = new SnapshotDiffJob(
+        System.currentTimeMillis(),
+        "test-job-buffer",
+        JobStatus.DONE,
+        "testVol",
+        "testBucket",
+        "fromSnap",
+        "toSnap",
+        false,
+        true,
+        500L,
+        SubStatus.OBJECT_ID_MAP_GEN_FSO,
+        75.5
+    );
+
+    // Test with direct allocator
+    try (CodecBuffer buffer = newCodec.toCodecBuffer(original, CodecBuffer.Allocator.getDirect())) {
+      SnapshotDiffJob decoded = newCodec.fromCodecBuffer(buffer);
+      assertSnapshotDiffJobEquals(original, decoded);
+    }
+
+    // Test with heap allocator
+    try (CodecBuffer buffer = newCodec.toCodecBuffer(original, CodecBuffer.Allocator.getHeap())) {
+      SnapshotDiffJob decoded = newCodec.fromCodecBuffer(buffer);
+      assertSnapshotDiffJobEquals(original, decoded);
+    }
+  }
+
+  @Test
+  public void testCodecBufferBackwardCompatibility() throws Exception {
+
+    SnapshotDiffJob original = new SnapshotDiffJob(
+        987654321L,
+        "compat-job",
+        JobStatus.FAILED,
+        "volX",
+        "buckY",
+        "oldSnap",
+        "newSnap",
+        true,
+        true,
+        0L,
+        null,
+        0.0
+    );
+    original.setReason("Test failure reason");
+
+
+    byte[] jsonData = oldCodec.toPersistedFormatImpl(original);
+
+    // Create a CodecBuffer from the JSON data
+    // This simulates reading old format data from storage
+    try (CodecBuffer buffer = CodecBuffer.Allocator.getHeap().apply(jsonData.length)) {
+      buffer.put(ByteBuffer.wrap(jsonData));
+
+      // The new codec should handle JSON fallback in fromCodecBuffer
+      SnapshotDiffJob decoded = newCodec.fromCodecBuffer(buffer);
+
+      assertEquals(original.getJobId(), decoded.getJobId());
+      assertEquals(original.getStatus(), decoded.getStatus());
+      assertEquals(original.getReason(), decoded.getReason());
+    }
+  }
+
+  private void assertSnapshotDiffJobEquals(SnapshotDiffJob expected, SnapshotDiffJob actual) {
+    assertEquals(expected.getCreationTime(), actual.getCreationTime());
+    assertEquals(expected.getJobId(), actual.getJobId());
+    assertEquals(expected.getStatus(), actual.getStatus());
+    assertEquals(expected.getVolume(), actual.getVolume());
+    assertEquals(expected.getBucket(), actual.getBucket());
+    assertEquals(expected.getFromSnapshot(), actual.getFromSnapshot());
+    assertEquals(expected.getToSnapshot(), actual.getToSnapshot());
+    assertEquals(expected.isForceFullDiff(), actual.isForceFullDiff());
+    assertEquals(expected.isNativeDiffDisabled(), actual.isNativeDiffDisabled());
+    assertEquals(expected.getTotalDiffEntries(), actual.getTotalDiffEntries());
+    assertEquals(expected.getSubStatus(), actual.getSubStatus());
+    assertEquals(expected.getKeysProcessedPct(), actual.getKeysProcessedPct(), 0.001);
+    assertEquals(expected.getReason(), actual.getReason());
+  }
 }