Merge pull request #6 from mercari/set-max-buffer

Set max buffer for kryo to serialize
mercari · Mar 25, 2022 · 42aef13 · 42aef13
2 parents 7cfa752 + 0f2a2c6
commit 42aef13
Show file tree

Hide file tree

Showing 6 changed files with 57 additions and 15 deletions.
diff --git a/README.md b/README.md
@@ -256,6 +256,8 @@ Struct{
 | cipher_algorithm                           | cipher algorithm used for data encryption (currently supports only one AEAD cipher: AES/GCM/NoPadding)                                                                                                                                         | string   | AES/GCM/NoPadding | AES/GCM/NoPadding                                                                                                                                                                                   | low        |
 | cipher_text_encoding                       | defines the encoding of the resulting ciphertext bytes (currently only supports 'base64')                                                                                                                                                      | string   | base64            | base64                                                                                                                                                                                              | low        |
 | path_delimiter                             | path delimiter used as field name separator when referring to nested fields in the input record                                                                                                                                                | string   | .                 | non-empty string                                                                                                                                                                                    | low        |
+| kryo_output_buffer_size                    | Initial buffer size for kryo to serialize.                                                                                                                                                                                                     | int      | 32                | int value                                                                                                                                                                                           | low        |
+| kryo_output_buffer_size_max                | Maximum buffer size for kryo to serialize. Default -1 corresponds to no upper limit (up to Integer.MAX_VALUE - 8 technically).                                                                                                                 | int      | -1                | int value                                                                                                                                                                                           | low        |
 
 ### Externalize configuration parameters
 

diff --git a/...cp/src/main/java/com/github/hpgrahsl/kafka/connect/transforms/kryptonite/CipherField.java b/...cp/src/main/java/com/github/hpgrahsl/kafka/connect/transforms/kryptonite/CipherField.java
@@ -95,6 +95,8 @@ public enum KeySource {
   public static final String CIPHER_MODE = "cipher_mode";
   public static final String KEY_SOURCE = "key_source";
   public static final String KMS_KEY_NAME = "kms_key_name";
+  public static final String KRYO_OUTPUT_BUFFER_SIZE = "kryo_output_buffer_size";
+  public static final String KRYO_OUTPUT_BUFFER_SIZE_MAX = "kryo_output_buffer_size_max";
 
   private static final String PATH_DELIMITER_DEFAULT = ".";
   private static final String FIELD_MODE_DEFAULT = "ELEMENT";
@@ -105,6 +107,8 @@ public enum KeySource {
   private static final String KEY_SOURCE_DEFAULT = "CONFIG";
   private static final String CIPHER_DATA_KEYS_DEFAULT = "[]";
   private static final String KMS_KEY_NAME_DEFAULT = null;
+  private static final int KRYO_OUTPUT_BUFFER_SIZE_DEFAULT = 32;
+  private static final int KRYO_OUTPUT_BUFFER_SIZE_MAX_DEFAULT = -1;
 
   public static final ConfigDef CONFIG_DEF =
       new ConfigDef()
@@ -207,7 +211,19 @@ public enum KeySource {
               KMS_KEY_NAME_DEFAULT,
               Importance.MEDIUM,
               "The GCP Cloud KMS key name for decrypting a data encryption key (DEK), "
-                  + "if the DEK is encrypted with a key encryption key (KEK)");
+                  + "if the DEK is encrypted with a key encryption key (KEK)")
+          .define(
+              KRYO_OUTPUT_BUFFER_SIZE,
+              Type.INT,
+              KRYO_OUTPUT_BUFFER_SIZE_DEFAULT,
+              Importance.LOW,
+              "Initial buffer size for kryo to serialize. Default is 32")
+          .define(
+              KRYO_OUTPUT_BUFFER_SIZE_MAX,
+              Type.INT,
+              KRYO_OUTPUT_BUFFER_SIZE_MAX_DEFAULT,
+              Importance.LOW,
+              "Maximum buffer size for kryo to serialize. Default -1 corresponds to no upper limit (up to Integer.MAX_VALUE - 8 technically).");
 
   private static final String PURPOSE = "(de)cipher record fields";
 
@@ -269,7 +285,9 @@ public void configure(Map<String, ?> props) {
               .stream()
               .collect(Collectors.toMap(FieldConfig::getName, Function.identity()));
       Kryptonite kryptonite = configureKryptonite(config);
-      SerdeProcessor serdeProcessor = new KryoSerdeProcessor();
+      SerdeProcessor serdeProcessor =
+          new KryoSerdeProcessor(
+              config.getInt(KRYO_OUTPUT_BUFFER_SIZE), config.getInt(KRYO_OUTPUT_BUFFER_SIZE_MAX));
       recordHandlerWithSchema =
           new SchemaawareRecordHandler(
               config,

diff --git a/.../src/main/java/com/github/hpgrahsl/kafka/connect/transforms/kryptonite/DataKeyConfig.java b/.../src/main/java/com/github/hpgrahsl/kafka/connect/transforms/kryptonite/DataKeyConfig.java
@@ -74,9 +74,6 @@ public int hashCode() {
 
   @Override
   public String toString() {
-    return "DataKeyConfig{"
-        + "identifier='"
-        + getIdentifier()
-        + "'}";
+    return "DataKeyConfig{" + "identifier='" + getIdentifier() + "'}";
   }
 }
diff --git a/.../src/main/java/com/github/hpgrahsl/kafka/connect/transforms/kryptonite/RecordHandler.java b/.../src/main/java/com/github/hpgrahsl/kafka/connect/transforms/kryptonite/RecordHandler.java
@@ -82,12 +82,7 @@ public Object processField(Object object, String matchedPath) {
       }
     } catch (Exception e) {
       throw new DataException(
-          "error: "
-              + cipherMode
-              + " of field path '"
-              + matchedPath
-              + "' failed unexpectedly",
-          e);
+          "error: " + cipherMode + " of field path '" + matchedPath + "' failed unexpectedly", e);
     }
   }
 

diff --git a/...va/com/github/hpgrahsl/kafka/connect/transforms/kryptonite/serdes/KryoSerdeProcessor.java b/...va/com/github/hpgrahsl/kafka/connect/transforms/kryptonite/serdes/KryoSerdeProcessor.java
@@ -20,7 +20,6 @@
 import com.esotericsoftware.kryo.Serializer;
 import com.esotericsoftware.kryo.io.Input;
 import com.esotericsoftware.kryo.io.Output;
-import java.io.ByteArrayOutputStream;
 import java.util.ArrayList;
 import java.util.List;
 import org.apache.kafka.connect.data.ConnectSchema;
@@ -34,15 +33,20 @@
 public class KryoSerdeProcessor implements SerdeProcessor {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(KryoSerdeProcessor.class);
+  private final int outputBufSize;
+  private final int outputBufMaxSize;
 
-  public KryoSerdeProcessor() {}
+  public KryoSerdeProcessor(int outputBufSize, int outputBufMaxSize) {
+    this.outputBufSize = outputBufSize;
+    this.outputBufMaxSize = outputBufMaxSize;
+  }
 
   public byte[] objectToBytes(Object object, Class<?> clazz) {
     return objectToBytes(object);
   }
 
   public byte[] objectToBytes(Object object) {
-    Output output = new Output(new ByteArrayOutputStream());
+    Output output = new Output(outputBufSize, outputBufMaxSize);
     KryoInstance.get().writeClassAndObject(output, object);
     return output.toBytes();
   }

diff --git a/...om/github/hpgrahsl/kafka/connect/transforms/kryptonite/serdes/KryoSerdeProcessorTest.java b/...om/github/hpgrahsl/kafka/connect/transforms/kryptonite/serdes/KryoSerdeProcessorTest.java
@@ -0,0 +1,26 @@
+package com.github.hpgrahsl.kafka.connect.transforms.kryptonite.serdes;
+
+import static org.junit.jupiter.api.Assertions.assertAll;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+import org.junit.jupiter.api.DisplayName;
+import org.junit.jupiter.api.Test;
+
+class KryoSerdeProcessorTest {
+
+  @Test
+  @DisplayName("serialize and deserialize with variable sized objects")
+  void serializeAndDeserializeWithVariableSized() {
+    KryoSerdeProcessor processor = new KryoSerdeProcessor(32, -1);
+
+    String nullString = null;
+    String small = "test";
+    String large = new String(new char[100000]).replace("\0", "a");
+
+    assertAll(
+        () ->
+            assertEquals(nullString, processor.bytesToObject(processor.objectToBytes(nullString))),
+        () -> assertEquals(small, processor.bytesToObject(processor.objectToBytes(small))),
+        () -> assertEquals(large, processor.bytesToObject(processor.objectToBytes(large))));
+  }
+}