Addressing comments for empty taskId

Signed-off-by: Vacha Shah <vachshah@amazon.com>
opensearch-project · Sep 14, 2023 · b5b236f · b5b236f
1 parent f76c872
commit b5b236f
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 5 deletions.
diff --git a/libs/core/src/main/java/org/opensearch/core/tasks/TaskId.java b/libs/core/src/main/java/org/opensearch/core/tasks/TaskId.java
@@ -56,7 +56,7 @@ public final class TaskId implements Writeable, ProtobufWriteable {
 
     public static final TaskId EMPTY_TASK_ID = new TaskId();
 
-    private final TaskIdProto.TaskId taskIdProto;
+    private TaskIdProto.TaskId taskIdProto;
 
     private final String nodeId;
     private final long id;
@@ -76,7 +76,11 @@ public TaskId(String nodeId, long id) {
     private TaskId() {
         nodeId = "";
         id = -1;
-        taskIdProto = TaskIdProto.TaskId.newBuilder().setId(id).build();
+        try {
+            taskIdProto = TaskIdProto.TaskId.parseFrom(new byte[0]);
+        } catch (IOException ex) {
+            throw new IllegalArgumentException("malformed task id ", ex);
+        }
     }
 
     public TaskId(String taskId) {
@@ -95,7 +99,11 @@ public TaskId(String taskId) {
         } else {
             nodeId = "";
             id = -1L;
-            taskIdProto = TaskIdProto.TaskId.newBuilder().setId(id).build();
+            try {
+                taskIdProto = TaskIdProto.TaskId.parseFrom(new byte[0]);
+            } catch (IOException ex) {
+                throw new IllegalArgumentException("malformed task id ", ex);
+            }
         }
     }
 

diff --git a/libs/core/src/test/java/org/opensearch/core/tasks/TaskIdTests.java b/libs/core/src/test/java/org/opensearch/core/tasks/TaskIdTests.java
@@ -98,8 +98,12 @@ private TaskId roundTripForProtobuf(TaskId taskId) throws IOException {
             taskId.writeTo(out);
             byte[] bytes = out.toByteArray();
             TaskId taskId2 = TaskId.readFromBytes(bytes);
-            assertEquals(taskId.getId(), taskId2.getTaskIdProto().getId());
-            assertEquals(taskId.getNodeId(), taskId2.getTaskIdProto().getNodeId());
+            if (taskId.equals(TaskId.EMPTY_TASK_ID)) {
+                assertEquals(0, taskId2.getTaskIdProto().getSerializedSize());
+            } else {
+                assertEquals(taskId.getId(), taskId2.getTaskIdProto().getId());
+                assertEquals(taskId.getNodeId(), taskId2.getTaskIdProto().getNodeId());
+            }
             return taskId2;
         }
     }