adding actuator health check support for Elasticsearch REST Clients

spring-boot-project/spring-boot-actuator-autoconfigure/pom.xml

@@ -269,6 +269,11 @@
 			<artifactId>elasticsearch</artifactId>
 			<optional>true</optional>
 		</dependency>
+		<dependency>
+		    <groupId>org.elasticsearch.client</groupId>
+		    <artifactId>elasticsearch-rest-client</artifactId>
+		    <optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>org.flywaydb</groupId>
 			<artifactId>flyway-core</artifactId>

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/elasticsearch/ElasticSearchRestHealthIndicatorAutoConfiguration.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2012-2018 the original author or authors.
+ *
+ * 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 org.springframework.boot.actuate.autoconfigure.elasticsearch;
+
+import java.util.Map;
+
+import org.elasticsearch.client.RestClient;
+
+import org.springframework.boot.actuate.autoconfigure.health.CompositeHealthIndicatorConfiguration;
+import org.springframework.boot.actuate.autoconfigure.health.ConditionalOnEnabledHealthIndicator;
+import org.springframework.boot.actuate.autoconfigure.health.HealthIndicatorAutoConfiguration;
+import org.springframework.boot.actuate.elasticsearch.ElasticsearchRestHealthIndicator;
+import org.springframework.boot.actuate.health.HealthIndicator;
+import org.springframework.boot.autoconfigure.AutoConfigureAfter;
+import org.springframework.boot.autoconfigure.AutoConfigureBefore;
+import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
+import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.elasticsearch.rest.RestClientAutoConfiguration;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
+
+/**
+ * {@link EnableAutoConfiguration Auto-configuration} for
+ * {@link ElasticsearchRestHealthIndicator} using the {@link RestClient}.
+ *
+ * @author Artsiom Yudovin
+ * @since 2.1.0
+ */
+
+@Configuration
+@ConditionalOnClass(RestClient.class)
+@ConditionalOnBean(RestClient.class)
+@ConditionalOnEnabledHealthIndicator("elasticsearch")
+@AutoConfigureBefore(HealthIndicatorAutoConfiguration.class)
+@AutoConfigureAfter({ RestClientAutoConfiguration.class,
+		ElasticSearchClientHealthIndicatorAutoConfiguration.class })
+public class ElasticSearchRestHealthIndicatorAutoConfiguration extends
+		CompositeHealthIndicatorConfiguration<ElasticsearchRestHealthIndicator, RestClient> {
+
+	private final Map<String, RestClient> clients;
+
+	public ElasticSearchRestHealthIndicatorAutoConfiguration(
+			Map<String, RestClient> clients) {
+		this.clients = clients;
+	}
+
+	@Bean
+	@ConditionalOnMissingBean(name = "elasticsearchRestHealthIndicator")
+	public HealthIndicator elasticsearchRestHealthIndicator() {
+		return createHealthIndicator(this.clients);
+	}
+
+	@Override
+	protected ElasticsearchRestHealthIndicator createHealthIndicator(RestClient client) {
+		return new ElasticsearchRestHealthIndicator(client);
+	}
+
+}

spring-boot-project/spring-boot-actuator/pom.xml

@@ -77,6 +77,11 @@
 			<artifactId>jest</artifactId>
 			<optional>true</optional>
 		</dependency>
+		<dependency>
+			<groupId>org.elasticsearch.client</groupId>
+			<artifactId>elasticsearch-rest-client</artifactId>
+			<optional>true</optional>
+		</dependency>
 		<dependency>
 			<groupId>io.undertow</groupId>
 			<artifactId>undertow-servlet</artifactId>
@@ -338,5 +343,6 @@
 			<artifactId>jsonassert</artifactId>
 			<scope>test</scope>
 		</dependency>
+
 	</dependencies>
 </project>

spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/elasticsearch/ElasticsearchRestHealthIndicator.java

@@ -0,0 +1,74 @@
+/*
+ * Copyright 2012-2018 the original author or authors.
+ *
+ * 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 org.springframework.boot.actuate.elasticsearch;
+
+import java.io.InputStreamReader;
+import java.nio.charset.StandardCharsets;
+
+import com.google.gson.JsonElement;
+import com.google.gson.JsonParser;
+import org.apache.http.HttpStatus;
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.client.RestClient;
+
+import org.springframework.boot.actuate.health.AbstractHealthIndicator;
+import org.springframework.boot.actuate.health.Health;
+import org.springframework.boot.actuate.health.HealthIndicator;
+
+/**
+ * {@link HealthIndicator} for an Elasticsearch cluster by REST.
+ *
+ * @author Artsiom Yudovin
+ * @since 2.1.0
+ */
+public class ElasticsearchRestHealthIndicator extends AbstractHealthIndicator {
+
+	private final RestClient client;
+
+	private final JsonParser jsonParser = new JsonParser();
+
+	public ElasticsearchRestHealthIndicator(RestClient client) {
+		super("Elasticsearch health check failed");
+		this.client = client;
+	}
+
+	@Override
+	protected void doHealthCheck(Health.Builder builder) throws Exception {
+		Response response = this.client
+				.performRequest(new Request("GET", "/_cluster/health/"));
+
+		if (response.getStatusLine().getStatusCode() != HttpStatus.SC_OK) {
+			builder.down();
+		}
+		else {
+			try (InputStreamReader reader = new InputStreamReader(
+					response.getEntity().getContent(), StandardCharsets.UTF_8)) {
+				JsonElement root = this.jsonParser.parse(reader);
+				JsonElement status = root.getAsJsonObject().get("status");
+				if (status.getAsString()
+						.equals(io.searchbox.cluster.Health.Status.RED.getKey())) {
+					builder.outOfService();
+				}
+				else {
+					builder.up();
+				}
+			}
+		}
+	}
+
+}

spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/elasticsearch/ElasticsearchRestHealthIndicatorTest.java

@@ -0,0 +1,127 @@
+/*
+ * Copyright 2012-2018 the original author or authors.
+ *
+ * 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 org.springframework.boot.actuate.elasticsearch;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+
+import org.apache.http.StatusLine;
+import org.apache.http.entity.BasicHttpEntity;
+import org.elasticsearch.client.Request;
+import org.elasticsearch.client.Response;
+import org.elasticsearch.client.RestClient;
+import org.junit.Test;
+
+import org.springframework.boot.actuate.health.Status;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.BDDMockito.mock;
+import static org.mockito.BDDMockito.when;
+
+/**
+ * Tests for {@link ElasticsearchRestHealthIndicator}.
+ *
+ * @author Artsiom Yudovin
+ */
+public class ElasticsearchRestHealthIndicatorTest {
+
+	private final RestClient restClient = mock(RestClient.class);
+
+	private final ElasticsearchRestHealthIndicator elasticsearchRestHealthIndicator = new ElasticsearchRestHealthIndicator(
+			this.restClient);
+
+	@Test
+	public void elasticsearchIsUp() throws IOException {
+		BasicHttpEntity httpEntity = new BasicHttpEntity();
+		httpEntity.setContent(
+				new ByteArrayInputStream(createJsonResult(200, "green").getBytes()));
+
+		Response response = mock(Response.class);
+		StatusLine statusLine = mock(StatusLine.class);
+
+		when(statusLine.getStatusCode()).thenReturn(200);
+		when(response.getStatusLine()).thenReturn(statusLine);
+		when(response.getEntity()).thenReturn(httpEntity);
+		when(this.restClient.performRequest(any(Request.class))).thenReturn(response);
+
+		assertThat(this.elasticsearchRestHealthIndicator.health().getStatus())
+				.isEqualTo(Status.UP);
+	}
+
+	@Test
+	public void elasticsearchIsDown() throws IOException {
+		when(this.restClient.performRequest(any(Request.class)))
+				.thenThrow(new IOException("Couldn't connect"));
+
+		assertThat(this.elasticsearchRestHealthIndicator.health().getStatus())
+				.isEqualTo(Status.DOWN);
+	}
+
+	@Test
+	public void elasticsearchIsDownByResponseCode() throws IOException {
+
+		Response response = mock(Response.class);
+		StatusLine statusLine = mock(StatusLine.class);
+
+		when(statusLine.getStatusCode()).thenReturn(500);
+		when(response.getStatusLine()).thenReturn(statusLine);
+		when(this.restClient.performRequest(any(Request.class))).thenReturn(response);
+
+		assertThat(this.elasticsearchRestHealthIndicator.health().getStatus())
+				.isEqualTo(Status.DOWN);
+	}
+
+	@Test
+	public void elasticsearchIsOutOfServiceByStatus() throws IOException {
+		BasicHttpEntity httpEntity = new BasicHttpEntity();
+		httpEntity.setContent(
+				new ByteArrayInputStream(createJsonResult(200, "red").getBytes()));
+
+		Response response = mock(Response.class);
+		StatusLine statusLine = mock(StatusLine.class);
+
+		when(statusLine.getStatusCode()).thenReturn(200);
+		when(response.getStatusLine()).thenReturn(statusLine);
+		when(response.getEntity()).thenReturn(httpEntity);
+		when(this.restClient.performRequest(any(Request.class))).thenReturn(response);
+
+		assertThat(this.elasticsearchRestHealthIndicator.health().getStatus())
+				.isEqualTo(Status.OUT_OF_SERVICE);
+	}
+
+	private String createJsonResult(int responseCode, String status) {
+		String json;
+		if (responseCode == 200) {
+			json = String.format("{\"cluster_name\":\"elasticsearch\","
+					+ "\"status\":\"%s\",\"timed_out\":false,\"number_of_nodes\":1,"
+					+ "\"number_of_data_nodes\":1,\"active_primary_shards\":0,"
+					+ "\"active_shards\":0,\"relocating_shards\":0,\"initializing_shards\":0,"
+					+ "\"unassigned_shards\":0,\"delayed_unassigned_shards\":0,"
+					+ "\"number_of_pending_tasks\":0,\"number_of_in_flight_fetch\":0,"
+					+ "\"task_max_waiting_in_queue_millis\":0,\"active_shards_percent_as_number\":100.0}",
+					status);
+		}
+		else {
+			json = "{\n" + "  \"error\": \"Server Error\",\n" + "  \"status\": "
+					+ responseCode + "\n" + "}";
+		}
+
+		return json;
+	}
+
+}