From 7ccf552a1389858fd5857589dee8441ea51e2d61 Mon Sep 17 00:00:00 2001
From: brharrington <brharrington@gmail.com>
Date: Sun, 16 Jun 2024 19:28:00 -0500
Subject: [PATCH] scheduler will now wait for threads on shutdown (#1137)

Updated the scheduler helper to have the shutdown method block
until the threads stop. This makes it easier for callers to
ensure they are not competing with in-progress tasks while
performing other shutdown work.
---
 .../com/netflix/spectator/impl/Scheduler.java | 30 +++++++++++++++----
 .../netflix/spectator/impl/SchedulerTest.java |  2 +-
 2 files changed, 26 insertions(+), 6 deletions(-)

diff --git a/spectator-api/src/main/java/com/netflix/spectator/impl/Scheduler.java b/spectator-api/src/main/java/com/netflix/spectator/impl/Scheduler.java
index 6affbfba0..2986dd19f 100644
--- a/spectator-api/src/main/java/com/netflix/spectator/impl/Scheduler.java
+++ b/spectator-api/src/main/java/com/netflix/spectator/impl/Scheduler.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2023 Netflix, Inc.
+ * Copyright 2014-2024 Netflix, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -170,15 +170,28 @@ public ScheduledFuture<?> schedule(Options options, Runnable task) {
 
   /**
    * Shutdown and cleanup resources associated with the scheduler. All threads will be
-   * interrupted, but this method does not block for them to all finish execution.
+   * interrupted and this method will block until they are shutdown.
    */
   public void shutdown() {
     lock.lock();
     try {
       shutdown = true;
+
+      // Interrupt threads to shutdown
+      for (Thread thread : threads) {
+        if (thread != null && thread.isAlive()) {
+          thread.interrupt();
+        }
+      }
+
+      // Wait for all threads to complete
       for (int i = 0; i < threads.length; ++i) {
-        if (threads[i] != null && threads[i].isAlive()) {
-          threads[i].interrupt();
+        if (threads[i] != null) {
+          try {
+            threads[i].join();
+          } catch (Exception e) {
+            LOGGER.debug("exception while shutting down thread {}", threads[i].getName(), e);
+          }
           threads[i] = null;
         }
       }
@@ -188,6 +201,13 @@ public void shutdown() {
   }
 
   private void startThreads() {
+    // Normally when a thread exits, it will try to restart, if shutting down
+    // exit early before trying to get the lock.
+    if (shutdown) {
+      return;
+    }
+
+    // Start threads if setting up the scheduler or if a thread failed.
     lock.lock();
     try {
       if (!shutdown) {
@@ -478,7 +498,7 @@ private final class Worker implements Runnable {
       try {
         // Note: do not use Thread.interrupted() because it will clear the interrupt
         // status of the thread.
-        while (!Thread.currentThread().isInterrupted()) {
+        while (!shutdown && !Thread.currentThread().isInterrupted()) {
           try {
             DelayedTask task = queue.take();
             stats.incrementActiveTaskCount();
diff --git a/spectator-api/src/test/java/com/netflix/spectator/impl/SchedulerTest.java b/spectator-api/src/test/java/com/netflix/spectator/impl/SchedulerTest.java
index 4c9823810..2e4330ffe 100644
--- a/spectator-api/src/test/java/com/netflix/spectator/impl/SchedulerTest.java
+++ b/spectator-api/src/test/java/com/netflix/spectator/impl/SchedulerTest.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2014-2019 Netflix, Inc.
+ * Copyright 2014-2024 Netflix, Inc.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.