fix: flaky Consul embedded test startup

embedded-tests/src/test/java/org/apache/druid/testing/embedded/consul/ConsulClusterResource.java

@@ -19,6 +19,7 @@
 
 package org.apache.druid.testing.embedded.consul;
 
+import com.fasterxml.jackson.databind.ObjectMapper;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.testing.embedded.EmbeddedDruidCluster;
@@ -31,7 +32,16 @@
 import org.testcontainers.utility.DockerImageName;
 
 import javax.annotation.Nullable;
+import javax.net.ssl.KeyManagerFactory;
+import javax.net.ssl.SSLContext;
+import javax.net.ssl.TrustManagerFactory;
+import java.io.FileInputStream;
+import java.io.IOException;
 import java.net.URI;
+import java.net.http.HttpClient;
+import java.net.http.HttpRequest;
+import java.net.http.HttpResponse;
+import java.security.KeyStore;
 import java.time.Duration;
 
 /**
@@ -41,9 +51,12 @@
 public class ConsulClusterResource extends TestcontainerResource<GenericContainer<?>>
 {
   private static final Logger log = new Logger(ConsulClusterResource.class);
+  private static final ObjectMapper JSON_MAPPER = new ObjectMapper();
   private static final int CONSUL_HTTP_PORT = 8500;
   private static final int CONSUL_HTTPS_PORT = 8501;
   private static final DockerImageName CONSUL_IMAGE = DockerImageName.parse("hashicorp/consul:1.18");
+  private static final Duration READINESS_TIMEOUT = Duration.ofSeconds(30);
+  private static final Duration READINESS_RETRY_DELAY = Duration.ofMillis(500);
 
   private final ConsulSecurityMode securityMode;
   private String consulHostForDruid;
@@ -70,6 +83,13 @@ public ConsulClusterResource(ConsulSecurityMode securityMode)
     this.securityMode = securityMode;
   }
 
+  @Override
+  public void start()
+  {
+    super.start();
+    waitForConsulApi();
+  }
+
   @Override
   protected GenericContainer<?> createContainer()
   {
@@ -106,6 +126,111 @@ protected GenericContainer<?> createContainer()
     }
   }
 
+  private void waitForConsulApi()
+  {
+    final long deadline = System.nanoTime() + READINESS_TIMEOUT.toNanos();
+    final HttpClient httpClient;
+    Exception lastException = null;
+
+    try {
+      httpClient = createHttpClient();
+    }
+    catch (Exception e) {
+      throw new RuntimeException("Failed to create Consul readiness client", e);
+    }
+
+    while (System.nanoTime() < deadline) {
+      try {
+        final HttpRequest request = HttpRequest.newBuilder(getHttpUri("/v1/status/leader"))
+                                               .timeout(Duration.ofSeconds(5))
+                                               .GET()
+                                               .build();
+        final HttpResponse<String> response = httpClient.send(request, HttpResponse.BodyHandlers.ofString());
+        if (isConsulLeaderReady(response.statusCode(), response.body())) {
+          log.info("Consul API is ready at [%s].", getHttpUri("/v1/status/leader"));
+          return;
+        }
+        lastException = new RuntimeException(
+            StringUtils.format(
+                "Consul leader endpoint returned status[%d] body[%s]",
+                response.statusCode(),
+                response.body()
+            )
+        );
+      }
+      catch (Exception e) {
+        lastException = e;
+      }
+
+      try {
+        Thread.sleep(READINESS_RETRY_DELAY.toMillis());
+      }
+      catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new RuntimeException("Interrupted while waiting for Consul API readiness", e);
+      }
+    }
+
+    throw new RuntimeException(
+        StringUtils.format("Consul API did not become ready within [%s]", READINESS_TIMEOUT),
+        lastException
+    );
+  }
+
+  static boolean isConsulLeaderReady(int statusCode, @Nullable String body)
+  {
+    if (statusCode != 200 || body == null || body.trim().isEmpty()) {
+      return false;
+    }
+
+    try {
+      final String leader = JSON_MAPPER.readValue(body, String.class);
+      return leader != null && !leader.trim().isEmpty();
+    }
+    catch (IOException e) {
+      return false;
+    }
+  }
+
+  private HttpClient createHttpClient() throws Exception
+  {
+    if (securityMode == ConsulSecurityMode.PLAIN) {
+      return HttpClient.newBuilder()
+                       .connectTimeout(Duration.ofSeconds(5))
+                       .build();
+    }
+
+    if (certBundle == null) {
+      throw new IllegalStateException("Consul TLS certificate bundle is not initialized");
+    }
+
+    final KeyStore trustStore = KeyStore.getInstance("PKCS12");
+    try (FileInputStream fis = new FileInputStream(certBundle.getTrustStorePath())) {
+      trustStore.load(fis, getStorePassword().toCharArray());
+    }
+
+    final TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
+    tmf.init(trustStore);
+
+    KeyManagerFactory kmf = null;
+    if (securityMode == ConsulSecurityMode.MTLS) {
+      final KeyStore keyStore = KeyStore.getInstance("PKCS12");
+      try (FileInputStream fis = new FileInputStream(certBundle.getKeyStorePath())) {
+        keyStore.load(fis, getStorePassword().toCharArray());
+      }
+      kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
+      kmf.init(keyStore, getStorePassword().toCharArray());
+    }
+
+    final SSLContext sslContext = SSLContext.getInstance("TLS");
+    sslContext.init(kmf == null ? null : kmf.getKeyManagers(), tmf.getTrustManagers(), null);
+
+    return HttpClient.newBuilder()
+                     .sslContext(sslContext)
+                     .connectTimeout(Duration.ofSeconds(5))
+                     .build();
+  }
+
   @Override
   public void onStarted(EmbeddedDruidCluster cluster)
   {

embedded-tests/src/test/java/org/apache/druid/testing/embedded/consul/ConsulClusterResourceTest.java

@@ -0,0 +1,58 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.apache.druid.testing.embedded.consul;
+
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+
+public class ConsulClusterResourceTest
+{
+  @Test
+  public void testLeaderEndpointWithLeaderIsReady()
+  {
+    Assertions.assertTrue(ConsulClusterResource.isConsulLeaderReady(200, "\"127.0.0.1:8300\""));
+  }
+
+  @Test
+  public void testLeaderEndpointWithoutLeaderIsNotReady()
+  {
+    Assertions.assertFalse(ConsulClusterResource.isConsulLeaderReady(200, "\"\""));
+  }
+
+  @Test
+  public void testBlankBodyIsNotReady()
+  {
+    Assertions.assertFalse(ConsulClusterResource.isConsulLeaderReady(200, ""));
+    Assertions.assertFalse(ConsulClusterResource.isConsulLeaderReady(200, "   "));
+    Assertions.assertFalse(ConsulClusterResource.isConsulLeaderReady(200, null));
+  }
+
+  @Test
+  public void testNonOkStatusIsNotReady()
+  {
+    Assertions.assertFalse(ConsulClusterResource.isConsulLeaderReady(503, "\"127.0.0.1:8300\""));
+  }
+
+  @Test
+  public void testMalformedBodyIsNotReady()
+  {
+    Assertions.assertFalse(ConsulClusterResource.isConsulLeaderReady(200, "127.0.0.1:8300"));
+  }
+}

extensions-contrib/consul-extensions/src/main/java/org/apache/druid/consul/discovery/ConsulDruidNodeAnnouncer.java

@@ -20,19 +20,22 @@
 package org.apache.druid.consul.discovery;
 
 import com.google.common.base.Preconditions;
+import com.google.common.base.Throwables;
 import com.google.inject.Inject;
 import org.apache.druid.concurrent.LifecycleLock;
 import org.apache.druid.discovery.DiscoveryDruidNode;
 import org.apache.druid.discovery.DruidNodeAnnouncer;
 import org.apache.druid.guice.ManageLifecycle;
 import org.apache.druid.java.util.common.ISE;
+import org.apache.druid.java.util.common.RetryUtils;
 import org.apache.druid.java.util.common.concurrent.Execs;
 import org.apache.druid.java.util.common.lifecycle.LifecycleStart;
 import org.apache.druid.java.util.common.lifecycle.LifecycleStop;
 import org.apache.druid.java.util.common.logger.Logger;
 import org.apache.druid.java.util.emitter.service.ServiceEmitter;
 
 import javax.annotation.Nullable;
+import java.io.IOException;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
@@ -48,6 +51,7 @@
 public class ConsulDruidNodeAnnouncer implements DruidNodeAnnouncer
 {
   private static final Logger LOGGER = new Logger(ConsulDruidNodeAnnouncer.class);
+  private static final int MAX_ANNOUNCE_REGISTRATION_TRIES = 3;
 
   private final ConsulApiClient consulApiClient;
   private final ConsulDiscoveryConfig config;
@@ -159,7 +163,7 @@ public void announce(DiscoveryDruidNode discoveryDruidNode)
       long registerStart = System.nanoTime();
 
       // Register in Consul, then track locally atomically in this block
-      consulApiClient.registerService(discoveryDruidNode);
+      registerServiceWithRetry(serviceId, discoveryDruidNode);
       announcedNodes.put(serviceId, discoveryDruidNode);
 
       long registerLatency = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - registerStart);
@@ -197,6 +201,32 @@ public void announce(DiscoveryDruidNode discoveryDruidNode)
     }
   }
 
+  private void registerServiceWithRetry(String serviceId, DiscoveryDruidNode discoveryDruidNode) throws Exception
+  {
+    RetryUtils.retry(
+        () -> {
+          consulApiClient.registerService(discoveryDruidNode);
+          return null;
+        },
+        ConsulDruidNodeAnnouncer::isTransientConsulFailure,
+        1,
+        MAX_ANNOUNCE_REGISTRATION_TRIES,
+        null,
+        "Registering Consul service [" + serviceId + "] failed"
+    );
+  }
+
+  private static boolean isTransientConsulFailure(Throwable throwable)
+  {
+    for (Throwable cause : Throwables.getCausalChain(throwable)) {
+      if (cause instanceof IOException) {
+        return true;
+      }
+    }
+
+    return false;
+  }
+
   @Override
   public void unannounce(DiscoveryDruidNode discoveryDruidNode)
   {

extensions-contrib/consul-extensions/src/test/java/org/apache/druid/consul/discovery/ConsulDruidNodeAnnouncerTest.java

@@ -22,6 +22,7 @@
 import org.apache.druid.discovery.DiscoveryDruidNode;
 import org.apache.druid.discovery.NodeRole;
 import org.apache.druid.server.DruidNode;
+import org.apache.http.NoHttpResponseException;
 import org.easymock.Capture;
 import org.easymock.EasyMock;
 import org.joda.time.Duration;
@@ -95,6 +96,32 @@ public void testAnnounce() throws Exception
     EasyMock.verify(mockConsulApiClient);
   }
 
+  @Test
+  public void testAnnounceRetriesTransientFailure() throws Exception
+  {
+    mockConsulApiClient.registerService(EasyMock.eq(testNode));
+    EasyMock.expectLastCall().andThrow(new NoHttpResponseException("Consul did not respond"));
+
+    mockConsulApiClient.registerService(EasyMock.eq(testNode));
+    EasyMock.expectLastCall().once();
+
+    mockConsulApiClient.passTtlCheck(EasyMock.anyString(), EasyMock.anyString());
+    EasyMock.expectLastCall().anyTimes();
+
+    mockConsulApiClient.deregisterService(EasyMock.anyString());
+    EasyMock.expectLastCall().once();
+
+    EasyMock.replay(mockConsulApiClient);
+
+    announcer = new ConsulDruidNodeAnnouncer(mockConsulApiClient, config);
+    announcer.start();
+
+    announcer.announce(testNode);
+    announcer.stop();
+
+    EasyMock.verify(mockConsulApiClient);
+  }
+
   @Test
   public void testUnannounce() throws Exception
   {