From bc4d842d0ef28e578b4ac86e52f87c6c0becd5ac Mon Sep 17 00:00:00 2001
From: Chenyang Ji <cyji@amazon.com>
Date: Mon, 12 Feb 2024 14:03:59 -0800
Subject: [PATCH] Add support for deep copying SearchRequest

Signed-off-by: Chenyang Ji <cyji@amazon.com>
---
 .../action/search/SearchRequest.java          | 22 +++++++++++++++++++
 .../action/search/SearchRequestTests.java     |  8 +++++++
 2 files changed, 30 insertions(+)

diff --git a/server/src/main/java/org/opensearch/action/search/SearchRequest.java b/server/src/main/java/org/opensearch/action/search/SearchRequest.java
index f738c182c06da..32a13585da037 100644
--- a/server/src/main/java/org/opensearch/action/search/SearchRequest.java
+++ b/server/src/main/java/org/opensearch/action/search/SearchRequest.java
@@ -39,6 +39,7 @@
 import org.opensearch.action.support.IndicesOptions;
 import org.opensearch.common.Nullable;
 import org.opensearch.common.annotation.PublicApi;
+import org.opensearch.common.io.stream.BytesStreamOutput;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.common.Strings;
 import org.opensearch.core.common.io.stream.StreamInput;
@@ -160,6 +161,27 @@ public SearchRequest(String[] indices, SearchSourceBuilder source) {
         this.source = source;
     }
 
+    /**
+     * Deep clone a SearchRequest
+     *
+     * @return a copy of the current SearchRequest
+     */
+    @Override
+    public SearchRequest clone() {
+        try (BytesStreamOutput out = new BytesStreamOutput()) {
+            try {
+                this.writeTo(out);
+            } catch (IOException e) {
+                throw new IllegalArgumentException(e);
+            }
+            try (StreamInput in = out.bytes().streamInput()) {
+                return new SearchRequest(in);
+            } catch (IOException e) {
+                throw new IllegalArgumentException(e);
+            }
+        }
+    }
+
     /**
      * Creates a new sub-search request starting from the original search request that is provided.
      * For internal use only, allows to fork a search request into multiple search requests that will be executed independently.
diff --git a/server/src/test/java/org/opensearch/action/search/SearchRequestTests.java b/server/src/test/java/org/opensearch/action/search/SearchRequestTests.java
index f025e3a63b9bf..a5143d43c5546 100644
--- a/server/src/test/java/org/opensearch/action/search/SearchRequestTests.java
+++ b/server/src/test/java/org/opensearch/action/search/SearchRequestTests.java
@@ -269,6 +269,14 @@ public void testDescriptionIncludesScroll() {
         );
     }
 
+    public void testClone() throws IOException {
+        SearchRequest searchRequest = createSearchRequest();
+        SearchRequest clonedRequest = searchRequest.clone();
+        assertEquals(searchRequest.hashCode(), clonedRequest.hashCode());
+        assertNotSame(searchRequest, clonedRequest);
+    }
+
+
     private String toDescription(SearchRequest request) {
         return request.createTask(0, "test", SearchAction.NAME, TaskId.EMPTY_TASK_ID, emptyMap()).getDescription();
     }