Cache status checks.

Also adds some tests for the /status/disable and /status/enable endpoints.

server/src/main/java/keywhiz/KeywhizConfig.java

@@ -21,6 +21,7 @@
 import io.dropwizard.db.DataSourceFactory;
 import io.dropwizard.db.ManagedDataSource;
 import java.io.IOException;
+import java.time.Duration;
 import java.util.Optional;
 import javax.validation.Valid;
 import javax.validation.constraints.NotNull;
@@ -81,6 +82,9 @@ public class KeywhizConfig extends Configuration {
   @JsonProperty
   private String migrationsDir;
 
+  @JsonProperty
+  private String statusCacheExpiry;
+
   public String getEnvironment() {
     return environment;
   }
@@ -126,6 +130,14 @@ public String getMigrationsDir() {
     return migrationsDir;
   }
 
+  public Duration getStatusCacheExpiry() {
+    if ((statusCacheExpiry == null) || (statusCacheExpiry.isEmpty())) {
+      // Default to 1 second
+      return Duration.ofSeconds(1);
+    }
+    return Duration.parse(statusCacheExpiry);
+  }
+
   /** @return LDAP configuration to authenticate admin users. Absent if fakeLdap is true. */
   public UserAuthenticatorFactory getUserAuthenticatorFactory() {
     return userAuth;

server/src/main/java/keywhiz/service/resources/StatusResource.java

@@ -3,17 +3,23 @@
 import com.codahale.metrics.annotation.ExceptionMetered;
 import com.codahale.metrics.annotation.Timed;
 import com.codahale.metrics.health.HealthCheck;
-import com.codahale.metrics.health.HealthCheckRegistry;
+import com.google.common.base.Supplier;
+import com.google.common.base.Suppliers;
 import io.dropwizard.setup.Environment;
+
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
 import java.util.SortedMap;
+import java.util.concurrent.TimeUnit;
 import javax.inject.Inject;
 import javax.ws.rs.GET;
 import javax.ws.rs.Path;
 import javax.ws.rs.Produces;
 import javax.ws.rs.core.Response;
+
+import keywhiz.KeywhizConfig;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -24,14 +30,40 @@
 @Path("/_status")
 @Produces(APPLICATION_JSON)
 public class StatusResource {
-  private Environment environment;
   private static final Logger logger = LoggerFactory.getLogger(SecretDeliveryResource.class);
+  Supplier<SortedMap<String, HealthCheck.Result>> memoizedCheck;
 
-  @Inject public StatusResource(Environment environment) {
-    this.environment = environment;
+  @Inject public StatusResource(KeywhizConfig keywhizConfig, Environment environment) {
+    Duration cacheExpiry = keywhizConfig.getStatusCacheExpiry();
+    memoizedCheck = Suppliers.memoizeWithExpiration(() -> environment.healthChecks().runHealthChecks(),
+        cacheExpiry.toMillis(), TimeUnit.MILLISECONDS);
+  }
+
+  @Timed @ExceptionMetered
+  @GET
+  public Response get() {
+    SortedMap<String, HealthCheck.Result> results = memoizedCheck.get();
+
+    List<String> failing = results.entrySet().stream()
+        .filter(r -> !r.getValue().isHealthy())
+        .map(Map.Entry::getKey)
+        .collect(toList());
+
+    if (!failing.isEmpty()) {
+      logger.warn("Health checks failed: {}", results);
+      String message = "failing health checks: " + Arrays.toString(failing.toArray());
+      StatusResponse sr = new StatusResponse("critical", message, results);
+      return Response.serverError().entity(sr).build();
+    }
+    StatusResponse sr = new StatusResponse("ok", "ok", results);
+    return Response.ok(sr).build();
   }
 
   public static class StatusResponse {
+    private String status;
+    private String message;
+    private SortedMap<String, HealthCheck.Result> results;
+
     public SortedMap<String, HealthCheck.Result> getResults() {
       return results;
     }
@@ -44,14 +76,11 @@ public String getStatus() {
       return status;
     }
 
-    private String status;
-    private String message;
-    private SortedMap<String, HealthCheck.Result> results;
-
     StatusResponse(String status, String message, SortedMap<String, HealthCheck.Result> results) {
       this.status = status;
       this.message = message;
       this.results = results;
+
     }
 
     @Override public String toString() {
@@ -62,25 +91,4 @@ public String getStatus() {
           '}';
     }
   }
-
-  @Timed @ExceptionMetered
-  @GET
-  public Response get() {
-    HealthCheckRegistry checks = this.environment.healthChecks();
-    SortedMap<String, HealthCheck.Result> results = checks.runHealthChecks();
-
-    List<String> failing = results.entrySet().stream()
-        .filter(r -> !r.getValue().isHealthy())
-        .map(Map.Entry::getKey)
-        .collect(toList());
-
-    if (!failing.isEmpty()) {
-      logger.warn("Health checks failed: {}", results);
-      String message = "failing health checks: " + Arrays.toString(failing.toArray());
-      StatusResponse sr = new StatusResponse("critical", message, results);
-      return Response.serverError().entity(sr).build();
-    }
-    StatusResponse sr = new StatusResponse("ok", "ok", results);
-    return Response.ok(sr).build();
-  }
 }

server/src/main/resources/keywhiz-development.yaml.docker

@@ -78,6 +78,8 @@ readonlyDatabase:
 migrationsDir:
   db/h2/migration
 
+statusCacheExpiry: PT1S
+
 userAuth:
   type: bcrypt
 

server/src/main/resources/keywhiz-development.yaml.h2

@@ -88,6 +88,8 @@ readonlyDatabase:
 migrationsDir:
   db/h2/migration
 
+statusCacheExpiry: PT1S
+
 userAuth:
   type: bcrypt
 

server/src/main/resources/keywhiz-development.yaml.mysql

@@ -88,6 +88,8 @@ readonlyDatabase:
 migrationsDir:
   db/mysql/migration
 
+statusCacheExpiry: PT1S
+
 userAuth:
   type: bcrypt
 

server/src/main/resources/keywhiz-development.yaml.postgres

@@ -86,6 +86,8 @@ readonlyDatabase:
 migrationsDir:
   db/postgres/migration
 
+statusCacheExpiry: PT1S
+
 userAuth:
   type: bcrypt
 

server/src/test/java/keywhiz/service/resources/StatusResourceIntegrationTest.java

@@ -0,0 +1,115 @@
+/*
+ * Copyright (C) 2015 Square, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package keywhiz.service.resources;
+
+import keywhiz.IntegrationTestRule;
+import keywhiz.TestClients;
+import okhttp3.*;
+import okhttp3.MediaType;
+import okhttp3.Request;
+import okhttp3.Response;
+import org.junit.Before;
+import org.junit.ClassRule;
+import org.junit.Test;
+import org.junit.rules.RuleChain;
+
+import static keywhiz.testing.HttpClients.testUrl;
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class StatusResourceIntegrationTest {
+  OkHttpClient httpClient, httpsClient;
+
+  @ClassRule
+  public static final RuleChain chain = IntegrationTestRule.rule();
+
+  @Before
+  public void setUp() throws Exception {
+    httpClient = new OkHttpClient();
+    httpsClient = TestClients.unauthenticatedClient();
+  }
+
+  private boolean isHealthy() throws Exception {
+    Request get = new Request.Builder()
+        .get()
+        .url(testUrl("/_status"))
+        .build();
+    okhttp3.Response statusResponse = httpsClient.newCall(get).execute();
+    return statusResponse.code() == 200;
+  }
+
+  @Test
+  public void successWhenHealthy() throws Exception {
+    Request request = new Request.Builder()
+        .url("http://localhost:8081/status/enable")
+        .post(RequestBody.create(MediaType.parse("text/plain"), ""))
+        .build();
+    Response disableResponse = httpClient.newCall(request).execute();
+    assertThat(disableResponse.code()).isEqualTo(200);
+    Thread.sleep(3500);
+    assertThat(isHealthy()).isTrue();
+  }
+
+  @Test
+  public void failsWhenUnhealthy() throws Exception {
+    Request request = new Request.Builder()
+        .url("http://localhost:8081/status/disable")
+        .post(RequestBody.create(MediaType.parse("text/plain"), ""))
+        .build();
+    Response disableResponse = httpClient.newCall(request).execute();
+    assertThat(disableResponse.code()).isEqualTo(200);
+
+    assertThat(isHealthy()).isFalse();
+  }
+
+  @Test
+  public void cachesStatusCheck() throws Exception {
+    // Start healthy
+    Request request = new Request.Builder()
+        .url("http://localhost:8081/status/enable")
+        .post(RequestBody.create(MediaType.parse("text/plain"), ""))
+        .build();
+    Response enableResponse = httpClient.newCall(request).execute();
+    assertThat(enableResponse.code()).isEqualTo(200);
+    Thread.sleep(3500);
+    assertThat(isHealthy()).isTrue();
+
+    // Make the service unhealthy
+    request = new Request.Builder()
+        .url("http://localhost:8081/status/disable")
+        .post(RequestBody.create(MediaType.parse("text/plain"), ""))
+        .build();
+    Response disableResponse = httpClient.newCall(request).execute();
+    assertThat(disableResponse.code()).isEqualTo(200);
+
+    // We should get back the cached healthy result. Sleep and query again should return unhealthy.
+    assertThat(isHealthy()).isTrue();
+    Thread.sleep(3500);
+    assertThat(isHealthy()).isFalse();
+
+    // Make the service healthy
+    request = new Request.Builder()
+        .url("http://localhost:8081/status/enable")
+        .post(RequestBody.create(MediaType.parse("text/plain"), ""))
+        .build();
+    enableResponse = httpClient.newCall(request).execute();
+    assertThat(enableResponse.code()).isEqualTo(200);
+
+    // We should get back the cached unhealthy result. Sleep and query again should return healthy.
+    assertThat(isHealthy()).isFalse();
+    Thread.sleep(3500);
+    assertThat(isHealthy()).isTrue();
+  }
+}

server/src/test/java/keywhiz/service/resources/StatusResponseTest.java → server/src/test/java/keywhiz/service/resources/StatusResourceTest.java

@@ -3,28 +3,36 @@
 import com.codahale.metrics.health.HealthCheck;
 import com.codahale.metrics.health.HealthCheckRegistry;
 import io.dropwizard.setup.Environment;
+
+import java.time.Duration;
 import java.util.TreeMap;
 import javax.ws.rs.InternalServerErrorException;
 import javax.ws.rs.core.Response;
+
+import keywhiz.KeywhizConfig;
 import org.junit.Before;
 import org.junit.Test;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
-public class StatusResponseTest {
+public class StatusResourceTest {
   HealthCheckRegistry registry;
   Environment environment;
   StatusResource status;
+  KeywhizConfig keywhizConfig;
 
   @Before
   public void setUp() throws Exception {
     this.registry = mock(HealthCheckRegistry.class);
     this.environment = mock(Environment.class);
-    this.status = new StatusResource(environment);
+    this.keywhizConfig = mock(KeywhizConfig.class);
 
     when(environment.healthChecks()).thenReturn(registry);
+    when(keywhizConfig.getStatusCacheExpiry()).thenReturn(Duration.ofSeconds(1));
+
+    this.status = new StatusResource(keywhizConfig, environment);
   }
 
   @Test

server/src/test/resources/keywhiz-test.yaml.docker

@@ -72,6 +72,8 @@ readonlyDatabase:
 migrationsDir:
   db/h2/migration
 
+statusCacheExpiry: PT3S
+
 userAuth:
   type: bcrypt
 

server/src/test/resources/keywhiz-test.yaml.h2

@@ -72,6 +72,8 @@ readonlyDatabase:
 migrationsDir:
   db/h2/migration
 
+statusCacheExpiry: PT3S
+
 userAuth:
   type: bcrypt
 

server/src/test/resources/keywhiz-test.yaml.mysql

@@ -72,6 +72,8 @@ readonlyDatabase:
 migrationsDir:
   db/mysql/migration
 
+statusCacheExpiry: PT3S
+
 userAuth:
   type: bcrypt
 

server/src/test/resources/keywhiz-test.yaml.postgres

@@ -70,6 +70,8 @@ readonlyDatabase:
 migrationsDir:
   db/postgres/migration
 
+statusCacheExpiry: PT3S
+
 userAuth:
   type: bcrypt