kubernetes-overlord-extension: Fix tasks not being shutdown (#16711)

Author: ac9817

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java

@@ -443,11 +443,17 @@ public Collection<TaskRunnerWorkItem> getPendingTasks()
   @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();
     }
   }
 

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java

@@ -653,6 +653,24 @@ public TaskLocation getLocation()
     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()
   {