Subscription IT: intro retry rule for flaky tests (#15698)

Author: VGalaxies

integration-test/src/test/java/org/apache/iotdb/subscription/it/Retry.java

@@ -0,0 +1,33 @@
+/*
+ * 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.iotdb.subscription.it;
+
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
+
+/** Marks a test method to specify how many times it should be retried (first run + retries). */
+@Retention(RetentionPolicy.RUNTIME)
+@Target(ElementType.METHOD)
+public @interface Retry {
+  /** Total execution count = 1 (initial run) + number of retries */
+  int times() default 3;
+}

integration-test/src/test/java/org/apache/iotdb/subscription/it/RetryRule.java

@@ -0,0 +1,69 @@
+/*
+ * 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.iotdb.subscription.it;
+
+import org.junit.rules.TestRule;
+import org.junit.runner.Description;
+import org.junit.runners.model.Statement;
+
+/** Controls retry logic for test failures based on the {@link Retry} annotation. */
+public class RetryRule implements TestRule {
+
+  @Override
+  public Statement apply(final Statement base, final Description description) {
+    // Read the annotation on the method; if absent, do not retry (times = 1)
+    final Retry retry = description.getAnnotation(Retry.class);
+    final int times = (retry != null ? retry.times() : 1);
+    return new RetryStatement(base, description, times);
+  }
+
+  private static class RetryStatement extends Statement {
+    private final Statement base;
+    private final Description description;
+    private final int times;
+
+    RetryStatement(final Statement base, final Description description, final int times) {
+      this.base = base;
+      this.description = description;
+      this.times = times;
+    }
+
+    @Override
+    public void evaluate() throws Throwable {
+      Throwable lastThrowable;
+      for (int i = 1; i <= times; i++) {
+        try {
+          base.evaluate();
+          return; // Return immediately on success
+        } catch (final Throwable t) {
+          lastThrowable = t;
+          System.err.printf(
+              "[%s] run %d/%d failed: %s%n",
+              description.getDisplayName(), i, times, t.getMessage());
+          if (i == times) {
+            // If it's the last attempt, and it still fails, throw the exception
+            throw lastThrowable;
+          }
+          // Otherwise, continue to the next retry
+        }
+      }
+    }
+  }
+}

integration-test/src/test/java/org/apache/iotdb/subscription/it/triple/regression/pushconsumer/mode/IoTDBSnapshotTSPatternDatasetPushConsumerIT.java

@@ -28,6 +28,8 @@
 import org.apache.iotdb.session.subscription.consumer.ConsumeResult;
 import org.apache.iotdb.session.subscription.consumer.SubscriptionPushConsumer;
 import org.apache.iotdb.session.subscription.payload.SubscriptionSessionDataSet;
+import org.apache.iotdb.subscription.it.Retry;
+import org.apache.iotdb.subscription.it.RetryRule;
 import org.apache.iotdb.subscription.it.triple.regression.AbstractSubscriptionRegressionIT;
 
 import org.apache.thrift.TException;
@@ -38,6 +40,7 @@
 import org.apache.tsfile.write.schema.MeasurementSchema;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
@@ -56,6 +59,9 @@
 @RunWith(IoTDBTestRunner.class)
 @Category({MultiClusterIT2SubscriptionRegressionConsumer.class})
 public class IoTDBSnapshotTSPatternDatasetPushConsumerIT extends AbstractSubscriptionRegressionIT {
+
+  @Rule public RetryRule retryRule = new RetryRule();
+
   private static final String database = "root.test.SnapshotTSPatternDatasetPushConsumer";
   private static final String database2 = "root.SnapshotTSPatternDatasetPushConsumer";
   private static final String device = database + ".d_0";
@@ -108,6 +114,7 @@ public void tearDown() throws Exception {
     subs.dropTopic(topicName);
     dropDB(database);
     dropDB(database2);
+    schemaList.clear();
     super.tearDown();
   }
 
@@ -127,6 +134,7 @@ private void insert_data(long timestamp)
   }
 
   @Test
+  @Retry
   public void do_test()
       throws InterruptedException,
           TException,

integration-test/src/test/java/org/apache/iotdb/subscription/it/triple/regression/pushconsumer/mode/IoTDBSnapshotTSPatternTsfilePushConsumerIT.java

@@ -27,6 +27,8 @@
 import org.apache.iotdb.session.subscription.consumer.AckStrategy;
 import org.apache.iotdb.session.subscription.consumer.ConsumeResult;
 import org.apache.iotdb.session.subscription.consumer.SubscriptionPushConsumer;
+import org.apache.iotdb.subscription.it.Retry;
+import org.apache.iotdb.subscription.it.RetryRule;
 import org.apache.iotdb.subscription.it.triple.regression.AbstractSubscriptionRegressionIT;
 
 import org.apache.thrift.TException;
@@ -42,6 +44,7 @@
 import org.apache.tsfile.write.schema.MeasurementSchema;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
@@ -63,6 +66,9 @@
 @RunWith(IoTDBTestRunner.class)
 @Category({MultiClusterIT2SubscriptionRegressionConsumer.class})
 public class IoTDBSnapshotTSPatternTsfilePushConsumerIT extends AbstractSubscriptionRegressionIT {
+
+  @Rule public RetryRule retryRule = new RetryRule();
+
   private static final String database = "root.test.SnapshotTSPatternTsfilePushConsumer";
   private static final String database2 = "root.SnapshotTSPatternTsfilePushConsumer";
   private static final String device = database + ".d_0";
@@ -122,6 +128,7 @@ public void tearDown() throws Exception {
     subs.dropTopic(topicName);
     dropDB(database);
     dropDB(database2);
+    schemaList.clear();
     super.tearDown();
   }
 
@@ -141,6 +148,7 @@ private void insert_data(long timestamp)
   }
 
   @Test
+  @Retry
   public void do_test()
       throws InterruptedException,
           TException,

integration-test/src/test/java/org/apache/iotdb/subscription/it/triple/regression/pushconsumer/multi/IoTDBOneConsumerMultiTopicsTsfileIT.java

@@ -26,6 +26,8 @@
 import org.apache.iotdb.session.subscription.consumer.AckStrategy;
 import org.apache.iotdb.session.subscription.consumer.ConsumeResult;
 import org.apache.iotdb.session.subscription.consumer.SubscriptionPushConsumer;
+import org.apache.iotdb.subscription.it.Retry;
+import org.apache.iotdb.subscription.it.RetryRule;
 import org.apache.iotdb.subscription.it.triple.regression.AbstractSubscriptionRegressionIT;
 
 import org.apache.thrift.TException;
@@ -41,6 +43,7 @@
 import org.apache.tsfile.write.schema.MeasurementSchema;
 import org.junit.After;
 import org.junit.Before;
+import org.junit.Rule;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
 import org.junit.runner.RunWith;
@@ -60,6 +63,9 @@
 @RunWith(IoTDBTestRunner.class)
 @Category({MultiClusterIT2SubscriptionRegressionConsumer.class})
 public class IoTDBOneConsumerMultiTopicsTsfileIT extends AbstractSubscriptionRegressionIT {
+
+  @Rule public RetryRule retryRule = new RetryRule();
+
   private static final String database = "root.test.OneConsumerMultiTopicsTsfile";
   private static final String database2 = "root.OneConsumerMultiTopicsTsfile";
   private static final String device = database + ".d_0";
@@ -122,6 +128,7 @@ private void insert_data(long timestamp, String device)
   }
 
   @Test
+  @Retry
   public void do_test()
       throws InterruptedException,
           TException,