Fix service re-registration after Consul connection failures

docs/modules/ROOT/pages/discovery.adoc

@@ -234,6 +234,10 @@ spring:
           enabled: true
           ttl: 10s
 ----
+Service re-registration after a heartbeat failure is disabled by default. To enable it, set
+`spring.cloud.consul.discovery.heartbeat.reregister-service-on-failure` to `true`.
+This option allows Spring Cloud Consul to attempt to re-register the service when a
+heartbeat failure meets the configured re-registration criteria.
 
 [[ttl-application-status]]
 ==== TTL Application Status

spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/discovery/TtlScheduler.java

@@ -33,6 +33,8 @@
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.scheduling.concurrent.ConcurrentTaskScheduler;
 import org.springframework.web.client.HttpStatusCodeException;
+import org.springframework.web.client.ResourceAccessException;
+import org.springframework.web.client.RestClientException;
 
 /**
  * Created by nicu on 11.03.2015.
@@ -154,6 +156,11 @@ private void possiblyReregisterIfFails(Runnable consulClientCall) {
 			try {
 				consulClientCall.run();
 			}
+			catch (ResourceAccessException e) {
+				log.warn("Unable to send heartbeat to Consul for service " + this.serviceId
+						+ ". Will retry on the next heartbeat. Reason: " + e.getMessage());
+				log.debug("Consul heartbeat failed for service " + this.serviceId, e);
+			}
 			catch (HttpStatusCodeException e) {
 				if (this.ttlScheduler.heartbeatProperties.isReregisterServiceOnFailure()
 						&& this.ttlScheduler.reregistrationPredicate.isEligible(e)) {
@@ -163,8 +170,17 @@ private void possiblyReregisterIfFails(Runnable consulClientCall) {
 						if (log.isInfoEnabled()) {
 							log.info("Re-register " + registeredService);
 						}
-						this.ttlScheduler.client.agentServiceRegister(
-								this.ttlScheduler.discoveryProperties.getAclToken(), registeredService);
+						try {
+							this.ttlScheduler.client.agentServiceRegister(
+									this.ttlScheduler.discoveryProperties.getAclToken(), registeredService);
+						}
+						catch (RestClientException registrationException) {
+							log.warn("Unable to re-register service " + this.serviceId
+									+ " with Consul. Will retry on the next heartbeat. Reason: "
+									+ registrationException.getMessage());
+							log.debug("Consul service re-registration failed for service " + this.serviceId,
+									registrationException);
+						}
 					}
 					else {
 						log.warn("The service to re-register is not found.");

spring-cloud-consul-discovery/src/main/java/org/springframework/cloud/consul/serviceregistry/ConsulServiceRegistry.java

@@ -62,13 +62,7 @@ public ConsulServiceRegistry(ConsulClient client, ConsulDiscoveryProperties prop
 	public void register(ConsulRegistration reg) {
 		log.info("Registering service with consul: " + reg.getService());
 		try {
-			ResponseEntity<Void> response = this.client.agentServiceRegister(this.properties.getAclToken(),
-					reg.getService());
-			NewService service = reg.getService();
-			if (this.heartbeatProperties.isEnabled() && this.ttlScheduler != null && service.getCheck() != null
-					&& service.getCheck().getTtl() != null) {
-				this.ttlScheduler.add(reg.getService());
-			}
+			this.client.agentServiceRegister(this.properties.getAclToken(), reg.getService());
 		}
 		catch (RuntimeException e) {
 			if (this.properties.isFailFast()) {
@@ -77,6 +71,12 @@ public void register(ConsulRegistration reg) {
 			}
 			log.warn("Failfast is false. Error registering service with consul: " + reg.getService(), e);
 		}
+
+		NewService service = reg.getService();
+		if (this.heartbeatProperties.isEnabled() && this.ttlScheduler != null && service.getCheck() != null
+				&& service.getCheck().getTtl() != null) {
+			this.ttlScheduler.add(service);
+		}
 	}
 
 	@Override

spring-cloud-consul-discovery/src/test/java/org/springframework/cloud/consul/discovery/ConsulHeartbeatTaskTests.java

@@ -35,12 +35,17 @@
 import org.springframework.http.HttpStatus;
 import org.springframework.web.client.HttpClientErrorException;
 import org.springframework.web.client.HttpServerErrorException;
+import org.springframework.web.client.ResourceAccessException;
+import org.springframework.web.client.RestClientException;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
 /**
@@ -217,4 +222,45 @@ public void disableReRegistration() {
 		assertThatThrownBy(consulHeartbeatTask::run).isSameAs(operationException);
 	}
 
+	@Test
+	public void connectionFailureDoesNotPreventSubsequentHeartbeatAttempts() {
+		TtlScheduler ttlScheduler = new TtlScheduler(heartbeatProperties, discoveryProperties, consulClient,
+				ReregistrationPredicate.DEFAULT, applicationStatusProviders);
+		ConsulHeartbeatTask consulHeartbeatTask = new ConsulHeartbeatTask(serviceId, ttlScheduler,
+				() -> Check.CheckStatus.PASSING);
+
+		given(consulClient.agentCheckPass("service:" + serviceId, null, null))
+			.willThrow(new ResourceAccessException("Consul is unavailable"));
+
+		assertThatCode(consulHeartbeatTask::run).doesNotThrowAnyException();
+		assertThatCode(consulHeartbeatTask::run).doesNotThrowAnyException();
+
+		verify(consulClient, times(2)).agentCheckPass("service:" + serviceId, null, null);
+	}
+
+	@Test
+	public void reRegistrationFailureDoesNotPreventSubsequentHeartbeatAttempts() {
+		TtlScheduler ttlScheduler = new TtlScheduler(heartbeatProperties, discoveryProperties, consulClient,
+				ReregistrationPredicate.DEFAULT, applicationStatusProviders);
+		ConsulHeartbeatTask consulHeartbeatTask = new ConsulHeartbeatTask(serviceId, ttlScheduler,
+				() -> Check.CheckStatus.PASSING);
+		heartbeatProperties.setReregisterServiceOnFailure(true);
+
+		NewService service = new NewService();
+		service.setId(serviceId);
+		ttlScheduler.add(service);
+
+		given(consulClient.agentCheckPass("service:" + serviceId, null, null))
+			.willThrow(new HttpServerErrorException(HttpStatus.SERVICE_UNAVAILABLE));
+
+		given(consulClient.agentServiceRegister(null, service))
+			.willThrow(new RestClientException("Consul is unavailable"));
+
+		assertThatCode(consulHeartbeatTask::run).doesNotThrowAnyException();
+		assertThatCode(consulHeartbeatTask::run).doesNotThrowAnyException();
+
+		verify(consulClient, atLeast(2)).agentCheckPass("service:" + serviceId, null, null);
+		verify(consulClient, atLeast(2)).agentServiceRegister(null, service);
+	}
+
 }

spring-cloud-consul-discovery/src/test/java/org/springframework/cloud/consul/serviceregistry/ConsulServiceRegistryRegistrationTests.java

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2013-present 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
+ *
+ *      https://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.cloud.consul.serviceregistry;
+
+import org.junit.jupiter.api.Test;
+
+import org.springframework.cloud.consul.ConsulClient;
+import org.springframework.cloud.consul.discovery.ConsulDiscoveryProperties;
+import org.springframework.cloud.consul.discovery.HeartbeatProperties;
+import org.springframework.cloud.consul.discovery.TtlScheduler;
+import org.springframework.cloud.consul.model.http.agent.NewService;
+
+import static org.assertj.core.api.Assertions.assertThatCode;
+import static org.mockito.BDDMockito.given;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+
+class ConsulServiceRegistryRegistrationTests {
+
+	@Test
+	void schedulesHeartbeatWhenInitialRegistrationFailsAndFailFastIsDisabled() {
+		ConsulClient consulClient = mock(ConsulClient.class);
+		ConsulDiscoveryProperties discoveryProperties = mock(ConsulDiscoveryProperties.class);
+		HeartbeatProperties heartbeatProperties = new HeartbeatProperties();
+		heartbeatProperties.setEnabled(true);
+		TtlScheduler ttlScheduler = mock(TtlScheduler.class);
+
+		given(discoveryProperties.isFailFast()).willReturn(false);
+
+		ConsulServiceRegistry serviceRegistry = new ConsulServiceRegistry(consulClient, discoveryProperties,
+				ttlScheduler, heartbeatProperties);
+
+		NewService service = new NewService();
+		service.setId("test-service");
+		service.setName("test-service");
+
+		NewService.Check check = new NewService.Check();
+		check.setTtl("10s");
+		service.setCheck(check);
+
+		ConsulRegistration registration = new ConsulRegistration(service, discoveryProperties);
+
+		given(consulClient.agentServiceRegister(discoveryProperties.getAclToken(), service))
+			.willThrow(new IllegalStateException("Consul is unavailable"));
+
+		assertThatCode(() -> serviceRegistry.register(registration)).doesNotThrowAnyException();
+
+		verify(ttlScheduler).add(service);
+	}
+
+	@Test
+	void doesNotScheduleHeartbeatWhenInitialRegistrationFailsAndFailFastIsEnabled() {
+		ConsulClient consulClient = mock(ConsulClient.class);
+		ConsulDiscoveryProperties discoveryProperties = mock(ConsulDiscoveryProperties.class);
+		HeartbeatProperties heartbeatProperties = new HeartbeatProperties();
+		heartbeatProperties.setEnabled(true);
+		TtlScheduler ttlScheduler = mock(TtlScheduler.class);
+
+		given(discoveryProperties.isFailFast()).willReturn(true);
+
+		ConsulServiceRegistry serviceRegistry = new ConsulServiceRegistry(consulClient, discoveryProperties,
+				ttlScheduler, heartbeatProperties);
+
+		NewService service = new NewService();
+		service.setId("test-service");
+		service.setName("test-service");
+
+		NewService.Check check = new NewService.Check();
+		check.setTtl("10s");
+		service.setCheck(check);
+
+		ConsulRegistration registration = new ConsulRegistration(service, discoveryProperties);
+
+		given(consulClient.agentServiceRegister(discoveryProperties.getAclToken(), service))
+			.willThrow(new IllegalStateException("Consul is unavailable"));
+
+		assertThatCode(() -> serviceRegistry.register(registration)).isInstanceOf(IllegalStateException.class)
+			.hasMessage("Consul is unavailable");
+
+		verify(ttlScheduler, org.mockito.Mockito.never()).add(service);
+	}
+
+}