From 488f1d83634e002c52ad9b7ff8c4cbb00efb4937 Mon Sep 17 00:00:00 2001
From: Nicholas Lippis <nick.lippis@imply.io>
Date: Wed, 29 Mar 2023 03:57:06 -0400
Subject: [PATCH] Do not print error message if pod not found when getting task
 location (#13971)

Do not print error message if pod not found when getting task location
---
 .../druid/k8s/overlord/KubernetesTaskRunner.java  |  5 +++++
 .../k8s/overlord/KubernetesTaskRunnerTest.java    | 15 +++++++++++++++
 2 files changed, 20 insertions(+)

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 ae2dbbb885a..7f9e92a9fa3 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
@@ -50,6 +50,7 @@ import org.apache.druid.k8s.overlord.common.DruidK8sConstants;
 import org.apache.druid.k8s.overlord.common.JobResponse;
 import org.apache.druid.k8s.overlord.common.K8sTaskId;
 import org.apache.druid.k8s.overlord.common.KubernetesPeonClient;
+import org.apache.druid.k8s.overlord.common.KubernetesResourceNotFoundException;
 import org.apache.druid.k8s.overlord.common.PeonPhase;
 import org.apache.druid.k8s.overlord.common.TaskAdapter;
 import org.apache.druid.tasklogs.TaskLogPusher;
@@ -469,6 +470,10 @@ public class KubernetesTaskRunner implements TaskLogStreamer, TaskRunner
             tlsEnabled
         );
       }
+      catch (KubernetesResourceNotFoundException e) {
+        log.debug(e, "Error getting task location for task %s", taskId);
+        return TaskLocation.unknown();
+      }
       catch (Exception e) {
         log.error(e, "Error getting task location for task %s", taskId);
         return TaskLocation.unknown();
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 6914f79d13f..7045abb9e68 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
@@ -52,6 +52,7 @@ import org.apache.druid.k8s.overlord.common.JobResponse;
 import org.apache.druid.k8s.overlord.common.K8sTaskAdapter;
 import org.apache.druid.k8s.overlord.common.K8sTaskId;
 import org.apache.druid.k8s.overlord.common.KubernetesPeonClient;
+import org.apache.druid.k8s.overlord.common.KubernetesResourceNotFoundException;
 import org.apache.druid.k8s.overlord.common.PeonPhase;
 import org.apache.druid.server.DruidNode;
 import org.apache.druid.server.log.StartupLoggingConfig;
@@ -427,6 +428,20 @@ public class KubernetesTaskRunnerTest
     assertEquals(TaskLocation.create("tweak", DruidK8sConstants.PORT, DruidK8sConstants.TLS_PORT, false), realLocation);
   }
 
+  @Test
+  public void testWorkItemGetLocation_withKubernetesResourceNotFoundException_returnsUnknownLocation()
+  {
+    KubernetesPeonClient client = mock(KubernetesPeonClient.class);
+
+    when(client.getMainJobPod(any())).thenThrow(KubernetesResourceNotFoundException.class);
+
+    Task task = mock(Task.class);
+    when(task.getId()).thenReturn("id");
+    KubernetesTaskRunner.K8sWorkItem k8sWorkItem = new KubernetesTaskRunner.K8sWorkItem(client, task, null);
+    TaskLocation location = k8sWorkItem.getLocation();
+    assertEquals(TaskLocation.unknown(), location);
+  }
+
   private Task makeTask()
   {
     return new TestableNoopTask(