(druid) branch master updated: kubernetes-overlord-extension: Fix tasks not being shutdown (#16711)

[georgew5656]

This is an automated email from the ASF dual-hosted git repository.

georgew5656 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/druid.git


The following commit(s) were added to refs/heads/master by this push:
     new 6cf6838eb97 kubernetes-overlord-extension: Fix tasks not being 
shutdown (#16711)
6cf6838eb97 is described below

commit 6cf6838eb974634ce5ae664d5ca877d8c804036a
Author: Adithya Chakilam <35785271+adithyachaki...@users.noreply.github.com>
AuthorDate: Mon Jul 15 16:35:11 2024 -0500

    kubernetes-overlord-extension: Fix tasks not being shutdown (#16711)
---
 .../druid/k8s/overlord/KubernetesTaskRunner.java       | 14 ++++++++++----
 .../druid/k8s/overlord/KubernetesTaskRunnerTest.java   | 18 ++++++++++++++++++
 2 files changed, 28 insertions(+), 4 deletions(-)

diff --git 
a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java
 
b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java
index 080a0fdaa98..c324b49e13a 100644
--- 
a/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java
+++ 
b/extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java
@@ -443,11 +443,17 @@ public class KubernetesTaskRunner implements 
TaskLogStreamer, TaskRunner
   @Override
   public TaskLocation getTaskLocation(String taskId)
   {
-    final KubernetesWorkItem workItem = tasks.get(taskId);
-    if (workItem == null) {
+    try {
+      final KubernetesWorkItem workItem = tasks.get(taskId);
+      if (workItem == null) {
+        return TaskLocation.unknown();
+      } else {
+        return workItem.getLocation();
+      }
+    }
+    catch (Exception e) {
+      log.warn("Unable to find location for task [%s]", taskId);
       return TaskLocation.unknown();
-    } else {
-      return workItem.getLocation();
     }
   }
 
diff --git 
a/extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java
 
b/extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java
index 391db70afb2..67a5278c6a3 100644
--- 
a/extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java
+++ 
b/extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java
@@ -654,6 +654,24 @@ public class KubernetesTaskRunnerTest extends 
EasyMockSupport
     Assert.assertEquals(TaskLocation.create("host", 0, 1, false), 
taskLocation);
   }
 
+  @Test
+  public void test_getTaskLocation_throws()
+  {
+    KubernetesWorkItem workItem = new KubernetesWorkItem(task, null)
+    {
+      @Override
+      public TaskLocation getLocation()
+      {
+        throw new RuntimeException();
+      }
+    };
+
+    runner.tasks.put(task.getId(), workItem);
+
+    TaskLocation taskLocation = runner.getTaskLocation(task.getId());
+    Assert.assertEquals(TaskLocation.unknown(), taskLocation);
+  }
+
   @Test
   public void test_getTaskLocation_noTaskFound()
   {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscr...@druid.apache.org
For additional commands, e-mail: commits-h...@druid.apache.org