Merge pull request #262 from GetFeedback/feature/introduce-kafka-streams-health-checks

Feature/introduce kafka streams health checks

Author: Stefano Guerrini

kahpp-spring-autoconfigure/build.gradle

@@ -30,7 +30,6 @@ dependencies{
     implementation "io.burt:jmespath-core:0.5.1"
     implementation "io.burt:jmespath-jackson:0.5.1"
     implementation "io.vavr:vavr:1.0.0-alpha-4"
-    implementation "com.deviceinsight.kafka:kafka-health-check:1.3.0"
     implementation "org.springframework.kafka:spring-kafka"
     implementation "org.springframework.boot:spring-boot-starter"
     implementation "org.springframework.boot:spring-boot-starter-validation"

kahpp-spring-autoconfigure/src/main/java/dev/vox/platform/kahpp/actuator/KafkaStreamsProducerState.java

@@ -0,0 +1,37 @@
+package dev.vox.platform.kahpp.actuator;
+
+import io.micrometer.core.instrument.FunctionCounter;
+import io.micrometer.core.instrument.MeterRegistry;
+import io.micrometer.core.instrument.search.Search;
+import java.util.Optional;
+import org.springframework.boot.actuate.health.AbstractHealthIndicator;
+import org.springframework.boot.actuate.health.Health;
+import org.springframework.boot.actuate.health.Status;
+import org.springframework.stereotype.Component;
+
+@Component
+public class KafkaStreamsProducerState extends AbstractHealthIndicator {
+
+  private final MeterRegistry meterRegistry;
+
+  public static final String KAFKA_PRODUCER_CONNECTION_METRIC_LABEL =
+      "kafka.producer.connection.count";
+
+  public KafkaStreamsProducerState(MeterRegistry meterRegistry) {
+    this.meterRegistry = meterRegistry;
+  }
+
+  @Override
+  protected void doHealthCheck(Health.Builder builder) {
+    Search connectionsSearch = meterRegistry.find(KAFKA_PRODUCER_CONNECTION_METRIC_LABEL);
+    Double kafkaConnections =
+        Optional.ofNullable(connectionsSearch.functionCounter())
+            .map(FunctionCounter::count)
+            .orElse(0d);
+    if (kafkaConnections > 0) {
+      builder.status(Status.UP);
+      return;
+    }
+    builder.status(Status.DOWN);
+  }
+}

kahpp-spring-autoconfigure/src/main/java/dev/vox/platform/kahpp/actuator/KafkaStreamsState.java

@@ -0,0 +1,47 @@
+package dev.vox.platform.kahpp.actuator;
+
+import java.util.List;
+import java.util.Objects;
+import org.apache.kafka.streams.KafkaStreams;
+import org.springframework.boot.actuate.health.AbstractHealthIndicator;
+import org.springframework.boot.actuate.health.Health;
+import org.springframework.boot.actuate.health.Status;
+import org.springframework.kafka.config.StreamsBuilderFactoryBean;
+import org.springframework.stereotype.Component;
+
+@Component
+public class KafkaStreamsState extends AbstractHealthIndicator {
+
+  private final List<StreamsBuilderFactoryBean> streamsBuilders;
+
+  public KafkaStreamsState(List<StreamsBuilderFactoryBean> streamsBuilders) {
+    this.streamsBuilders = List.copyOf(streamsBuilders);
+  }
+
+  @Override
+  @SuppressWarnings("PMD.CloseResource")
+  protected void doHealthCheck(Health.Builder builder) {
+    List<KafkaStreams> streamsList =
+        streamsBuilders.stream()
+            .map(StreamsBuilderFactoryBean::getKafkaStreams)
+            .filter(Objects::nonNull)
+            .toList();
+
+    if (streamsList.isEmpty()) {
+      builder.status(Status.UNKNOWN);
+      return;
+    }
+
+    for (KafkaStreams streams : streamsList) {
+      if (streams.state().hasNotStarted()) {
+        builder.status(Status.UNKNOWN);
+        return;
+      } else if (!streams.state().isRunningOrRebalancing()) {
+        builder.status(Status.DOWN);
+        return;
+      }
+    }
+
+    builder.status(Status.UP);
+  }
+}

kahpp-spring-autoconfigure/src/main/resources/application.properties

@@ -14,4 +14,6 @@ management.endpoint.metrics.enabled=true
 management.endpoint.health.show-details=always
 management.metrics.export.jmx.domain=dev.vox
 management.endpoints.web.exposure.include=metrics,health,prometheus
+management.endpoint.health.group.readiness.include=readinessState,kafkaStreamsProducerState,kafkaStreamsState
+management.endpoint.health.group.liveness.include=livenessState,kafkaStreamsProducerState,kafkaStreamsState
 spring.config.import=${KAHPP_CONFIG_LOCATION:file:/kahpp/application.yaml}

kahpp-spring-autoconfigure/src/test/java/dev/vox/platform/kahpp/actuator/KafkaStreamsProducerStateTest.java

@@ -0,0 +1,43 @@
+package dev.vox.platform.kahpp.actuator;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import io.micrometer.core.instrument.*;
+import io.micrometer.core.instrument.simple.SimpleMeterRegistry;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.boot.actuate.health.Health;
+import org.springframework.boot.actuate.health.Status;
+
+@ExtendWith(MockitoExtension.class)
+class KafkaStreamsProducerStateTest {
+
+  private MeterRegistry meterRegistry;
+
+  @BeforeEach
+  public void setUp() {
+    meterRegistry = new SimpleMeterRegistry();
+  }
+
+  @Test
+  void statusDownWhenProducerIsNotActive() {
+    meterRegistry
+        .more()
+        .counter(KafkaStreamsProducerState.KAFKA_PRODUCER_CONNECTION_METRIC_LABEL, Tags.empty(), 0);
+    Health.Builder builder = new Health.Builder();
+    new KafkaStreamsProducerState(meterRegistry).doHealthCheck(builder);
+    assertThat(builder.build().getStatus()).isEqualTo(Status.DOWN);
+  }
+
+  @Test
+  void statusUpWhenProducerIsActive() {
+    meterRegistry
+        .more()
+        .counter(KafkaStreamsProducerState.KAFKA_PRODUCER_CONNECTION_METRIC_LABEL, Tags.empty(), 1);
+    Health.Builder builder = new Health.Builder();
+    new KafkaStreamsProducerState(meterRegistry).doHealthCheck(builder);
+    assertThat(builder.build().getStatus()).isEqualTo(Status.UP);
+  }
+}

kahpp-spring-autoconfigure/src/test/java/dev/vox/platform/kahpp/actuator/KafkaStreamsStateTest.java

@@ -0,0 +1,66 @@
+package dev.vox.platform.kahpp.actuator;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.lenient;
+import static org.mockito.Mockito.when;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.Map;
+import org.apache.kafka.streams.KafkaStreams;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.Mock;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.boot.actuate.health.Health;
+import org.springframework.boot.actuate.health.Status;
+import org.springframework.kafka.config.StreamsBuilderFactoryBean;
+
+@ExtendWith(MockitoExtension.class)
+class KafkaStreamsStateTest {
+
+  @Mock private StreamsBuilderFactoryBean streamsBuilderFactoryBean;
+
+  @Mock private KafkaStreams kafkaStreams;
+
+  Map<KafkaStreams.State, Status> kafkaStateMapping =
+      Map.of(
+          KafkaStreams.State.CREATED, Status.UNKNOWN,
+          KafkaStreams.State.RUNNING, Status.UP,
+          KafkaStreams.State.REBALANCING, Status.UP,
+          KafkaStreams.State.ERROR, Status.DOWN,
+          KafkaStreams.State.PENDING_ERROR, Status.DOWN,
+          KafkaStreams.State.PENDING_SHUTDOWN, Status.DOWN,
+          KafkaStreams.State.NOT_RUNNING, Status.DOWN);
+
+  @BeforeEach
+  public void setUp() {
+    lenient().when(streamsBuilderFactoryBean.getKafkaStreams()).thenReturn(kafkaStreams);
+  }
+
+  @Test
+  void checkIfStatusIsUnknownWhenStreamsIsNull() {
+    Health.Builder builder = new Health.Builder();
+    new KafkaStreamsState(List.of()).doHealthCheck(builder);
+    assertThat(builder.build().getStatus()).isEqualTo(Status.UNKNOWN);
+  }
+
+  @Test
+  @SuppressWarnings("PMD.JUnitTestsShouldIncludeAssert")
+  void checkAllKafkaStreamsStatuses() {
+    KafkaStreamsState kafkaStreamsState = new KafkaStreamsState(List.of(streamsBuilderFactoryBean));
+    Health.Builder builder = new Health.Builder();
+
+    new KafkaStreamsState(List.of(streamsBuilderFactoryBean));
+    Arrays.stream(KafkaStreams.State.values())
+        .forEach(
+            state -> {
+              when(kafkaStreams.state()).thenReturn(state);
+              kafkaStreamsState.doHealthCheck(builder);
+              assertThat(builder.build().getStatus())
+                  .as("Testing %s Kafka Stream State", state)
+                  .isEqualTo(kafkaStateMapping.get(state));
+            });
+  }
+}