Retry locator operation if locator no available

Author: acogoluegnes

src/main/java/com/rabbitmq/stream/impl/ConsumersCoordinator.java

@@ -115,14 +115,11 @@ Runnable subscribe(
     return subscriptionTracker::cancel;
   }
 
-  private Client locator() {
-    return environment.locator();
-  }
-
   // package protected for testing
   List<Client.Broker> findBrokersForStream(String stream) {
     // FIXME make sure locator is not null (retry)
-    Map<String, Client.StreamMetadata> metadata = locator().metadata(stream);
+    Map<String, Client.StreamMetadata> metadata =
+        this.environment.locatorOperation(c -> c.metadata(stream));
     if (metadata.size() == 0 || metadata.get(stream) == null) {
       // this is not supposed to happen
       throw new StreamDoesNotExistException(stream);

src/main/java/com/rabbitmq/stream/impl/HashRoutingStrategy.java

@@ -38,8 +38,7 @@ class HashRoutingStrategy implements RoutingStrategy {
       ToIntFunction<String> hash) {
     this.routingKeyExtractor = routingKeyExtractor;
     this.env = env;
-    // TODO use async retry to get locator
-    List<String> ps = this.env.locator().partitions(superStream);
+    List<String> ps = this.env.locatorOperation(c -> c.partitions(superStream));
     this.partitions =
         new CopyOnWriteArrayList<>(
             ps.stream().map(Collections::singletonList).collect(Collectors.toList()));

src/main/java/com/rabbitmq/stream/impl/ProducersCoordinator.java

@@ -95,7 +95,8 @@ private Runnable registerAgentTracker(AgentTracker tracker, String stream) {
   }
 
   private Client.Broker getBrokerForProducer(String stream) {
-    Map<String, Client.StreamMetadata> metadata = this.environment.locator().metadata(stream);
+    Map<String, Client.StreamMetadata> metadata =
+        this.environment.locatorOperation(c -> c.metadata(stream));
     if (metadata.size() == 0 || metadata.get(stream) == null) {
       throw new StreamDoesNotExistException(stream);
     }

src/main/java/com/rabbitmq/stream/impl/RoutingKeyRoutingStrategy.java

@@ -44,7 +44,7 @@ public List<String> route(Message message) {
             routingKey,
             routingKey1 -> {
               // TODO retry on locator lookup
-              return env.locator().route(routingKey1, superStream);
+              return env.locatorOperation(c -> c.route(routingKey1, superStream));
             });
     return streams;
   }

src/main/java/com/rabbitmq/stream/impl/StreamEnvironment.java

@@ -55,6 +55,7 @@
 import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.LongConsumer;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 import javax.net.ssl.SSLException;
 import org.slf4j.Logger;
@@ -477,15 +478,53 @@ Runnable registerProducer(StreamProducer producer, String reference, String stre
     return producersCoordinator.registerProducer(producer, reference, stream);
   }
 
-  // FIXME make the locator available as a completable future (with retry)
-  // this would make client code more robust
   Client locator() {
     if (this.locator == null) {
-      throw new StreamException("No connection available");
+      throw new LocatorNotAvailableException();
     }
     return this.locator;
   }
 
+  <T> T locatorOperation(Function<Client, T> operation) {
+    return locatorOperation(operation, () -> locator(), this.recoveryBackOffDelayPolicy);
+  }
+
+  static <T> T locatorOperation(
+      Function<Client, T> operation,
+      Supplier<Client> clientSupplier,
+      BackOffDelayPolicy backOffDelayPolicy) {
+    int maxAttempt = 3;
+    int attempt = 0;
+    boolean executed = false;
+    Exception lastException = null;
+    T result = null;
+    while (attempt < maxAttempt) {
+      try {
+        result = operation.apply(clientSupplier.get());
+        executed = true;
+        break;
+      } catch (LocatorNotAvailableException e) {
+        attempt++;
+        try {
+          Thread.sleep(backOffDelayPolicy.delay(attempt).toMillis());
+        } catch (InterruptedException ex) {
+          lastException = ex;
+          Thread.currentThread().interrupt();
+          break;
+        }
+      }
+    }
+    if (!executed) {
+      if (lastException == null) {
+        throw new LocatorNotAvailableException();
+      } else {
+        throw new StreamException(
+            "Could not execute operation after " + maxAttempt + " attempts", lastException);
+      }
+    }
+    return result;
+  }
+
   Clock clock() {
     return this.clock;
   }
@@ -561,4 +600,11 @@ public Consumer<Context> postMessageProcessingCallback() {
       return postMessageProcessingCallback;
     }
   }
+
+  static class LocatorNotAvailableException extends StreamException {
+
+    public LocatorNotAvailableException() {
+      super("Locator not available");
+    }
+  }
 }

src/main/java/com/rabbitmq/stream/impl/StreamStreamCreator.java

@@ -69,7 +69,7 @@ public void create() {
     }
     this.environment.maybeInitializeLocator();
     Client.Response response =
-        environment.locator().create(stream, streamParametersBuilder.build());
+        environment.locatorOperation(c -> c.create(stream, streamParametersBuilder.build()));
     if (!response.isOk()
         && response.getResponseCode() != Constants.RESPONSE_CODE_STREAM_ALREADY_EXISTS) {
       throw new StreamException(

src/main/java/com/rabbitmq/stream/impl/SuperStreamConsumer.java

@@ -29,7 +29,7 @@ class SuperStreamConsumer implements Consumer {
   SuperStreamConsumer(
       StreamConsumerBuilder builder, String superStream, StreamEnvironment environment) {
     this.superStream = superStream;
-    for (String partition : environment.locator().partitions(superStream)) {
+    for (String partition : environment.locatorOperation(c -> c.partitions(superStream))) {
       Consumer consumer = builder.duplicate().superStream(null).stream(partition).build();
       consumers.put(partition, consumer);
       LOGGER.debug("Created consumer on stream '{}' for super stream '{}'", partition, superStream);

src/main/java/com/rabbitmq/stream/impl/SuperStreamProducer.java

@@ -61,11 +61,12 @@ public MessageBuilder messageBuilder() {
   @Override
   public long getLastPublishingId() {
     if (this.name != null && !this.name.isEmpty()) {
-      List<String> streams = this.environment.locator().partitions(superStream);
+      List<String> streams = this.environment.locatorOperation(c -> c.partitions(superStream));
       long publishingId = 0;
       boolean first = true;
       for (String partition : streams) {
-        long pubId = this.environment.locator().queryPublisherSequence(this.name, partition);
+        long pubId =
+            this.environment.locatorOperation(c -> c.queryPublisherSequence(this.name, partition));
         if (first) {
           publishingId = pubId;
           first = false;

src/test/java/com/rabbitmq/stream/impl/ConsumersCoordinatorTest.java

@@ -130,6 +130,7 @@ public Client.ClientParameters shutdownListener(
         };
     mocks = MockitoAnnotations.openMocks(this);
     when(environment.locator()).thenReturn(locator);
+    when(environment.locatorOperation(any())).thenCallRealMethod();
     when(environment.clientParametersCopy()).thenReturn(clientParameters);
     when(environment.addressResolver()).thenReturn(address -> address);
 

src/test/java/com/rabbitmq/stream/impl/ProducersCoordinatorTest.java

@@ -105,6 +105,7 @@ public Client.ClientParameters metadataListener(
         };
     mocks = MockitoAnnotations.openMocks(this);
     when(environment.locator()).thenReturn(locator);
+    when(environment.locatorOperation(any())).thenCallRealMethod();
     when(environment.clientParametersCopy()).thenReturn(clientParameters);
     when(environment.addressResolver()).thenReturn(address -> address);
     when(trackingConsumer.stream()).thenReturn("stream");

src/test/java/com/rabbitmq/stream/impl/StreamEnvironmentTest.java

@@ -373,34 +373,21 @@ void environmentPublishersConsumersShouldCloseSuccessfullyWhenBrokerIsDown() thr
   void locatorShouldReconnectIfConnectionIsLost(TestInfo info) throws Exception {
     try (Environment environment =
         environmentBuilder
-            .recoveryBackOffDelayPolicy(BackOffDelayPolicy.fixed(Duration.ofSeconds(2)))
+            .recoveryBackOffDelayPolicy(BackOffDelayPolicy.fixed(Duration.ofSeconds(1)))
             .build()) {
       String s = streamName(info);
       environment.streamCreator().stream(s).create();
       environment.deleteStream(s);
       Host.killConnection("rabbitmq-stream-locator");
-      assertThatThrownBy(() -> environment.streamCreator().stream("whatever").create())
-          .isInstanceOf(StreamException.class);
-      assertThatThrownBy(() -> environment.deleteStream("whatever"))
-          .isInstanceOf(StreamException.class);
-      assertThatThrownBy(() -> environment.producerBuilder().stream(stream).build())
-          .isInstanceOf(StreamException.class);
-      assertThatThrownBy(() -> environment.consumerBuilder().stream(stream).build())
-          .isInstanceOf(StreamException.class);
-
-      Producer producer = null;
-      int timeout = 10_000;
-      int waited = 0;
-      int interval = 1_000;
-      while (producer == null && waited < timeout) {
-        try {
-          Thread.sleep(interval);
-          waited += interval;
-          producer = environment.producerBuilder().stream(stream).build();
-        } catch (StreamException e) {
-        }
+      environment.streamCreator().stream(s).create();
+      try {
+        Producer producer = environment.producerBuilder().stream(s).build();
+        Consumer consumer = environment.consumerBuilder().stream(s).build();
+        producer.close();
+        consumer.close();
+      } finally {
+        environment.deleteStream(s);
       }
-      assertThat(producer).isNotNull();
     }
   }
 

src/test/java/com/rabbitmq/stream/impl/StreamEnvironmentUnitTest.java

@@ -13,18 +13,25 @@
 // info@rabbitmq.com.
 package com.rabbitmq.stream.impl;
 
+import static com.rabbitmq.stream.impl.TestUtils.latchAssert;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.*;
 
 import com.rabbitmq.stream.BackOffDelayPolicy;
+import com.rabbitmq.stream.StreamException;
 import com.rabbitmq.stream.impl.Client.ClientParameters;
+import com.rabbitmq.stream.impl.StreamEnvironment.LocatorNotAvailableException;
 import io.netty.buffer.ByteBufAllocator;
 import java.net.URI;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executors;
 import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 import org.junit.jupiter.api.AfterEach;
@@ -173,4 +180,92 @@ void shouldNotOpenConnectionWhenLazyInitIsEnabled(
             cf);
     verify(cf, times(expectedConnectionCreation)).apply(any(Client.ClientParameters.class));
   }
+
+  @Test
+  void locatorOperationShouldReturnOperationResultIfNoProblem() {
+    AtomicInteger counter = new AtomicInteger();
+    int result =
+        StreamEnvironment.locatorOperation(
+            c -> counter.incrementAndGet(), () -> null, BackOffDelayPolicy.fixed(Duration.ZERO));
+    assertThat(result).isEqualTo(1);
+  }
+
+  @Test
+  void locatorOperationShouldRetryAndReturnResultIfLocatorException() {
+    AtomicInteger counter = new AtomicInteger();
+    int result =
+        StreamEnvironment.locatorOperation(
+            c -> {
+              if (counter.incrementAndGet() < 2) {
+                throw new LocatorNotAvailableException();
+              } else {
+                return counter.get();
+              }
+            },
+            () -> null,
+            BackOffDelayPolicy.fixed(Duration.ofMillis(10)));
+    assertThat(result).isEqualTo(2);
+  }
+
+  @Test
+  void locatorOperationShouldThrowLocatorExceptionWhenRetryExhausts() {
+    AtomicInteger counter = new AtomicInteger();
+    assertThatThrownBy(
+            () ->
+                StreamEnvironment.locatorOperation(
+                    c -> {
+                      counter.incrementAndGet();
+                      throw new LocatorNotAvailableException();
+                    },
+                    () -> null,
+                    BackOffDelayPolicy.fixed(Duration.ofMillis(10))))
+        .isInstanceOf(LocatorNotAvailableException.class);
+    assertThat(counter).hasValue(3);
+  }
+
+  @Test
+  void locatorOperationShouldThrowInterruptedExceptionAsCauseIfInterrupted()
+      throws InterruptedException {
+    CountDownLatch latch = new CountDownLatch(1);
+    AtomicReference<Exception> exception = new AtomicReference<>();
+    Thread thread =
+        new Thread(
+            () -> {
+              try {
+                StreamEnvironment.locatorOperation(
+                    c -> {
+                      latch.countDown();
+                      throw new LocatorNotAvailableException();
+                    },
+                    () -> null,
+                    BackOffDelayPolicy.fixed(Duration.ofMinutes(10)));
+              } catch (StreamException e) {
+                exception.set(e);
+              }
+            });
+    thread.start();
+    latchAssert(latch).completes();
+    Thread.sleep(100);
+    thread.interrupt();
+    Thread.sleep(100);
+    assertThat(exception.get())
+        .isInstanceOf(StreamException.class)
+        .hasCauseInstanceOf(InterruptedException.class);
+  }
+
+  @Test
+  void locatorOperationShouldNotRetryAndReThrowUnexpectedException() {
+    AtomicInteger counter = new AtomicInteger();
+    assertThatThrownBy(
+            () ->
+                StreamEnvironment.locatorOperation(
+                    c -> {
+                      counter.incrementAndGet();
+                      throw new RuntimeException();
+                    },
+                    () -> null,
+                    BackOffDelayPolicy.fixed(Duration.ofMillis(10))))
+        .isInstanceOf(RuntimeException.class);
+    assertThat(counter).hasValue(1);
+  }
 }

src/test/java/com/rabbitmq/stream/impl/StreamProducerUnitTest.java

@@ -111,6 +111,7 @@ void init() {
         .thenCallRealMethod();
     when(env.scheduledExecutorService()).thenReturn(executorService);
     when(env.locator()).thenReturn(client);
+    when(env.locatorOperation(any())).thenCallRealMethod();
     when(env.clock()).thenReturn(clock);
     when(env.codec()).thenReturn(new SimpleCodec());
     doAnswer(