Fixes #25666: Add retry-with-backoff to getServiceStatus for transient failure recovery

Author: RajdeepKushwaha5

openmetadata-service/src/main/java/org/openmetadata/service/clients/pipeline/airflow/AirflowRESTClient.java

@@ -458,26 +458,36 @@ public PipelineServiceClientResponse getServiceStatusInternal() {
       if (response.statusCode() == 200) {
         JSONObject responseJSON = new JSONObject(response.body());
         String ingestionVersion = responseJSON.getString("version");
-        return validServerClientVersions(ingestionVersion, SERVER_VERSION)
-            ? buildHealthyStatus(ingestionVersion)
-            : buildUnhealthyStatus(
-                buildVersionMismatchErrorMessage(ingestionVersion, SERVER_VERSION));
+        if (validServerClientVersions(ingestionVersion, SERVER_VERSION)) {
+          return buildHealthyStatus(ingestionVersion);
+        }
+        // Version mismatch is a config error, not transient — use 422 to skip retries
+        return new PipelineServiceClientResponse()
+            .withCode(422)
+            .withReason(buildVersionMismatchErrorMessage(ingestionVersion, SERVER_VERSION))
+            .withPlatform(this.getPlatform());
       }
 
-      // Auth error when accessing the APIs
+      // Auth error when accessing the APIs — not retryable (use 401/403 code)
       if (response.statusCode() == 401 || response.statusCode() == 403) {
-        return buildUnhealthyStatus(
-            String.format(
-                "Authentication failed for user [%s] trying to access the Airflow APIs at [%s]",
-                this.username, serviceURL.toString()));
+        return new PipelineServiceClientResponse()
+            .withCode(response.statusCode())
+            .withReason(
+                String.format(
+                    "Authentication failed for user [%s] trying to access the Airflow APIs at [%s]",
+                    this.username, serviceURL.toString()))
+            .withPlatform(this.getPlatform());
       }
 
-      // APIs URL not found
+      // APIs URL not found — not retryable (use 404 code)
       if (response.statusCode() == 404) {
-        return buildUnhealthyStatus(
-            String.format(
-                "Airflow APIs not found at [%s]. Please validate if the OpenMetadata Airflow plugin is installed correctly. %s",
-                serviceURL.toString(), DOCS_LINK));
+        return new PipelineServiceClientResponse()
+            .withCode(response.statusCode())
+            .withReason(
+                String.format(
+                    "Airflow APIs not found at [%s]. Please validate if the OpenMetadata Airflow plugin is installed correctly. %s",
+                    serviceURL.toString(), DOCS_LINK))
+            .withPlatform(this.getPlatform());
       }
 
       return buildUnhealthyStatus(

openmetadata-service/src/test/java/org/openmetadata/service/clients/pipeline/airflow/AirflowRESTClientTest.java

@@ -224,14 +224,14 @@ void getServiceStatusReportsAuthFailuresAndMissingPlugins() throws Exception {
       server.enqueue("GET", healthPath, 401, "{\"detail\":\"denied\"}");
       AirflowRESTClient authFailureClient = newClient(server, basePath);
       PipelineServiceClientResponse authFailure = authFailureClient.getServiceStatusInternal();
-      assertEquals(500, authFailure.getCode());
+      assertEquals(401, authFailure.getCode());
       assertTrue(authFailure.getReason().contains("Authentication failed"));
 
       server.enqueue("GET", healthPath, 200, "{\"version\":\"" + version + "\"}");
       server.enqueue("GET", healthPath, 404, "{\"detail\":\"missing\"}");
       AirflowRESTClient missingPluginClient = newClient(server, basePath);
       PipelineServiceClientResponse missingPlugin = missingPluginClient.getServiceStatusInternal();
-      assertEquals(500, missingPlugin.getCode());
+      assertEquals(404, missingPlugin.getCode());
       assertTrue(missingPlugin.getReason().contains("Airflow APIs not found"));
     }
   }
@@ -249,7 +249,7 @@ void getServiceStatusReportsVersionMismatch() throws Exception {
       AirflowRESTClient client = newClient(server, basePath);
       PipelineServiceClientResponse status = client.getServiceStatusInternal();
 
-      assertEquals(500, status.getCode());
+      assertEquals(422, status.getCode());
       assertTrue(
           status.getReason().contains("upgrade your server")
               || status.getReason().contains("upgrade your ingestion client"));
@@ -500,6 +500,54 @@ void runPipelineThrowsDeploymentExceptionWhenAirflowReturnsAnError() throws Exce
     }
   }
 
+  @Test
+  void getServiceStatusRecoversFromTransientFailure() throws Exception {
+    try (AirflowTestServer server = new AirflowTestServer()) {
+      String basePath = "/airflow";
+      String healthPath = basePath + "/pluginsv2/api/v2/openmetadata/health-auth";
+      String version = PipelineServiceClient.getServerVersion();
+
+      // First call: detection probe succeeds
+      server.enqueue("GET", healthPath, 200, "{\"version\":\"" + version + "\"}");
+      // Second call (first getServiceStatus attempt): transient 500
+      server.enqueue("GET", healthPath, 500, "{\"error\":\"selector manager closed\"}");
+      // Third call (retry): healthy response
+      server.enqueue("GET", healthPath, 200, "{\"version\":\"" + version + "\"}");
+
+      AirflowRESTClient client = newClient(server, basePath);
+
+      PipelineServiceClientResponse status = client.getServiceStatus();
+
+      assertEquals(200, status.getCode());
+      assertEquals(version, status.getVersion());
+      List<RequestRecord> healthRequests = server.requests("GET", healthPath);
+      assertEquals(3, healthRequests.size());
+    }
+  }
+
+  @Test
+  void getServiceStatusDoesNotRetryOnAuthError() throws Exception {
+    try (AirflowTestServer server = new AirflowTestServer()) {
+      String basePath = "/airflow";
+      String healthPath = basePath + "/pluginsv2/api/v2/openmetadata/health-auth";
+      String version = PipelineServiceClient.getServerVersion();
+
+      // Detection probe succeeds
+      server.enqueue("GET", healthPath, 200, "{\"version\":\"" + version + "\"}");
+      // Auth failure — should NOT be retried
+      server.enqueue("GET", healthPath, 401, "{\"error\":\"unauthorized\"}");
+
+      AirflowRESTClient client = newClient(server, basePath);
+
+      PipelineServiceClientResponse status = client.getServiceStatus();
+
+      assertEquals(401, status.getCode());
+      // Only 2 requests: detection probe + single 401 (no retry)
+      List<RequestRecord> healthRequests = server.requests("GET", healthPath);
+      assertEquals(2, healthRequests.size());
+    }
+  }
+
   private static AirflowRESTClient newClient(AirflowTestServer server, String basePath)
       throws KeyStoreException {
     Parameters parameters = new Parameters();
@@ -512,7 +560,12 @@ private static AirflowRESTClient newClient(AirflowTestServer server, String base
     config.setApiEndpoint(server.url(basePath));
     config.setMetadataApiEndpoint("http://localhost:8585/api");
     config.setParameters(parameters);
-    return new AirflowRESTClient(config);
+    return new AirflowRESTClient(config) {
+      @Override
+      protected long getRetryBackoffMillis() {
+        return 0L;
+      }
+    };
   }
 
   private static IngestionPipeline ingestionPipeline(String name, boolean enabled) {
@@ -527,9 +580,18 @@ private static Map<String, List<String>> cookieHeaders(String... values) {
   }
 
   private static void invokePrivate(Object target, String methodName) throws Exception {
-    Method method = target.getClass().getDeclaredMethod(methodName);
-    method.setAccessible(true);
-    method.invoke(target);
+    Class<?> clazz = target.getClass();
+    while (clazz != null) {
+      try {
+        Method method = clazz.getDeclaredMethod(methodName);
+        method.setAccessible(true);
+        method.invoke(target);
+        return;
+      } catch (NoSuchMethodException e) {
+        clazz = clazz.getSuperclass();
+      }
+    }
+    throw new NoSuchMethodException(target.getClass().getName() + "." + methodName + "()");
   }
 
   private record RequestRecord(

openmetadata-spec/src/main/java/org/openmetadata/service/clients/pipeline/PipelineServiceClient.java

@@ -27,7 +27,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Properties;
-import java.util.function.Supplier;
 import java.util.regex.Pattern;
 import lombok.Getter;
 import lombok.Setter;
@@ -66,7 +65,7 @@ public abstract class PipelineServiceClient implements PipelineServiceClientInte
   protected static final String CONTENT_HEADER = "Content-Type";
   protected static final String CONTENT_TYPE = "application/json";
   private static final Integer MAX_ATTEMPTS = 3;
-  private static final Integer BACKOFF_TIME_SECONDS = 5;
+  private static final long DEFAULT_BACKOFF_MILLIS = 5_000L;
   private static final String DISABLED_STATUS = "disabled";
 
   protected static final String SERVER_VERSION;
@@ -196,38 +195,45 @@ private Response getHostIpInternal() {
 
   /**
    * Check the pipeline service status with an exception backoff to make sure we don't raise any
-   * false positives.
+   * false positives. Delegates to {@link #getServiceStatus()}, which already retries transient
+   * failures.
    */
   public String getServiceStatusBackoff() {
+    PipelineServiceClientResponse status = getServiceStatus();
+    return status.getCode() != 200 ? UNHEALTHY_STATUS : HEALTHY_STATUS;
+  }
+
+  /**
+   * Check the status of pipeline service to ensure it is healthy. Retries up to {@link
+   * #MAX_ATTEMPTS} times with backoff so that a single transient failure (e.g. Airflow restart,
+   * network blip) does not permanently mark the agent as unavailable. Responses with status codes
+   * ≥ 500 are retried; 4xx responses (auth, config, version mismatch) are returned immediately.
+   * Thrown exceptions are also retried by the underlying Resilience4j configuration.
+   */
+  public PipelineServiceClientResponse getServiceStatus() {
+    if (!pipelineServiceClientEnabled) {
+      return buildHealthyStatus(DISABLED_STATUS).withPlatform(DISABLED_STATUS);
+    }
+
     RetryConfig retryConfig =
-        RetryConfig.<String>custom()
+        RetryConfig.<PipelineServiceClientResponse>custom()
             .maxAttempts(MAX_ATTEMPTS)
-            .waitDuration(Duration.ofMillis(BACKOFF_TIME_SECONDS * 1_000L))
-            .retryOnResult(response -> !HEALTHY_STATUS.equals(response))
+            .waitDuration(Duration.ofMillis(getRetryBackoffMillis()))
+            .retryOnResult(response -> response != null && response.getCode() >= 500)
             .failAfterMaxAttempts(false)
             .build();
 
     Retry retry = Retry.of("getServiceStatus", retryConfig);
 
-    Supplier<String> responseSupplier =
-        () -> {
-          try {
-            PipelineServiceClientResponse status = getServiceStatus();
-            return status.getCode() != 200 ? UNHEALTHY_STATUS : HEALTHY_STATUS;
-          } catch (Exception e) {
-            throw new RuntimeException(e);
-          }
-        };
-
-    return retry.executeSupplier(responseSupplier);
+    PipelineServiceClientResponse response = retry.executeSupplier(this::getServiceStatusInternal);
+    return response != null
+        ? response
+        : buildUnhealthyStatus("Pipeline service returned no response");
   }
 
-  /* Check the status of pipeline service to ensure it is healthy */
-  public PipelineServiceClientResponse getServiceStatus() {
-    if (pipelineServiceClientEnabled) {
-      return getServiceStatusInternal();
-    }
-    return buildHealthyStatus(DISABLED_STATUS).withPlatform(DISABLED_STATUS);
+  /** Returns the wait duration between retry attempts in milliseconds. */
+  protected long getRetryBackoffMillis() {
+    return DEFAULT_BACKOFF_MILLIS;
   }
 
   public List<PipelineStatus> getQueuedPipelineStatus(IngestionPipeline ingestionPipeline) {