From 4c4a954c1c7aa93785ac14eabe61bb721a38ea94 Mon Sep 17 00:00:00 2001
From: Jonathan Klabunde Tomer <jkt@signal.org>
Date: Wed, 19 Nov 2025 15:04:11 -0800
Subject: [PATCH] update shutdown gauge when delayed shutdown starts

Otherwise we will report that we are not shutting down while k8s correctly
notes that we are unhealthy and it will look like something is wrong.
---
 .../metrics/ApplicationShutdownMonitor.java   | 24 ++++++++++++++++---
 .../textsecuregcm/metrics/MetricsUtil.java    |  2 +-
 2 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/ApplicationShutdownMonitor.java b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/ApplicationShutdownMonitor.java
index ab4c36319..d3e5477cd 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/ApplicationShutdownMonitor.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/ApplicationShutdownMonitor.java
@@ -11,15 +11,19 @@ import static org.whispersystems.textsecuregcm.metrics.MetricsUtil.name;
 import io.micrometer.core.instrument.Gauge;
 import io.micrometer.core.instrument.MeterRegistry;
 import java.util.concurrent.atomic.AtomicBoolean;
-import org.eclipse.jetty.util.component.LifeCycle;
+import org.eclipse.jetty.util.component.AbstractLifeCycle;
+import org.eclipse.jetty.util.thread.ShutdownThread;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 /**
  * A managed monitor that reports whether the application is shutting down as a metric. That metric can then be used in
  * conjunction with other indicators to conditionally fire or suppress alerts.
  */
-public class ApplicationShutdownMonitor implements LifeCycle.Listener {
+public class ApplicationShutdownMonitor extends AbstractLifeCycle {
 
   private final AtomicBoolean shuttingDown = new AtomicBoolean(false);
+  private final Logger logger = LoggerFactory.getLogger(ApplicationShutdownMonitor.class);
 
   public ApplicationShutdownMonitor(final MeterRegistry meterRegistry) {
     // without a strong reference to the gauge’s value supplier, shutdown garbage collection
@@ -29,8 +33,22 @@ public class ApplicationShutdownMonitor implements LifeCycle.Listener {
         .register(meterRegistry);
   }
 
+  public void register() {
+    // Force this component to get shut down before Dropwizard's
+    // DelayedShutdownHandler, which initiates the delayed-shutdown process
+    // without an additional chance for us to hook it
+    logger.info("registering shutdown monitor");
+    try {
+      start();                    // jetty won't stop an unstarted lifecycle
+      ShutdownThread.register(0, this);
+    } catch (Exception e) {
+      logger.error("failed to start application shutdown monitor", e);
+    }
+  }
+
   @Override
-  public void lifeCycleStopping(final LifeCycle event) {
+  public void doStop() {
+    logger.info("setting shutdown flag");
     shuttingDown.set(true);
   }
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
index 7b8d169cc..58521ed5c 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
@@ -83,7 +83,7 @@ public class MetricsUtil {
     environment.lifecycle().addServerLifecycleListener(
         server -> JettySslHandshakeMetrics.addToAllConnectors(server, Metrics.globalRegistry));
 
-    environment.lifecycle().addEventListener(new ApplicationShutdownMonitor(Metrics.globalRegistry));
+    new ApplicationShutdownMonitor(Metrics.globalRegistry).register();
     environment.lifecycle().addEventListener(
         new MicrometerRegistryManager(Metrics.globalRegistry, shutdownWaitDuration));