From dcf06f943910dc8c338ec45198bcd9090e439787 Mon Sep 17 00:00:00 2001
From: Naorem Khogendro Singh <nsingh@yugabyte.com>
Date: Wed, 8 Jan 2025 16:35:51 -0800
Subject: [PATCH] [PLAT-16436] Refactor platform scheduler to remove duplicate
 code

Summary: Minor refactoring to reuse code.

Test Plan: Schedulers are running fine with the change.

Reviewers: cwang, amalyshev

Reviewed By: amalyshev

Subscribers: yugaware

Differential Revision: https://phorge.dev.yugabyte.com/D41106
---
 .../yugabyte/yw/common/PlatformScheduler.java | 107 ++++++++----------
 1 file changed, 47 insertions(+), 60 deletions(-)

diff --git a/managed/src/main/java/com/yugabyte/yw/common/PlatformScheduler.java b/managed/src/main/java/com/yugabyte/yw/common/PlatformScheduler.java
index 3035fbdb0f05..d5bf01652dfb 100644
--- a/managed/src/main/java/com/yugabyte/yw/common/PlatformScheduler.java
+++ b/managed/src/main/java/com/yugabyte/yw/common/PlatformScheduler.java
@@ -5,6 +5,7 @@
 import com.yugabyte.yw.models.HighAvailabilityConfig;
 import java.time.Duration;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.function.Function;
 import javax.inject.Inject;
 import javax.inject.Singleton;
 import lombok.extern.slf4j.Slf4j;
@@ -31,45 +32,39 @@ public PlatformScheduler(
     this.shutdownHookHandler = shutdownHookHandler;
   }
 
-  public Cancellable schedule(
-      String name, Duration initialDelay, Duration interval, Runnable runnable) {
+  private Cancellable createShutdownAwareSchedule(
+      String name, Runnable runnable, Function<Runnable, Cancellable> scheduleFactory) {
     final AtomicBoolean isRunning = new AtomicBoolean();
     final Object lock = new Object();
-
-    Cancellable cancellable =
-        actorSystem
-            .scheduler()
-            .scheduleWithFixedDelay(
-                initialDelay,
-                interval,
-                () -> {
-                  boolean shouldRun = false;
-                  synchronized (lock) {
-                    // Synchronized block in shutdown and this should be serialized.
-                    shouldRun =
-                        !shutdownHookHandler.isShutdown()
-                            && !HighAvailabilityConfig.isFollower()
-                            && isRunning.compareAndSet(false, true);
-                  }
-                  if (shouldRun) {
-                    try {
-                      runnable.run();
-                    } finally {
-                      isRunning.set(false);
-                      if (shutdownHookHandler.isShutdown()) {
-                        synchronized (lock) {
-                          lock.notify();
-                        }
-                      }
-                    }
-                  } else {
-                    log.warn(
-                        "Previous run of scheduler {} is in progress, is being shut down, or YBA is"
-                            + " in follower mode.",
-                        name);
-                  }
-                },
-                executionContext);
+    Runnable wrappedRunnable =
+        () -> {
+          boolean shouldRun = false;
+          synchronized (lock) {
+            // Synchronized block in shutdown and this should be serialized.
+            shouldRun =
+                !shutdownHookHandler.isShutdown()
+                    && !HighAvailabilityConfig.isFollower()
+                    && isRunning.compareAndSet(false, true);
+          }
+          if (shouldRun) {
+            try {
+              runnable.run();
+            } finally {
+              isRunning.set(false);
+              if (shutdownHookHandler.isShutdown()) {
+                synchronized (lock) {
+                  lock.notify();
+                }
+              }
+            }
+          } else {
+            log.warn(
+                "Previous run of scheduler {} is in progress, is being shut down, or YBA is in"
+                    + " follower mode.",
+                name);
+          }
+        };
+    Cancellable cancellable = scheduleFactory.apply(wrappedRunnable);
     shutdownHookHandler.addShutdownHook(
         cancellable,
         can -> {
@@ -93,29 +88,21 @@ public Cancellable schedule(
     return cancellable;
   }
 
-  public Cancellable scheduleOnce(String name, Duration initialDelay, Runnable runnable) {
-    final AtomicBoolean isRunning = new AtomicBoolean();
+  public Cancellable schedule(
+      String name, Duration initialDelay, Duration interval, Runnable runnable) {
+    return createShutdownAwareSchedule(
+        name,
+        runnable,
+        r ->
+            actorSystem
+                .scheduler()
+                .scheduleWithFixedDelay(initialDelay, interval, r, executionContext));
+  }
 
-    return actorSystem
-        .scheduler()
-        .scheduleOnce(
-            initialDelay,
-            () -> {
-              boolean shouldRun = false;
-              synchronized (isRunning) {
-                shouldRun =
-                    isRunning.compareAndSet(false, true) && !HighAvailabilityConfig.isFollower();
-              }
-              if (shouldRun) {
-                try {
-                  runnable.run();
-                } finally {
-                  isRunning.set(false);
-                }
-              } else {
-                log.warn("Scheduler {} did not run because YBA is in follower mode.", name);
-              }
-            },
-            executionContext);
+  public Cancellable scheduleOnce(String name, Duration initialDelay, Runnable runnable) {
+    return createShutdownAwareSchedule(
+        name,
+        runnable,
+        r -> actorSystem.scheduler().scheduleOnce(initialDelay, r, executionContext));
   }
 }