Subscription: support TsFile message format & improve exception handling logic

example/session/src/main/java/org/apache/iotdb/SubscriptionSessionExample.java

@@ -22,38 +22,51 @@
 import org.apache.iotdb.isession.SessionDataSet;
 import org.apache.iotdb.isession.util.Version;
 import org.apache.iotdb.rpc.subscription.config.ConsumerConstant;
+import org.apache.iotdb.rpc.subscription.config.TopicConstant;
 import org.apache.iotdb.session.Session;
-import org.apache.iotdb.session.subscription.SubscriptionMessage;
 import org.apache.iotdb.session.subscription.SubscriptionPullConsumer;
 import org.apache.iotdb.session.subscription.SubscriptionSession;
-import org.apache.iotdb.session.subscription.SubscriptionSessionDataSet;
-import org.apache.iotdb.session.subscription.SubscriptionSessionDataSets;
+import org.apache.iotdb.session.subscription.payload.SubscriptionMessage;
+import org.apache.iotdb.session.subscription.payload.SubscriptionSessionDataSet;
 
-import java.time.Duration;
+import org.apache.tsfile.read.TsFileReader;
+
+import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 import java.util.Properties;
+import java.util.concurrent.locks.LockSupport;
 
 public class SubscriptionSessionExample {
 
   private static Session session;
 
-  private static final String LOCAL_HOST = "127.0.0.1";
+  private static final String HOST = "127.0.0.1";
+  private static final int PORT = 6667;
+
+  private static final String TOPIC_1 = "topic1";
+  private static final String TOPIC_2 = "`topic2`";
+
+  public static final long SLEEP_NS = 1_000_000_000L;
+  public static final long POLL_TIMEOUT_MS = 10_000L;
+  private static final int MAX_RETRY_TIMES = 3;
 
-  public static void main(String[] args) throws Exception {
+  private static void prepareData() throws Exception {
+    // Open session
     session =
         new Session.Builder()
-            .host(LOCAL_HOST)
-            .port(6667)
+            .host(HOST)
+            .port(PORT)
             .username("root")
             .password("root")
             .version(Version.V_1_0)
             .build();
     session.open(false);
 
     // Insert some historical data
-    long currentTime = System.currentTimeMillis();
-    for (int i = 0; i < 100; ++i) {
+    final long currentTime = System.currentTimeMillis();
+    for (int i = 0; i < 10000; ++i) {
       session.executeNonQueryStatement(
           String.format("insert into root.db.d1(time, s1, s2) values (%s, 1, 2)", i));
       session.executeNonQueryStatement(
@@ -63,31 +76,60 @@ public static void main(String[] args) throws Exception {
     }
     session.executeNonQueryStatement("flush");
 
-    // Create topic
-    final String topic1 = "topic1";
-    final String topic2 = "`topic2`";
-    try (SubscriptionSession subscriptionSession = new SubscriptionSession(LOCAL_HOST, 6667)) {
+    // Close session
+    session.close();
+    session = null;
+  }
+
+  private static void dataQuery() throws Exception {
+    // Open session
+    session =
+        new Session.Builder()
+            .host(HOST)
+            .port(PORT)
+            .username("root")
+            .password("root")
+            .version(Version.V_1_0)
+            .build();
+    session.open(false);
+
+    // Query
+    final SessionDataSet dataSet = session.executeQueryStatement("select ** from root.**");
+    while (dataSet.hasNext()) {
+      System.out.println(dataSet.next());
+    }
+
+    // Close session
+    session.close();
+    session = null;
+  }
+
+  private static void subscriptionExample1() throws Exception {
+    // Create topics
+    try (final SubscriptionSession subscriptionSession = new SubscriptionSession(HOST, PORT)) {
       subscriptionSession.open();
-      subscriptionSession.createTopic(topic1);
-      subscriptionSession.createTopic(topic2);
+      subscriptionSession.createTopic(TOPIC_1);
     }
 
+    int retryCount = 0;
     // Subscription: property-style ctor
-    Properties config = new Properties();
+    final Properties config = new Properties();
     config.put(ConsumerConstant.CONSUMER_ID_KEY, "c1");
     config.put(ConsumerConstant.CONSUMER_GROUP_ID_KEY, "cg1");
-    SubscriptionPullConsumer consumer1 = new SubscriptionPullConsumer(config);
+    final SubscriptionPullConsumer consumer1 = new SubscriptionPullConsumer(config);
     consumer1.open();
-    consumer1.subscribe(topic1);
+    consumer1.subscribe(TOPIC_1);
     while (true) {
-      Thread.sleep(1000); // Wait for some time
-      List<SubscriptionMessage> messages = consumer1.poll(Duration.ofMillis(10000));
+      LockSupport.parkNanos(SLEEP_NS); // wait some time
+      final List<SubscriptionMessage> messages = consumer1.poll(POLL_TIMEOUT_MS);
       if (messages.isEmpty()) {
-        break;
+        retryCount++;
+        if (retryCount >= MAX_RETRY_TIMES) {
+          break;
+        }
       }
-      for (SubscriptionMessage message : messages) {
-        SubscriptionSessionDataSets payload = (SubscriptionSessionDataSets) message.getPayload();
-        for (SubscriptionSessionDataSet dataSet : payload) {
+      for (final SubscriptionMessage message : messages) {
+        for (final SubscriptionSessionDataSet dataSet : message.getSessionDataSetsHandler()) {
           System.out.println(dataSet.getColumnNames());
           System.out.println(dataSet.getColumnTypes());
           while (dataSet.hasNext()) {
@@ -99,51 +141,75 @@ public static void main(String[] args) throws Exception {
     }
 
     // Show topics and subscriptions
-    try (SubscriptionSession subscriptionSession = new SubscriptionSession(LOCAL_HOST, 6667)) {
+    try (final SubscriptionSession subscriptionSession = new SubscriptionSession(HOST, PORT)) {
       subscriptionSession.open();
       subscriptionSession.getTopics().forEach((System.out::println));
       subscriptionSession.getSubscriptions().forEach((System.out::println));
     }
 
-    consumer1.unsubscribe(topic1);
+    consumer1.unsubscribe(TOPIC_1);
     consumer1.close();
+  }
 
-    // Subscription: builder-style ctor
-    try (SubscriptionPullConsumer consumer2 =
-        new SubscriptionPullConsumer.Builder()
-            .consumerId("c2")
-            .consumerGroupId("cg2")
-            .autoCommit(false)
-            .buildPullConsumer()) {
-      consumer2.open();
-      consumer2.subscribe(topic2);
-      while (true) {
-        Thread.sleep(1000); // wait some time
-        List<SubscriptionMessage> messages =
-            consumer2.poll(Collections.singleton(topic2), Duration.ofMillis(10000));
-        if (messages.isEmpty()) {
-          break;
-        }
-        for (SubscriptionMessage message : messages) {
-          SubscriptionSessionDataSets payload = (SubscriptionSessionDataSets) message.getPayload();
-          for (SubscriptionSessionDataSet dataSet : payload) {
-            System.out.println(dataSet.getColumnNames());
-            System.out.println(dataSet.getColumnTypes());
-            while (dataSet.hasNext()) {
-              System.out.println(dataSet.next());
-            }
-          }
-        }
-        consumer2.commitSync(messages);
-      }
-      consumer2.unsubscribe(topic2);
+  private static void subscriptionExample2() throws Exception {
+    try (final SubscriptionSession subscriptionSession = new SubscriptionSession(HOST, PORT)) {
+      subscriptionSession.open();
+      final Properties config = new Properties();
+      config.put(TopicConstant.FORMAT_KEY, TopicConstant.FORMAT_TS_FILE_HANDLER_VALUE);
+      subscriptionSession.createTopic(TOPIC_2, config);
     }
 
-    // Query
-    SessionDataSet dataSet = session.executeQueryStatement("select ** from root.**");
-    while (dataSet.hasNext()) {
-      System.out.println(dataSet.next());
+    final List<Thread> threads = new ArrayList<>();
+    for (int i = 0; i < 8; ++i) {
+      final int idx = i;
+      final Thread thread =
+          new Thread(
+              () -> {
+                int retryCount = 0;
+                // Subscription: builder-style ctor
+                try (final SubscriptionPullConsumer consumer2 =
+                    new SubscriptionPullConsumer.Builder()
+                        .consumerId("c" + idx)
+                        .consumerGroupId("cg2")
+                        .autoCommit(false)
+                        .buildPullConsumer()) {
+                  consumer2.open();
+                  consumer2.subscribe(TOPIC_2);
+                  while (true) {
+                    LockSupport.parkNanos(SLEEP_NS); // wait some time
+                    final List<SubscriptionMessage> messages =
+                        consumer2.poll(Collections.singleton(TOPIC_2), POLL_TIMEOUT_MS);
+                    if (messages.isEmpty()) {
+                      retryCount++;
+                      if (retryCount >= MAX_RETRY_TIMES) {
+                        break;
+                      }
+                    }
+                    for (final SubscriptionMessage message : messages) {
+                      try (final TsFileReader reader = message.getTsFileHandler().openReader()) {
+                        // do something...
+                      }
+                    }
+                    consumer2.commitSync(messages);
+                  }
+                  consumer2.unsubscribe(TOPIC_2);
+                } catch (final IOException e) {
+                  throw new RuntimeException(e);
+                }
+              });
+      thread.start();
+      threads.add(thread);
     }
-    session.close();
+
+    for (final Thread thread : threads) {
+      thread.join();
+    }
+  }
+
+  public static void main(final String[] args) throws Exception {
+    prepareData();
+    dataQuery();
+    subscriptionExample1();
+    subscriptionExample2();
   }
 }

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

@@ -0,0 +1,30 @@
+/*
+ * 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;
+
+public class IoTDBSubscriptionITConstant {
+
+  public static final long AWAITILITY_POLL_DELAY_SECOND = 1L;
+  public static final long AWAITILITY_POLL_INTERVAL_SECOND = 2L;
+  public static final long AWAITILITY_AT_MOST_SECOND = 240L;
+
+  public static final long SLEEP_NS = 1_000_000_000L;
+  public static final long POLL_TIMEOUT_MS = 10_000L;
+}

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

@@ -36,8 +36,11 @@ public void setUp() {
     senderEnv = MultiEnvFactory.getEnv(0);
     receiverEnv = MultiEnvFactory.getEnv(1);
 
+    // enable auto create schema
     senderEnv.getConfig().getCommonConfig().setAutoCreateSchemaEnabled(true);
     receiverEnv.getConfig().getCommonConfig().setAutoCreateSchemaEnabled(true);
+
+    // for IoTDBSubscriptionConsumerGroupIT
     receiverEnv.getConfig().getCommonConfig().setPipeAirGapReceiverEnabled(true);
 
     // 10 min, assert that the operations will not time out