Merge branch 'pg--cat-1551-update-components-semantic-tokens' into pg--cat-1659-update-semantic-tokens-shared-folder

Author: purnimagarg1

datahub-frontend/app/config/GracefulShutdownModule.java

@@ -0,0 +1,229 @@
+package config;
+
+import akka.Done;
+import akka.actor.CoordinatedShutdown;
+import com.google.inject.AbstractModule;
+import com.google.inject.Inject;
+import com.typesafe.config.Config;
+import java.io.IOException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.atomic.AtomicBoolean;
+import javax.inject.Provider;
+import javax.inject.Singleton;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Guice module managing graceful shutdown for the DataHub Frontend pod.
+ *
+ * <p>When enabled via FRONTEND_GRACEFUL_SHUTDOWN_ENABLED, this module coordinates with Akka's
+ * CoordinatedShutdown to gracefully handle Kubernetes pod termination:
+ *
+ * <ol>
+ *   <li>AWS Spot/K8s sends SIGTERM to the pod
+ *   <li>Akka CoordinatedShutdown triggers, calling registered phase tasks
+ *   <li>In before-service-unbind phase: FrontendShutdownHook flips isShuttingDown flag to true
+ *   <li>HealthCheckController reads this flag and returns 503, signaling readiness to LB
+ *   <li>Load balancer stops routing new requests and drains existing connections
+ *   <li>Remaining phases (service-requests-done: 65s, service-stop: 15s) allow graceful cleanup
+ *   <li>SIGKILL arrives after terminationGracePeriodSeconds (120s) if not already terminated
+ * </ol>
+ *
+ * <p>Thread-safe via AtomicBoolean to prevent race conditions between shutdown hook and concurrent
+ * request handling.
+ *
+ * <p>Disabled by default (FRONTEND_GRACEFUL_SHUTDOWN_ENABLED=false) for backward compatibility.
+ * When disabled, a no-op FrontendShutdownHook is instantiated, ensuring graceful shutdown can be
+ * toggled without code changes.
+ */
+@Singleton
+public class GracefulShutdownModule extends AbstractModule {
+
+  /**
+   * Flag indicating if the service is shutting down.
+   *
+   * <p>Set to true by FrontendShutdownHook during the before-service-unbind phase of Akka
+   * CoordinatedShutdown. Read by HealthCheckController to return 503 responses, signaling
+   * Kubernetes that the pod should stop receiving new requests.
+   *
+   * <p>Uses AtomicBoolean for thread-safe reads/writes from concurrent request handlers and the
+   * shutdown hook.
+   */
+  private final AtomicBoolean isShuttingDown = new AtomicBoolean(false);
+
+  @Override
+  protected void configure() {
+    bind(GracefulShutdownModule.class).toInstance(this);
+    bind(FrontendShutdownHook.class)
+        .toProvider(FrontendShutdownHookProvider.class)
+        .asEagerSingleton();
+  }
+
+  /**
+   * Checks if the service is currently shutting down.
+   *
+   * <p>This is called by HealthCheckController to determine if the readiness probe should return
+   * 503 (Service Unavailable). Non-blocking atomic read.
+   *
+   * @return true if shutdown has been initiated, false otherwise
+   */
+  public boolean isShuttingDown() {
+    return isShuttingDown.get();
+  }
+
+  /**
+   * Marks the service as shutting down (one-way operation for production use).
+   *
+   * <p>Called by FrontendShutdownHook during the before-service-unbind phase of Akka
+   * CoordinatedShutdown to signal that the service is shutting down and should not accept new
+   * requests.
+   *
+   * <p>This is a one-way latch that only transitions from false→true, matching the GMS handler
+   * pattern (GracefulShutdownHandler.onApplicationClosed). This prevents external code from
+   * reversing a shutdown state.
+   */
+  public void markShuttingDown() {
+    isShuttingDown.set(true);
+  }
+
+  /**
+   * Sets the shutdown flag to an arbitrary value.
+   *
+   * <p><b>@VisibleForTesting</b> — Used by tests to reset state between test runs. This method
+   * exists only for test cleanup, allowing tests to verify behavior in both shutdown and running
+   * states.
+   *
+   * <p>Production code should use {@link #markShuttingDown()} instead, which enforces one-way
+   * semantics and prevents accidental state reversals.
+   *
+   * @param value true to mark service as shutting down, false to reset state (tests only)
+   */
+  public void setShuttingDown(boolean value) {
+    isShuttingDown.set(value);
+  }
+
+  /**
+   * Provider factory for FrontendShutdownHook that conditionally enables graceful shutdown.
+   *
+   * <p>This indirection allows graceful shutdown to be toggled via configuration (
+   * FRONTEND_GRACEFUL_SHUTDOWN_ENABLED) without code changes or complex conditional bean
+   * definitions. Returns either an active hook (when enabled) or a no-op hook (when disabled).
+   *
+   * <p><b>Design Pattern (No-Op Construction)</b>: When disabled, a FrontendShutdownHook is still
+   * instantiated but with null arguments. The constructor's null-check (line 171) detects this and
+   * skips task registration, creating a safe no-op. This avoids the complexity of conditional bean
+   * definitions while maintaining clear intent through documentation.
+   */
+  @Singleton
+  public static class FrontendShutdownHookProvider implements Provider<FrontendShutdownHook> {
+
+    private final Config config;
+    private final CoordinatedShutdown coordinatedShutdown;
+    private final CloseableHttpClient httpClient;
+    private final GracefulShutdownModule module;
+
+    @Inject
+    public FrontendShutdownHookProvider(
+        Config config,
+        CoordinatedShutdown coordinatedShutdown,
+        CloseableHttpClient httpClient,
+        GracefulShutdownModule module) {
+      this.config = config;
+      this.coordinatedShutdown = coordinatedShutdown;
+      this.httpClient = httpClient;
+      this.module = module;
+    }
+
+    @Override
+    public FrontendShutdownHook get() {
+      if (config.getBoolean("frontend.graceful_shutdown_enabled")) {
+        return new FrontendShutdownHook(coordinatedShutdown, httpClient, module);
+      }
+      // Return a no-op hook if graceful shutdown is disabled.
+      // Passing null for coordinatedShutdown signals the constructor (see line 171)
+      // to skip all task registration, making this hook a no-op.
+      // This is the no-op construction pattern — allows feature toggle without code changes.
+      return new FrontendShutdownHook(null, null, module);
+    }
+  }
+
+  /**
+   * Registers Akka CoordinatedShutdown phase tasks to gracefully shut down the frontend.
+   *
+   * <p>Akka CoordinatedShutdown provides a multi-phase shutdown mechanism that coordinates resource
+   * cleanup with the termination signal. The phases execute in order:
+   *
+   * <ol>
+   *   <li><b>before-service-unbind (10s timeout)</b>: Signal intent to clients (e.g., via status
+   *       pages, server header). Sets isShuttingDown=true so HealthCheckController returns 503.
+   *   <li><b>service-requests-done (65s timeout)</b>: Wait for in-flight HTTP requests to complete.
+   *       Play Framework drains the HTTP server during this phase.
+   *   <li><b>service-stop (15s timeout)</b>: Cleanup remaining resources (WebSocket clients, HTTP
+   *       caches, etc.).
+   * </ol>
+   *
+   * <p><b>Timing Budget Analysis</b>: Akka phases total 10 + 65 + 15 = 90 seconds. Kubernetes
+   * terminationGracePeriodSeconds is set to 120s. The preStop hook (70s sleep) and Akka shutdown
+   * overlap after SIGTERM, so the effective maximum is max(70, 90) = 90 seconds, safely within the
+   * 120s K8s budget. AWS Spot instances provide 120 seconds before SIGKILL.
+   *
+   * <p><b>No-op behavior</b>: If coordinatedShutdown is null (feature disabled), no tasks are
+   * registered, and the hook becomes a no-op. This allows graceful shutdown to be toggled via
+   * configuration alone.
+   */
+  @Singleton
+  public static class FrontendShutdownHook {
+
+    private static final Logger log = LoggerFactory.getLogger(FrontendShutdownHook.class);
+    private final CoordinatedShutdown coordinatedShutdown;
+    private final CloseableHttpClient httpClient;
+    private final GracefulShutdownModule module;
+
+    public FrontendShutdownHook(
+        CoordinatedShutdown coordinatedShutdown,
+        CloseableHttpClient httpClient,
+        GracefulShutdownModule module) {
+      this.coordinatedShutdown = coordinatedShutdown;
+      this.httpClient = httpClient;
+      this.module = module;
+
+      if (coordinatedShutdown != null) {
+        // Phase 1: before-service-unbind (10s default timeout)
+        // Signal clients that shutdown is starting and flip the readiness flag.
+        // HealthCheckController reads isShuttingDown and returns 503 (Service Unavailable),
+        // which Kubernetes interprets as readiness failure. The load balancer stops routing
+        // new traffic while existing connections are allowed to drain.
+        coordinatedShutdown.addTask(
+            CoordinatedShutdown.PhaseBeforeServiceUnbind(),
+            "mark-unhealthy",
+            () ->
+                CompletableFuture.runAsync(
+                        () -> {
+                          log.info("Frontend shutdown initiated - stopping new connections soon");
+                          module.markShuttingDown();
+                        })
+                    .thenApply(v -> Done.done()));
+
+        // Phase 3: service-stop (15s default timeout)
+        // After in-flight requests drain, close long-lived connections (HTTP client for GMS calls).
+        coordinatedShutdown.addTask(
+            CoordinatedShutdown.PhaseServiceStop(),
+            "close-http-clients",
+            () ->
+                CompletableFuture.runAsync(
+                        () -> {
+                          try {
+                            log.info("Frontend shutdown initiated - shutting down open resources");
+                            if (httpClient != null) {
+                              httpClient.close();
+                            }
+                          } catch (IOException e) {
+                            log.error("Error closing CloseableHttpClient during shutdown", e);
+                          }
+                        })
+                    .thenApply(v -> Done.done()));
+      }
+    }
+  }
+}

datahub-frontend/app/controllers/Application.java

@@ -12,6 +12,7 @@
 import com.linkedin.metadata.utils.BasePathUtils;
 import com.linkedin.util.Pair;
 import com.typesafe.config.Config;
+import config.GracefulShutdownModule;
 import java.io.InputStream;
 import java.net.URI;
 import java.net.http.HttpClient;
@@ -50,16 +51,22 @@ public class Application extends Controller {
 
   private final Config config;
   private final Environment environment;
+  private final GracefulShutdownModule shutdownModule;
 
   private final String basePath;
   private final String gaTrackingId;
   private final List<String> streamingPathPrefixes;
 
   @Inject
-  public Application(HttpClient httpClient, Environment environment, @Nonnull Config config) {
+  public Application(
+      HttpClient httpClient,
+      Environment environment,
+      @Nonnull Config config,
+      GracefulShutdownModule shutdownModule) {
     this.httpClient = httpClient;
     this.config = config;
     this.environment = environment;
+    this.shutdownModule = shutdownModule;
     this.basePath = config.getString("datahub.basePath");
     this.gaTrackingId =
         config.hasPath("analytics.google.tracking.id")
@@ -122,6 +129,9 @@ private Result serveAsset(@Nullable String path) {
 
   @Nonnull
   public Result healthcheck() {
+    if (shutdownModule.isShuttingDown()) {
+      return status(SERVICE_UNAVAILABLE, "Shutting down");
+    }
     return ok("GOOD");
   }
 

datahub-frontend/conf/application.conf

@@ -27,6 +27,7 @@ play.modules.disabled += "play.api.mvc.CookiesModule"
 play.modules.enabled += "auth.cookie.CustomCookiesModule"
 play.modules.enabled += "auth.AuthModule"
 play.modules.enabled += "modules.StartupModule"
+play.modules.enabled += "config.GracefulShutdownModule"
 
 # Debug configuration dumping, if set to true might log secrets too.
 # Do not enable Akka log-config-on-start (CVE-2023-45865: env vars could be logged).
@@ -321,3 +322,37 @@ play.server.https.keyStore {
   type = ${?PLAY_HTTPS_KEYSTORE_TYPE}
   password = ${?PLAY_HTTPS_KEYSTORE_PASSWORD}
 }
+
+# Graceful shutdown settings
+# Enable/disable graceful shutdown hook for marking service as unhealthy during shutdown
+# When enabled, the frontend will flip its readiness flag on SIGTERM, allowing Kubernetes
+# to drain existing connections while rejecting new requests.
+# Default: false (backward compatible, graceful shutdown must be explicitly enabled)
+frontend.graceful_shutdown_enabled = false
+frontend.graceful_shutdown_enabled = ${?FRONTEND_GRACEFUL_SHUTDOWN_ENABLED}
+
+akka.coordinated-shutdown.phases {
+  # Akka CoordinatedShutdown phases execute sequentially: before-service-unbind → service-requests-done → service-stop
+  # Total duration: 10s + 65s + 15s = 90 seconds
+  #
+  # Timing budget analysis:
+  # - Kubernetes preStop hook (70s sleep) and Akka shutdown (90s total) overlap after SIGTERM
+  # - Effective max shutdown time: max(70s, 90s) = 90s, safely within K8s terminationGracePeriodSeconds (120s)
+  # - AWS Spot instance termination window: 120s before SIGKILL
+  # - Buffer: 30s between Akka completion (90s) and K8s SIGKILL (120s)
+
+  # Phase 1: before-service-unbind (signal clients, flip readiness flag)
+  # time to signal clients before connections close
+  before-service-unbind.timeout = 10s
+  before-service-unbind.timeout = ${?FRONTEND_BEFORE_SERVICE_UNBIND_TIMEOUT}
+
+  # Phase 2: service-requests-done (drain in-flight HTTP requests)
+  # time for in-flight HTTP requests to drain — equivalent to Spring's timeout-per-shutdown-phase
+  service-requests-done.timeout = 65s
+  service-requests-done.timeout = ${?FRONTEND_SERVICE_REQUESTS_DONE_TIMEOUT}
+
+  # Phase 3: service-stop (cleanup resources)
+  # time for cleanup tasks (WebSocket clients, HTTP caches, etc.)
+  service-stop.timeout = 15s
+  service-stop.timeout = ${?FRONTEND_SERVICE_STOP_TIMEOUT}
+}

datahub-frontend/test/app/ApplicationTest.java

@@ -7,6 +7,7 @@
 import static org.mockito.Mockito.*;
 import static play.mvc.Http.Status.MOVED_PERMANENTLY;
 import static play.mvc.Http.Status.OK;
+import static play.mvc.Http.Status.SERVICE_UNAVAILABLE;
 import static play.test.Helpers.fakeRequest;
 import static play.test.Helpers.route;
 
@@ -24,6 +25,7 @@
 import com.nimbusds.jwt.JWT;
 import com.nimbusds.jwt.JWTClaimsSet;
 import com.nimbusds.jwt.JWTParser;
+import config.GracefulShutdownModule;
 import controllers.routes;
 import java.io.IOException;
 import java.net.HttpURLConnection;
@@ -1466,6 +1468,25 @@ public void testIndexWhenResourceNotFound() {
     assertEquals("no-cache", result.headers().get("Cache-Control"));
   }
 
+  @Test
+  public void testHealthCheckReturns503WhenShuttingDown() {
+    // Get the singleton GracefulShutdownModule instance from the app injector
+    GracefulShutdownModule shutdownModule = app.injector().instanceOf(GracefulShutdownModule.class);
+    // Set shutdown flag via instance method
+    shutdownModule.setShuttingDown(true);
+
+    try {
+      Http.RequestBuilder request = fakeRequest(Helpers.GET, "/health");
+      Result result = route(app, request);
+      assertEquals(SERVICE_UNAVAILABLE, result.status());
+      String content = Helpers.contentAsString(result);
+      assertEquals("Shutting down", content);
+    } finally {
+      // Reset the shutdown flag after test
+      shutdownModule.setShuttingDown(false);
+    }
+  }
+
   /**
    * Test module that provides a mock Application controller that simulates resource loading failure
    */
@@ -1486,7 +1507,10 @@ protected controllers.Application provideFailingApplicationController(
       Environment mockEnvironment = mock(Environment.class);
       when(mockEnvironment.resourceAsStream("public/index.html")).thenReturn(null);
 
-      return new controllers.Application(mockHttpClient, mockEnvironment, config);
+      // Mock GracefulShutdownModule for the test
+      GracefulShutdownModule mockShutdownModule = mock(GracefulShutdownModule.class);
+      return new controllers.Application(
+          mockHttpClient, mockEnvironment, config, mockShutdownModule);
     }
 
     @Provides