Repository: skiptests/astraea
Pull request #564 (2022-08-09 12:46:03)

Enable to unsubscribe/resubscribe topics/partitions
Astraea consumer 建立後便固定了訂閱的資訊，這樣在測試consumer re-balance時會有些不方便，因此這隻PR讓consumer可以暫時解除訂閱，如此就能自由的觸發re-balance


Patch:
diff --git a/app/src/main/java/org/astraea/app/consumer/Builder.java b/app/src/main/java/org/astraea/app/consumer/Builder.java
index 72988b9e..2e517a0c 100644
--- a/app/src/main/java/org/astraea/app/consumer/Builder.java
+++ b/app/src/main/java/org/astraea/app/consumer/Builder.java
@@ -24,6 +24,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiConsumer;
 import org.apache.kafka.clients.consumer.ConsumerConfig;
 import org.astraea.app.admin.TopicPartition;
@@ -115,8 +116,9 @@ public abstract class Builder<Key, Value> {
   /** @return consumer instance. The different builders may return inherited consumer interface. */
   public abstract Consumer<Key, Value> build();
 
-  protected static class BaseConsumer<Key, Value> implements Consumer<Key, Value> {
+  protected abstract static class BaseConsumer<Key, Value> implements Consumer<Key, Value> {
     protected final org.apache.kafka.clients.consumer.Consumer<Key, Value> kafkaConsumer;
+    private final AtomicBoolean subscribed = new AtomicBoolean(true);
 
     public BaseConsumer(org.apache.kafka.clients.consumer.Consumer<Key, Value> kafkaConsumer) {
       this.kafkaConsumer = kafkaConsumer;
@@ -143,6 +145,18 @@ public abstract class Builder<Key, Value> {
     public void close() {
       kafkaConsumer.close();
     }
+
+    @Override
+    public void resubscribe() {
+      if (subscribed.compareAndSet(false, true)) doResubscribe();
+    }
+
+    @Override
+    public void unsubscribe() {
+      if (subscribed.compareAndSet(true, false)) kafkaConsumer.unsubscribe();
+    }
+
+    protected abstract void doResubscribe();
   }
 
   public enum SeekStrategy {
diff --git a/app/src/main/java/org/astraea/app/consumer/Consumer.java b/app/src/main/java/org/astraea/app/consumer/Consumer.java
index 593fe3e0..dcb4ea6e 100644
--- a/app/src/main/java/org/astraea/app/consumer/Consumer.java
+++ b/app/src/main/java/org/astraea/app/consumer/Consumer.java
@@ -49,6 +49,12 @@ public interface Consumer<Key, Value> extends AutoCloseable {
   @Override
   void close();
 
+  /** resubscribe partitions or rejoin the consumer group. */
+  void resubscribe();
+
+  /** unsubscribe all partitions. */
+  void unsubscribe();
+
   /**
    * Create a consumer builder by setting specific topics
    *
diff --git a/app/src/main/java/org/astraea/app/consumer/PartitionsBuilder.java b/app/src/main/java/org/astraea/app/consumer/PartitionsBuilder.java
index 9a267c66..9c5a90ca 100644
--- a/app/src/main/java/org/astraea/app/consumer/PartitionsBuilder.java
+++ b/app/src/main/java/org/astraea/app/consumer/PartitionsBuilder.java
@@ -100,14 +100,23 @@ public class PartitionsBuilder<Key, Value> extends Builder<Key, Value> {
 
     seekStrategy.apply(kafkaConsumer, seekValue);
 
-    return new AssignedConsumerImpl<>(kafkaConsumer);
+    return new AssignedConsumerImpl<>(kafkaConsumer, partitions);
   }
 
   private static class AssignedConsumerImpl<Key, Value> extends Builder.BaseConsumer<Key, Value>
       implements AssignedConsumer<Key, Value> {
 
-    public AssignedConsumerImpl(Consumer<Key, Value> kafkaConsumer) {
+    private final Set<TopicPartition> partitions;
+
+    public AssignedConsumerImpl(
+        Consumer<Key, Value> kafkaConsumer, Set<TopicPartition> partitions) {
       super(kafkaConsumer);
+      this.partitions = partitions;
+    }
+
+    @Override
+    protected void doResubscribe() {
+      kafkaConsumer.assign(partitions.stream().map(TopicPartition::to).collect(toList()));
     }
   }
 }
diff --git a/app/src/main/java/org/astraea/app/consumer/TopicsBuilder.java b/app/src/main/java/org/astraea/app/consumer/TopicsBuilder.java
index bf54a2be..523b3aef 100644
--- a/app/src/main/java/org/astraea/app/consumer/TopicsBuilder.java
+++ b/app/src/main/java/org/astraea/app/consumer/TopicsBuilder.java
@@ -122,15 +122,21 @@ public class TopicsBuilder<Key, Value> extends Builder<Key, Value> {
 
     seekStrategy.apply(kafkaConsumer, seekValue);
 
-    return new SubscribedConsumerImpl<>(kafkaConsumer);
+    return new SubscribedConsumerImpl<>(kafkaConsumer, topics, listener);
   }
 
   private static class SubscribedConsumerImpl<Key, Value> extends Builder.BaseConsumer<Key, Value>
       implements SubscribedConsumer<Key, Value> {
+    private final Set<String> topics;
+    private final ConsumerRebalanceListener listener;
 
     public SubscribedConsumerImpl(
-        org.apache.kafka.clients.consumer.Consumer<Key, Value> kafkaConsumer) {
+        org.apache.kafka.clients.consumer.Consumer<Key, Value> kafkaConsumer,
+        Set<String> topics,
+        ConsumerRebalanceListener listener) {
       super(kafkaConsumer);
+      this.topics = topics;
+      this.listener = listener;
     }
 
     @Override
@@ -151,5 +157,10 @@ public class TopicsBuilder<Key, Value> extends Builder<Key, Value> {
     public Optional<String> groupInstanceId() {
       return kafkaConsumer.groupMetadata().groupInstanceId();
     }
+
+    @Override
+    protected void doResubscribe() {
+      kafkaConsumer.subscribe(topics, ConsumerRebalanceListener.of(listener));
+    }
   }
 }
diff --git a/app/src/main/java/org/astraea/app/performance/ConsumerThread.java b/app/src/main/java/org/astraea/app/performance/ConsumerThread.java
index 17fde14f..930bac37 100644
--- a/app/src/main/java/org/astraea/app/performance/ConsumerThread.java
+++ b/app/src/main/java/org/astraea/app/performance/ConsumerThread.java
@@ -63,10 +63,17 @@ public interface ConsumerThread extends AbstractThread {
               var report = reports.get(index);
               var closeLatch = closeLatches.get(index);
               var closed = closeFlags.get(index);
+              var subscribed = new AtomicBoolean(true);
               executors.execute(
                   () -> {
                     try {
                       while (!closed.get()) {
+                        if (subscribed.get()) consumer.resubscribe();
+                        else {
+                          consumer.unsubscribe();
+                          Utils.sleep(Duration.ofSeconds(1));
+                          continue;
+                        }
                         consumer
                             .poll(Duration.ofSeconds(1))
                             .forEach(
@@ -102,6 +109,16 @@ public interface ConsumerThread extends AbstractThread {
                   return closeLatch.getCount() == 0;
                 }
 
+                @Override
+                public void resubscribe() {
+                  subscribed.set(true);
+                }
+
+                @Override
+                public void unsubscribe() {
+                  subscribed.set(false);
+                }
+
                 @Override
                 public Report report() {
                   return report;
@@ -118,6 +135,10 @@ public interface ConsumerThread extends AbstractThread {
         .collect(Collectors.toUnmodifiableList());
   }
 
+  void resubscribe();
+
+  void unsubscribe();
+
   /** @return report of this thread */
   Report report();
 }


Test patch:
diff --git a/app/src/test/java/org/astraea/app/consumer/ConsumerTest.java b/app/src/test/java/org/astraea/app/consumer/ConsumerTest.java
index efa70bc1..ca056022 100644
--- a/app/src/test/java/org/astraea/app/consumer/ConsumerTest.java
+++ b/app/src/test/java/org/astraea/app/consumer/ConsumerTest.java
@@ -347,4 +347,60 @@ public class ConsumerTest extends RequireBrokerCluster {
                 .build(),
         "seek value should >= 0");
   }
+
+  @Test
+  void testResubscribeTopics() {
+    var topic = Utils.randomString(10);
+    produceData(topic, 100);
+    try (var consumer =
+        Consumer.forTopics(Set.of(topic))
+            .bootstrapServers(bootstrapServers())
+            .fromBeginning()
+            .build()) {
+      Assertions.assertNotEquals(0, consumer.poll(Duration.ofSeconds(5)).size());
+      consumer.unsubscribe();
+      Assertions.assertThrows(
+          IllegalStateException.class, () -> consumer.poll(Duration.ofSeconds(2)));
+      // unsubscribe is idempotent op
+      consumer.unsubscribe();
+      consumer.unsubscribe();
+      consumer.unsubscribe();
+
+      consumer.resubscribe();
+      Assertions.assertNotEquals(0, consumer.poll(Duration.ofSeconds(5)).size());
+
+      // resubscribe is idempotent op
+      consumer.resubscribe();
+      consumer.resubscribe();
+      consumer.resubscribe();
+    }
+  }
+
+  @Test
+  void testResubscribePartitions() {
+    var topic = Utils.randomString(10);
+    produceData(topic, 100);
+    try (var consumer =
+        Consumer.forPartitions(Set.of(TopicPartition.of(topic, 0)))
+            .bootstrapServers(bootstrapServers())
+            .fromBeginning()
+            .build()) {
+      Assertions.assertNotEquals(0, consumer.poll(Duration.ofSeconds(5)).size());
+      consumer.unsubscribe();
+      Assertions.assertThrows(
+          IllegalStateException.class, () -> consumer.poll(Duration.ofSeconds(2)));
+      // unsubscribe is idempotent op
+      consumer.unsubscribe();
+      consumer.unsubscribe();
+      consumer.unsubscribe();
+
+      consumer.resubscribe();
+      Assertions.assertNotEquals(0, consumer.poll(Duration.ofSeconds(5)).size());
+
+      // resubscribe is idempotent op
+      consumer.resubscribe();
+      consumer.resubscribe();
+      consumer.resubscribe();
+    }
+  }
 }