Make the shutdown hook configurable

[#151302620] https://www.pivotaltracker.com/story/show/151302620
cloudfoundry · Oct 10, 2017 · 4520026 · 4520026
1 parent b36a57d
commit 4520026
Show file tree

Hide file tree

Showing 5 changed files with 37 additions and 7 deletions.
diff --git a/server/src/main/java/org/cloudfoundry/identity/uaa/health/HealthzEndpoint.java b/server/src/main/java/org/cloudfoundry/identity/uaa/health/HealthzEndpoint.java
@@ -29,18 +29,21 @@
  */
 @Controller
 public class HealthzEndpoint {
-    private static final int SLEEP_UPON_SHUTDOWN = 5000;
     private static Log logger = LogFactory.getLog(HealthzEndpoint.class);
     private volatile boolean stopping = false;
     private final Thread shutdownhook;
+    private final long sleepTime;
 
-    public HealthzEndpoint() {
+    public HealthzEndpoint(long sleepTime) {
+        this.sleepTime = sleepTime;
         shutdownhook = new Thread(() -> {
             stopping = true;
             logger.warn("Shutdown hook received, future requests to this endpoint will return 503");
-            logger.debug("Healthz is sleeping shutdown thread for "+SLEEP_UPON_SHUTDOWN+" ms.");
             try {
-                Thread.sleep(SLEEP_UPON_SHUTDOWN);
+                if (sleepTime>0) {
+                    logger.debug("Healthz is sleeping shutdown thread for "+sleepTime+" ms.");
+                    Thread.sleep(sleepTime);
+                }
             } catch (InterruptedException e) {
                 logger.warn("Shutdown sleep interrupted.", e);
             }
@@ -60,4 +63,7 @@ public String getHealthz(HttpServletResponse response) throws Exception {
         }
     }
 
+    public long getSleepTime() {
+        return sleepTime;
+    }
 }
diff --git a/server/src/test/java/org/cloudfoundry/identity/uaa/web/HealthzEndpointTests.java b/server/src/test/java/org/cloudfoundry/identity/uaa/web/HealthzEndpointTests.java
@@ -23,7 +23,9 @@
 
 public class HealthzEndpointTests {
 
-    private HealthzEndpoint endpoint = new HealthzEndpoint();
+    private static final int SLEEP_UPON_SHUTDOWN = 150;
+
+    private HealthzEndpoint endpoint = new HealthzEndpoint(SLEEP_UPON_SHUTDOWN);
     private MockHttpServletResponse response = new MockHttpServletResponse();
 
     @Test
@@ -39,7 +41,18 @@ public void shutdown_sends_stopping() throws Exception {
         assertEquals("stopping\n", endpoint.getHealthz(response));
         assertEquals(503, response.getStatus());
         long after = System.currentTimeMillis();
-        assertThat(after, Matchers.greaterThanOrEqualTo(now+5000));
+        assertThat(after, Matchers.greaterThanOrEqualTo(now+SLEEP_UPON_SHUTDOWN));
+    }
+
+    @Test
+    public void shutdown_without_sleep() throws Exception {
+        long now = System.currentTimeMillis();
+        endpoint = new HealthzEndpoint(-1);
+        runShutdownHook();
+        assertEquals("stopping\n", endpoint.getHealthz(response));
+        assertEquals(503, response.getStatus());
+        long after = System.currentTimeMillis();
+        assertThat(after, Matchers.lessThanOrEqualTo(now+SLEEP_UPON_SHUTDOWN));
     }
 
     protected void runShutdownHook() {

diff --git a/uaa/src/main/webapp/WEB-INF/spring-servlet.xml b/uaa/src/main/webapp/WEB-INF/spring-servlet.xml
@@ -479,7 +479,9 @@
         <property name="globalLinks" ref="globalLinks"/>
     </bean>
 
-    <bean id="healthzEndpoint" class="org.cloudfoundry.identity.uaa.health.HealthzEndpoint" />
+    <bean id="healthzEndpoint" class="org.cloudfoundry.identity.uaa.health.HealthzEndpoint">
+        <constructor-arg name="sleepTime" value="${uaa.shutdown.sleep:-1}"/>
+    </bean>
 
     <context:annotation-config />
 

diff --git a/uaa/src/test/java/org/cloudfoundry/identity/uaa/login/BootstrapTests.java b/uaa/src/test/java/org/cloudfoundry/identity/uaa/login/BootstrapTests.java
@@ -17,6 +17,7 @@
 import org.cloudfoundry.identity.uaa.authentication.manager.AuthzAuthenticationManager;
 import org.cloudfoundry.identity.uaa.authentication.manager.PeriodLockoutPolicy;
 import org.cloudfoundry.identity.uaa.constants.OriginKeys;
+import org.cloudfoundry.identity.uaa.health.HealthzEndpoint;
 import org.cloudfoundry.identity.uaa.home.HomeController;
 import org.cloudfoundry.identity.uaa.impl.config.IdentityZoneConfigurationBootstrap;
 import org.cloudfoundry.identity.uaa.impl.config.YamlServletProfileInitializer;
@@ -201,6 +202,9 @@ public void defaults_and_required_properties() throws Exception {
 
         context = getServletContext(profiles, false, new String[] {"login.yml", "uaa.yml", "required_configuration.yml"}, "file:./src/main/webapp/WEB-INF/spring-servlet.xml");
 
+        HealthzEndpoint hend = context.getBean(HealthzEndpoint.class);
+        assertEquals(-1, hend.getSleepTime());
+
         UaaMetricsFilter metricsFilter = context.getBean(UaaMetricsFilter.class);
         assertTrue(metricsFilter.isEnabled());
 
@@ -451,6 +455,9 @@ public void all_properties_set() throws Exception {
         String profiles = System.getProperty("spring.profiles.active");
         context = getServletContext(profiles, false, new String[] {"login.yml", "uaa.yml", "test/bootstrap/all-properties-set.yml"}, "file:./src/main/webapp/WEB-INF/spring-servlet.xml");
 
+        HealthzEndpoint hend = context.getBean(HealthzEndpoint.class);
+        assertEquals(5000, hend.getSleepTime());
+
         UaaMetricsFilter metricsFilter = context.getBean(UaaMetricsFilter.class);
         assertFalse(metricsFilter.isEnabled());
 

diff --git a/uaa/src/test/resources/test/bootstrap/all-properties-set.yml b/uaa/src/test/resources/test/bootstrap/all-properties-set.yml
@@ -565,6 +565,8 @@ smtp:
   starttls: true
 uaa:
   url: https://uaa.some.test.domain.com:555/uaa
+  shutdown:
+    sleep: 5000
   limitedFunctionality:
     enabled: true
     whitelist: