Updated wait strategy and added a scheduler for polling to check function status

Author: priyanshu-ctds

tests/src/test/java/com/datastax/oss/pulsar/functions/transforms/tests/AbstractDockerTest.java

@@ -25,12 +25,13 @@
 import com.google.gson.Gson;
 import com.google.gson.reflect.TypeToken;
 import java.math.BigDecimal;
+import java.time.Duration;
 import java.time.LocalDate;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
-import java.util.concurrent.TimeUnit;
+import java.util.concurrent.*;
 import lombok.Value;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.avro.Conversions;
@@ -595,42 +596,8 @@ private void deployFunction(
                 new Gson().fromJson(userConfig, new TypeToken<Map<String, Object>>() {}.getType()))
             .build();
 
-    for (int i = 0; i < 600; i++) {
-      try {
-        admin.functions().createFunction(functionConfig, null);
-        break;
-      } catch (PulsarAdminException e) {
-        if (e.getMessage() != null && e.getMessage().contains("Leader not yet ready")) {
-          log.info("Function worker leader not ready yet, retrying... ({}/600)", i + 1);
-          Thread.sleep(500);
-        } else {
-          throw e;
-        }
-      }
-    }
-
-    FunctionStatus functionStatus = null;
-    for (int i = 0; i < 600; i++) {
-      functionStatus = admin.functions().getFunctionStatus("public", "default", functionName);
-      if (functionStatus.getNumRunning() == 1) {
-        break;
-      }
-      log.info("Function status: {}", functionStatus);
-      functionStatus
-          .getInstances()
-          .forEach(
-              f -> {
-                log.info("Function instance status: {}", f);
-                if (!StringUtils.isEmpty(f.getStatus().getError())) {
-                  log.error("Function errored out " + f);
-                }
-              });
-      Thread.sleep(500);
-    }
-
-    if (functionStatus.getNumRunning() != 1) {
-      fail("Function didn't start in time");
-    }
+    admin.functions().createFunction(functionConfig, null);
+    waitForFunctionRunning(functionName);
   }
 
   @Value
@@ -644,4 +611,53 @@ private static class Pojo2 {
     String c;
     String d;
   }
+
+  private void waitForFunctionRunning(String functionName) throws InterruptedException {
+    ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
+    CompletableFuture<Void> future = new CompletableFuture<>();
+    long deadline = System.currentTimeMillis() + Duration.ofMinutes(5).toMillis();
+    scheduler.scheduleWithFixedDelay(
+        () -> {
+          try {
+            FunctionStatus status =
+                admin.functions().getFunctionStatus("public", "default", functionName);
+
+            status.getInstances().stream()
+                .filter(f -> !StringUtils.isEmpty(f.getStatus().getError()))
+                .findFirst()
+                .ifPresent(
+                    f ->
+                        future.completeExceptionally(
+                            new RuntimeException("Function errored out: " + f)));
+
+            if (future.isCompletedExceptionally()) {
+              scheduler.shutdown();
+              return;
+            }
+            if (status.getNumRunning() == 1) {
+              future.complete(null);
+              scheduler.shutdown();
+              return;
+            }
+            if (System.currentTimeMillis() > deadline) {
+              future.completeExceptionally(
+                  new RuntimeException("Function didn't start in time. Status: " + status));
+              scheduler.shutdown();
+              return;
+            }
+            log.info("Function not running yet, status: {}", status);
+          } catch (PulsarAdminException e) {
+            future.completeExceptionally(e);
+            scheduler.shutdown();
+          }
+        }, 0, 2, TimeUnit.SECONDS);
+
+    try {
+      future.get();
+    } catch (ExecutionException e) {
+      fail(e.getCause().getMessage());
+    } finally {
+      scheduler.shutdownNow();
+    }
+  }
 }

tests/src/test/java/com/datastax/oss/pulsar/functions/transforms/tests/PulsarContainer.java

@@ -53,7 +53,9 @@ public void start() {
             .waitingFor(
                 (new WaitAllStrategy())
                     .withStrategy(Wait.defaultWaitStrategy())
-                    .withStrategy(Wait.forLogMessage(".*Created namespace public/default.*", 1)))
+                    .withStrategy(Wait.forLogMessage(".*Created namespace public/default.*", 1))
+                    .withStrategy(Wait.forLogMessage(".*Function worker service started.*", 1))
+                    .withStrategy(Wait.forLogMessage(".*FunctionMetaDataManager done becoming leader.*", 1)))
             .withLogConsumer(
                 (f) -> {
                   String text = f.getUtf8String().trim();