From 1926449ecd6e2c94a59127fa6cc3cc45ff1fb6c2 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Thu, 5 Oct 2023 13:57:12 -0400
Subject: [PATCH] Adding cancellation start time in client side TaskInfo
 (#10326) (#10384)

(cherry picked from commit 71b69482a82b647bc6b1bfbc011c9ea96511f1e4)

Signed-off-by: Sagar Upadhyaya <sagar.upadhyaya.121@gmail.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../org/opensearch/client/tasks/TaskInfo.java  | 18 ++++++++++++++++--
 .../client/tasks/CancelTasksResponseTests.java |  8 +++++++-
 2 files changed, 23 insertions(+), 3 deletions(-)

diff --git a/client/rest-high-level/src/main/java/org/opensearch/client/tasks/TaskInfo.java b/client/rest-high-level/src/main/java/org/opensearch/client/tasks/TaskInfo.java
index 51ac62830446f..75badc4e3dbf2 100644
--- a/client/rest-high-level/src/main/java/org/opensearch/client/tasks/TaskInfo.java
+++ b/client/rest-high-level/src/main/java/org/opensearch/client/tasks/TaskInfo.java
@@ -54,6 +54,7 @@ public class TaskInfo {
     private long runningTimeNanos;
     private boolean cancellable;
     private boolean cancelled;
+    private Long cancellationStartTime;
     private TaskId parentTaskId;
     private final Map<String, Object> status = new HashMap<>();
     private final Map<String, String> headers = new HashMap<>();
@@ -127,6 +128,14 @@ void setCancelled(boolean cancelled) {
         this.cancelled = cancelled;
     }
 
+    public Long getCancellationStartTime() {
+        return this.cancellationStartTime;
+    }
+
+    public void setCancellationStartTime(Long cancellationStartTime) {
+        this.cancellationStartTime = cancellationStartTime;
+    }
+
     public TaskId getParentTaskId() {
         return parentTaskId;
     }
@@ -180,6 +189,7 @@ private void noOpParse(Object s) {}
         parser.declareString(TaskInfo::setParentTaskId, new ParseField("parent_task_id"));
         parser.declareObject(TaskInfo::setHeaders, (p, c) -> p.mapStrings(), new ParseField("headers"));
         parser.declareObject(TaskInfo::setResourceStats, (p, c) -> p.map(), new ParseField("resource_stats"));
+        parser.declareLong(TaskInfo::setCancellationStartTime, new ParseField("cancellation_time_millis"));
         PARSER = (XContentParser p, Void v, String name) -> parser.parse(p, new TaskInfo(new TaskId(name)), null);
     }
 
@@ -199,7 +209,8 @@ && isCancelled() == taskInfo.isCancelled()
             && Objects.equals(getParentTaskId(), taskInfo.getParentTaskId())
             && Objects.equals(status, taskInfo.status)
             && Objects.equals(getHeaders(), taskInfo.getHeaders())
-            && Objects.equals(getResourceStats(), taskInfo.getResourceStats());
+            && Objects.equals(getResourceStats(), taskInfo.getResourceStats())
+            && Objects.equals(getCancellationStartTime(), taskInfo.cancellationStartTime);
     }
 
     @Override
@@ -216,7 +227,8 @@ public int hashCode() {
             getParentTaskId(),
             status,
             getHeaders(),
-            getResourceStats()
+            getResourceStats(),
+            getCancellationStartTime()
         );
     }
 
@@ -250,6 +262,8 @@ public String toString() {
             + headers
             + ", resource_stats="
             + resourceStats
+            + ", cancellationStartTime="
+            + cancellationStartTime
             + '}';
     }
 }
diff --git a/client/rest-high-level/src/test/java/org/opensearch/client/tasks/CancelTasksResponseTests.java b/client/rest-high-level/src/test/java/org/opensearch/client/tasks/CancelTasksResponseTests.java
index 835a93b5b09ce..faf5024d0c173 100644
--- a/client/rest-high-level/src/test/java/org/opensearch/client/tasks/CancelTasksResponseTests.java
+++ b/client/rest-high-level/src/test/java/org/opensearch/client/tasks/CancelTasksResponseTests.java
@@ -84,6 +84,10 @@ protected CancelTasksResponseTests.ByNodeCancelTasksResponse createServerTestIns
         for (int i = 0; i < 4; i++) {
             boolean cancellable = randomBoolean();
             boolean cancelled = cancellable == true ? randomBoolean() : false;
+            Long cancellationStartTime = null;
+            if (cancelled) {
+                cancellationStartTime = randomNonNegativeLong();
+            }
             tasks.add(
                 new org.opensearch.tasks.TaskInfo(
                     new TaskId(NODE_ID, (long) i),
@@ -97,7 +101,8 @@ protected CancelTasksResponseTests.ByNodeCancelTasksResponse createServerTestIns
                     cancelled,
                     new TaskId("node1", randomLong()),
                     Collections.singletonMap("x-header-of", "some-value"),
-                    null
+                    null,
+                    cancellationStartTime
                 )
             );
         }
@@ -135,6 +140,7 @@ protected void assertInstances(
             assertEquals(ti.isCancelled(), taskInfo.isCancelled());
             assertEquals(ti.getParentTaskId().getNodeId(), taskInfo.getParentTaskId().getNodeId());
             assertEquals(ti.getParentTaskId().getId(), taskInfo.getParentTaskId().getId());
+            assertEquals(ti.getCancellationStartTime(), taskInfo.getCancellationStartTime());
             FakeTaskStatus status = (FakeTaskStatus) ti.getStatus();
             assertEquals(status.code, taskInfo.getStatus().get("code"));
             assertEquals(status.status, taskInfo.getStatus().get("status"));