From d90dff95b1cf3941ffe9c7df0164253509b5ebfc Mon Sep 17 00:00:00 2001
From: Chris Eager <chris@signal.org>
Date: Wed, 20 Dec 2023 12:21:19 -0600
Subject: [PATCH] Move MeterRegistry#close calls to after the lifecycle has
 fully stopped

---
 .../textsecuregcm/metrics/MetricsUtil.java    |  3 +-
 .../metrics/MicrometerRegistryManager.java    | 30 ++++++++++++-------
 2 files changed, 22 insertions(+), 11 deletions(-)

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 416f8339a..2a7ddbe5d 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MetricsUtil.java
@@ -60,7 +60,8 @@ public class MetricsUtil {
       Metrics.addRegistry(dogstatsdMeterRegistry);
     }
 
-    environment.lifecycle().manage(new MicrometerRegistryManager(Metrics.globalRegistry));
+    environment.lifecycle().addEventListener(
+        new MicrometerRegistryManager(Metrics.globalRegistry, config.getDatadogConfiguration().pollingFrequency()));
     environment.lifecycle().manage(new ApplicationShutdownMonitor(Metrics.globalRegistry));
   }
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MicrometerRegistryManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MicrometerRegistryManager.java
index 4fc55043c..0bf6798d0 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MicrometerRegistryManager.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/metrics/MicrometerRegistryManager.java
@@ -5,25 +5,35 @@
 
 package org.whispersystems.textsecuregcm.metrics;
 
-import io.dropwizard.lifecycle.Managed;
 import io.micrometer.core.instrument.MeterRegistry;
+import java.time.Duration;
+import org.eclipse.jetty.util.component.LifeCycle;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
-public class MicrometerRegistryManager implements Managed {
+public class MicrometerRegistryManager implements LifeCycle.Listener {
+
+  private static final Logger logger = LoggerFactory.getLogger(MicrometerRegistryManager.class);
+  private static final Duration BUFFER = Duration.ofSeconds(5);
 
   private final MeterRegistry meterRegistry;
+  private final Duration waitDuration;
 
-  public MicrometerRegistryManager(final MeterRegistry meterRegistry) {
+  public MicrometerRegistryManager(final MeterRegistry meterRegistry, final Duration pollingFrequency) {
     this.meterRegistry = meterRegistry;
+    this.waitDuration = pollingFrequency.plus(BUFFER);
   }
 
   @Override
-  public void start() throws Exception {
-
+  public void lifeCycleStopped(final LifeCycle event) {
+    try {
+      logger.info("Waiting for {} to ensure final metrics are polled and published", waitDuration);
+      Thread.sleep(waitDuration.toMillis());
+    } catch (final InterruptedException e) {
+      logger.warn("Waiting interrupted", e);
+    } finally {
+      meterRegistry.close();
+    }
   }
 
-  @Override
-  public void stop() throws Exception {
-    // closing the registry publishes one final set of metrics
-    meterRegistry.close();
-  }
 }