From da7482f9a14e2c26dfc0c0d907d2193568a8e76b Mon Sep 17 00:00:00 2001
From: lambcode <5820125+lambcode@users.noreply.github.com>
Date: Sun, 23 Jun 2019 03:46:50 -0600
Subject: [PATCH] Fixes ThrottledStorageComponent not passing through health
 check (#2632)

---
 .../server/internal/ZipkinHealthIndicator.java       |  3 +--
 .../internal/throttle/ThrottledStorageComponent.java |  7 ++++++-
 .../throttle/ThrottledStorageComponentTest.kt        | 12 +++++++++++-
 3 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/zipkin-server/src/main/java/zipkin2/server/internal/ZipkinHealthIndicator.java b/zipkin-server/src/main/java/zipkin2/server/internal/ZipkinHealthIndicator.java
index aabd737c1d6..ebed8ff1f95 100644
--- a/zipkin-server/src/main/java/zipkin2/server/internal/ZipkinHealthIndicator.java
+++ b/zipkin-server/src/main/java/zipkin2/server/internal/ZipkinHealthIndicator.java
@@ -27,8 +27,7 @@ final class ZipkinHealthIndicator extends CompositeHealthIndicator {
   }
 
   void addComponent(Component component) {
-    String healthName = component.getClass().getSimpleName();
-    healthName = healthName.replace("AutoValue_", "");
+    String healthName = component.toString();
     addHealthIndicator(healthName, new ComponentHealthIndicator(component));
   }
 
diff --git a/zipkin-server/src/main/java/zipkin2/server/internal/throttle/ThrottledStorageComponent.java b/zipkin-server/src/main/java/zipkin2/server/internal/throttle/ThrottledStorageComponent.java
index d95e845da3e..6ea8b46f6c2 100644
--- a/zipkin-server/src/main/java/zipkin2/server/internal/throttle/ThrottledStorageComponent.java
+++ b/zipkin-server/src/main/java/zipkin2/server/internal/throttle/ThrottledStorageComponent.java
@@ -30,6 +30,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 import zipkin2.Call;
+import zipkin2.CheckResult;
 import zipkin2.Span;
 import zipkin2.storage.SpanConsumer;
 import zipkin2.storage.SpanStore;
@@ -97,8 +98,12 @@ public ThrottledStorageComponent(StorageComponent delegate, MeterRegistry regist
     delegate.close();
   }
 
+  @Override public CheckResult check() {
+    return delegate.check();
+  }
+
   @Override public String toString() {
-    return "Throttled(" + delegate + ")";
+    return "Throttled(" + delegate.toString() + ")";
   }
 
   static final class ThrottledSpanConsumer implements SpanConsumer {
diff --git a/zipkin-server/src/test/kotlin/zipkin2/server/internal/throttle/ThrottledStorageComponentTest.kt b/zipkin-server/src/test/kotlin/zipkin2/server/internal/throttle/ThrottledStorageComponentTest.kt
index bde579f2efe..95fa44a1758 100644
--- a/zipkin-server/src/test/kotlin/zipkin2/server/internal/throttle/ThrottledStorageComponentTest.kt
+++ b/zipkin-server/src/test/kotlin/zipkin2/server/internal/throttle/ThrottledStorageComponentTest.kt
@@ -16,7 +16,11 @@ package zipkin2.server.internal.throttle
 import com.linecorp.armeria.common.metric.NoopMeterRegistry
 import org.assertj.core.api.Assertions.assertThat
 import org.junit.Test
+import org.mockito.Mockito.mock
+import org.mockito.Mockito.times
+import org.mockito.Mockito.verify
 import zipkin2.storage.InMemoryStorage
+import zipkin2.storage.StorageComponent
 
 class ThrottledStorageComponentTest {
   val delegate = InMemoryStorage.newBuilder().build()
@@ -43,6 +47,12 @@ class ThrottledStorageComponentTest {
 
   @Test fun niceToString() {
     assertThat(ThrottledStorageComponent(delegate, registry, 1, 2, 1))
-      .hasToString("Throttled(InMemoryStorage{traceCount=0})");
+      .hasToString("Throttled(InMemoryStorage{traceCount=0})")
+  }
+
+  @Test fun delegatesCheck() {
+    val mock = mock(StorageComponent::class.java)
+    ThrottledStorageComponent(mock, registry, 1, 2, 1).check()
+    verify(mock, times(1)).check()
   }
 }