Add super stream SAC balancing test

References rabbitmq/rabbitmq-server#3753

Author: acogoluegnes

src/main/java/com/rabbitmq/stream/ConsumerUpdateListener.java

@@ -21,6 +21,8 @@ interface Context {
 
     Consumer consumer();
 
+    String stream();
+
     Status status();
 
     Status previousStatus();

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

@@ -250,6 +250,7 @@ private SubscriptionTracker(
       } else {
         Map<String, String> properties = new ConcurrentHashMap<>(subscriptionProperties.size() + 1);
         properties.putAll(subscriptionProperties);
+        // we propagate the subscription name, used for monitoring
         properties.put("name", this.offsetTrackingReference);
         this.subscriptionProperties = Collections.unmodifiableMap(properties);
       }

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

@@ -27,11 +27,13 @@
 import java.time.Duration;
 import java.util.Collections;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.LongConsumer;
 import java.util.function.LongSupplier;
@@ -40,7 +42,10 @@
 
 class StreamConsumer implements Consumer {
 
+  private static final AtomicLong ID_SEQUENCE = new AtomicLong(0);
+
   private static final Logger LOGGER = LoggerFactory.getLogger(StreamConsumer.class);
+  private final long id;
   private final Runnable closingTrackingCallback;
   private final AtomicBoolean closed = new AtomicBoolean(false);
   private final String name;
@@ -67,6 +72,7 @@ class StreamConsumer implements Consumer {
       Map<String, String> subscriptionProperties,
       ConsumerUpdateListener consumerUpdateListener) {
 
+    this.id = ID_SEQUENCE.getAndIncrement();
     try {
       this.name = name;
       this.stream = stream;
@@ -289,12 +295,12 @@ OffsetSpecification consumerUpdate(boolean active) {
 
   private static class DefaultConsumerUpdateContext implements ConsumerUpdateListener.Context {
 
-    private final Consumer consumer;
+    private final StreamConsumer consumer;
     private final ConsumerUpdateListener.Status status;
     private final ConsumerUpdateListener.Status previousStatus;
 
     private DefaultConsumerUpdateContext(
-        Consumer consumer,
+        StreamConsumer consumer,
         ConsumerUpdateListener.Status status,
         ConsumerUpdateListener.Status previousStatus) {
       this.consumer = consumer;
@@ -307,6 +313,11 @@ public Consumer consumer() {
       return this.consumer;
     }
 
+    @Override
+    public String stream() {
+      return this.consumer.stream;
+    }
+
     @Override
     public ConsumerUpdateListener.Status status() {
       return this.status;
@@ -316,6 +327,20 @@ public ConsumerUpdateListener.Status status() {
     public ConsumerUpdateListener.Status previousStatus() {
       return this.previousStatus;
     }
+
+    @Override
+    public String toString() {
+      return "DefaultConsumerUpdateContext{"
+          + "consumer="
+          + consumer
+          + ", stream="
+          + stream()
+          + ", status="
+          + status
+          + ", previousStatus="
+          + previousStatus
+          + '}';
+    }
   }
 
   private boolean canTrack() {
@@ -387,4 +412,26 @@ enum Status {
     NOT_AVAILABLE,
     CLOSED
   }
+
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    StreamConsumer that = (StreamConsumer) o;
+    return id == that.id && stream.equals(that.stream);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hash(id, stream);
+  }
+
+  @Override
+  public String toString() {
+    return "StreamConsumer{" + "id=" + id + ", stream='" + stream + '\'' + '}';
+  }
 }

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

@@ -133,12 +133,15 @@ public Consumer build() {
       throw new IllegalArgumentException("Stream and superStream cannot be set at the same time");
     }
     if (this.messageHandler == null) {
-      throw new IllegalArgumentException("A message handle must be set");
+      throw new IllegalArgumentException("A message handler must be set");
     }
     if (this.name == null
         && (this.autoTrackingStrategy != null || this.manualTrackingStrategy != null)) {
       throw new IllegalArgumentException("A name must be set if a tracking strategy is specified");
     }
+    if (Utils.isSac(this.subscriptionProperties) && this.name == null) {
+      throw new IllegalArgumentException("A name must be set if single active consumer is enabled");
+    }
 
     this.environment.maybeInitializeLocator();
     TrackingConfiguration trackingConfiguration;
@@ -179,6 +182,9 @@ public Consumer build() {
               this.consumerUpdateListener);
       environment.addConsumer((StreamConsumer) consumer);
     } else {
+      if (Utils.isSac(this.subscriptionProperties)) {
+        this.subscriptionProperties.put("super-stream", this.superStream);
+      }
       consumer =
           new SuperStreamConsumer(this, this.superStream, this.environment, trackingConfiguration);
     }

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

@@ -18,11 +18,14 @@
 import static com.rabbitmq.stream.impl.TestUtils.deleteSuperStreamTopology;
 import static com.rabbitmq.stream.impl.TestUtils.latchAssert;
 import static com.rabbitmq.stream.impl.TestUtils.localhost;
+import static com.rabbitmq.stream.impl.TestUtils.waitAtMost;
 import static org.assertj.core.api.Assertions.assertThat;
 
 import com.rabbitmq.client.Connection;
 import com.rabbitmq.client.ConnectionFactory;
 import com.rabbitmq.stream.Consumer;
+import com.rabbitmq.stream.ConsumerUpdateListener;
+import com.rabbitmq.stream.ConsumerUpdateListener.Status;
 import com.rabbitmq.stream.Environment;
 import com.rabbitmq.stream.EnvironmentBuilder;
 import com.rabbitmq.stream.OffsetSpecification;
@@ -31,10 +34,14 @@
 import java.nio.charset.StandardCharsets;
 import java.util.Collections;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -229,4 +236,91 @@ void autoOffsetTrackingShouldStoreOnAllPartitions() throws Exception {
         p -> assertThat(client.queryOffset(consumerName, p)).isGreaterThan(almostLastOffset));
     consumer.close();
   }
+
+  @Test
+  void sacShouldSpreadAcrossPartitions() throws Exception {
+    declareSuperStreamTopology(connection, superStream, partitionCount);
+    List<String> partitions =
+        IntStream.range(0, partitionCount)
+            .mapToObj(i -> superStream + "-" + i)
+            .collect(Collectors.toList());
+    Map<String, Status> consumerStates = new ConcurrentHashMap<>();
+    String consumerName = "my-app";
+    Function<String, Consumer> consumerCreator =
+        consumer -> {
+          return environment
+              .consumerBuilder()
+              .singleActiveConsumer()
+              .superStream(superStream)
+              .offset(OffsetSpecification.first())
+              .name(consumerName)
+              .manualTrackingStrategy()
+              .builder()
+              .messageHandler((context, message) -> {})
+              .consumerUpdateListener(
+                  new ConsumerUpdateListener() {
+                    @Override
+                    public OffsetSpecification update(Context context) {
+                      consumerStates.put(consumer + context.stream(), context.status());
+                      //                    System.out.println(context);
+                      return null;
+                    }
+                  })
+              .build();
+        };
+
+    Consumer consumer0 = consumerCreator.apply("0");
+
+    waitAtMost(
+        () ->
+            consumerStates.get("0" + partitions.get(0)) == Status.ACTIVE
+                && consumerStates.get("0" + partitions.get(1)) == Status.ACTIVE
+                && consumerStates.get("0" + partitions.get(2)) == Status.ACTIVE);
+
+    Consumer consumer1 = consumerCreator.apply("1");
+
+    waitAtMost(
+        () ->
+            consumerStates.get("0" + partitions.get(0)) == Status.ACTIVE
+                && consumerStates.get("1" + partitions.get(1)) == Status.ACTIVE
+                && consumerStates.get("0" + partitions.get(2)) == Status.ACTIVE);
+
+    Consumer consumer2 = consumerCreator.apply("2");
+
+    waitAtMost(
+        () ->
+            consumerStates.get("0" + partitions.get(0)) == Status.ACTIVE
+                && consumerStates.get("1" + partitions.get(1)) == Status.ACTIVE
+                && consumerStates.get("2" + partitions.get(2)) == Status.ACTIVE);
+
+    waitAtMost(() -> consumerStates.size() == 9);
+    assertThat(consumerStates)
+        .containsEntry("0" + partitions.get(0), Status.ACTIVE)
+        .containsEntry("0" + partitions.get(1), Status.PASSIVE)
+        .containsEntry("0" + partitions.get(2), Status.PASSIVE)
+        .containsEntry("1" + partitions.get(0), Status.PASSIVE)
+        .containsEntry("1" + partitions.get(1), Status.ACTIVE)
+        .containsEntry("1" + partitions.get(2), Status.PASSIVE)
+        .containsEntry("2" + partitions.get(0), Status.PASSIVE)
+        .containsEntry("2" + partitions.get(1), Status.PASSIVE)
+        .containsEntry("2" + partitions.get(2), Status.ACTIVE);
+
+    consumer0.close();
+
+    waitAtMost(
+        () ->
+            consumerStates.get("1" + partitions.get(0)) == Status.ACTIVE
+                && consumerStates.get("2" + partitions.get(1)) == Status.ACTIVE
+                && consumerStates.get("1" + partitions.get(2)) == Status.ACTIVE);
+
+    consumer1.close();
+
+    waitAtMost(
+        () ->
+            consumerStates.get("2" + partitions.get(0)) == Status.ACTIVE
+                && consumerStates.get("2" + partitions.get(1)) == Status.ACTIVE
+                && consumerStates.get("2" + partitions.get(2)) == Status.ACTIVE);
+
+    consumer2.close();
+  }
 }