From d579c51ffd5997809db29c867a47a5772bd7ff39 Mon Sep 17 00:00:00 2001
From: Peter Alfonsi <petealft@amazon.com>
Date: Tue, 27 Feb 2024 13:50:41 -0800
Subject: [PATCH] Adds IRC key serializer

Signed-off-by: Peter Alfonsi <petealft@amazon.com>
---
 .../indices/IRCKeyWriteableSerializer.java    | 63 +++++++++++++++++++
 .../IRCKeyWriteableSerializerTests.java       | 54 ++++++++++++++++
 2 files changed, 117 insertions(+)
 create mode 100644 server/src/main/java/org/opensearch/indices/IRCKeyWriteableSerializer.java
 create mode 100644 server/src/test/java/org/opensearch/indices/IRCKeyWriteableSerializerTests.java

diff --git a/server/src/main/java/org/opensearch/indices/IRCKeyWriteableSerializer.java b/server/src/main/java/org/opensearch/indices/IRCKeyWriteableSerializer.java
new file mode 100644
index 0000000000000..56d16e84f3f42
--- /dev/null
+++ b/server/src/main/java/org/opensearch/indices/IRCKeyWriteableSerializer.java
@@ -0,0 +1,63 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.indices;
+
+import org.opensearch.OpenSearchException;
+import org.opensearch.common.cache.serializer.Serializer;
+import org.opensearch.common.io.stream.BytesStreamOutput;
+import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.core.common.io.stream.BytesStreamInput;
+
+import java.io.IOException;
+import java.util.Arrays;
+
+/**
+ * This class serializes the IndicesRequestCache.Key using its writeTo method.
+ */
+public class IRCKeyWriteableSerializer implements Serializer<IndicesRequestCache.Key, byte[]> {
+
+
+    public IRCKeyWriteableSerializer() {
+    }
+
+    @Override
+    public byte[] serialize(IndicesRequestCache.Key object) {
+        try {
+            BytesStreamOutput os = new BytesStreamOutput();
+            object.writeTo(os);
+            return BytesReference.toBytes(os.bytes());
+        } catch (IOException e) {
+            throw new OpenSearchException(e);
+        }
+    }
+
+    @Override
+    public IndicesRequestCache.Key deserialize(byte[] bytes) {
+        if (bytes == null) {
+            return null;
+        }
+        try {
+            BytesStreamInput is = new BytesStreamInput(bytes, 0, bytes.length);
+            return new IndicesRequestCache.Key(is);
+        } catch (IOException e) {
+            throw new OpenSearchException(e);
+        }
+    }
+
+    @Override
+    public boolean equals(IndicesRequestCache.Key object, byte[] bytes) {
+        // Deserialization is much slower than serialization for keys of order 1 KB,
+        // while time to serialize is fairly constant (per byte)
+        if (bytes.length < 5000) {
+            return Arrays.equals(serialize(object), bytes);
+        } else {
+            return object.equals(deserialize(bytes));
+        }
+    }
+}
diff --git a/server/src/test/java/org/opensearch/indices/IRCKeyWriteableSerializerTests.java b/server/src/test/java/org/opensearch/indices/IRCKeyWriteableSerializerTests.java
new file mode 100644
index 0000000000000..a4a97b5c8470f
--- /dev/null
+++ b/server/src/test/java/org/opensearch/indices/IRCKeyWriteableSerializerTests.java
@@ -0,0 +1,54 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+package org.opensearch.indices;
+
+import org.opensearch.common.Randomness;
+import org.opensearch.core.common.bytes.BytesArray;
+import org.opensearch.core.common.bytes.BytesReference;
+import org.opensearch.core.index.shard.ShardId;
+import org.opensearch.index.IndexService;
+import org.opensearch.index.shard.IndexShard;
+import org.opensearch.test.OpenSearchSingleNodeTestCase;
+
+import java.util.Random;
+import java.util.UUID;
+
+public class IRCKeyWriteableSerializerTests extends OpenSearchSingleNodeTestCase {
+
+    public void testSerializer() throws Exception {
+        IndexService indexService = createIndex("test");
+        IndexShard indexShard = indexService.getShardOrNull(0);
+        IRCKeyWriteableSerializer ser = new IRCKeyWriteableSerializer();
+
+        int NUM_KEYS = 1000;
+        int[] valueLengths = new int[] { 1000, 6000 }; // test both branches in equals()
+        Random rand = Randomness.get();
+        for (int valueLength : valueLengths) {
+            for (int i = 0; i < NUM_KEYS; i++) {
+                IndicesRequestCache.Key key = getRandomIRCKey(valueLength, rand, indexShard.shardId());
+                byte[] serialized = ser.serialize(key);
+                assertTrue(ser.equals(key, serialized));
+                IndicesRequestCache.Key deserialized = ser.deserialize(serialized);
+                assertTrue(key.equals(deserialized));
+            }
+        }
+    }
+
+    private IndicesRequestCache.Key getRandomIRCKey(
+        int valueLength,
+        Random random,
+        ShardId shard
+    ) {
+        byte[] value = new byte[valueLength];
+        for (int i = 0; i < valueLength; i++) {
+            value[i] = (byte) (random.nextInt(126 - 32) + 32);
+        }
+        BytesReference keyValue = new BytesArray(value);
+        return new IndicesRequestCache.Key(shard, keyValue, UUID.randomUUID().toString()); // same UUID source as used in real key
+    }
+}